Coverage for cli/commands/models_cmd.py: 90.91%

1"""Model weight management commands.""" 

2 

3import sys 

4from typing import Any 

5 

6import click 

7 

8from ..config import GCOConfig 

9from ..output import get_output_formatter 

10 

11pass_config = click.make_pass_decorator(GCOConfig, ensure=True) 

12 

13 

14@click.group() 

15@pass_config 

16def models(config: Any) -> None: 

17 """Manage model weights in the central S3 bucket.""" 

18 pass 

19 

20 

21@models.command("upload") 

22@click.argument("local_path") 

23@click.option("--name", "-n", required=True, help="Model name in the registry") 

24@pass_config 

25def models_upload(config: Any, local_path: Any, name: Any) -> None: 

26 """Upload model weights to the central S3 bucket. 

27 

28 Models uploaded here are available to inference endpoints in all regions. 

29 The inference_monitor syncs them to local EFS via an init container. 

30 

31 Examples: 

32 gco models upload ./my-model/ --name llama3-8b 

33 gco models upload ./weights.safetensors --name my-model 

34 """ 

35 from ..models import get_model_manager 

36 

37 formatter = get_output_formatter(config) 

38 

39 try: 

40 manager = get_model_manager(config) 

41 if config.output_format == "table": 

42 formatter.print_info(f"Uploading {local_path} as '{name}'...") 

43 result = manager.upload(local_path, name) 

44 

45 if config.output_format == "table": 

46 formatter.print_success( 

47 f"Uploaded {result['files_uploaded']} file(s) to {result['s3_uri']}" 

48 ) 

49 formatter.print_info( 

50 f"Use --model-source {result['s3_uri']} when deploying inference endpoints" 

51 ) 

52 else: 

53 formatter.print(result) 

54 

55 except Exception as e: 

56 formatter.print_error(f"Failed to upload model: {e}") 

57 sys.exit(1) 

58 

59 

60@models.command("upload-regional") 

61@click.argument("local_path") 

62@click.option( 

63 "--region", 

64 "-r", 

65 required=True, 

66 help="Target region whose regional bucket receives the objects", 

67) 

68@click.option( 

69 "--prefix", 

70 default="uploads", 

71 show_default=True, 

72 help="S3 prefix for uploaded objects", 

73) 

74@pass_config 

75def models_upload_regional(config: Any, local_path: Any, region: Any, prefix: Any) -> None: 

76 """Upload local files or a directory to a region's regional bucket. 

77 

78 Objects are written to that region's general-purpose 

79 gco-regional-shared-<account>-<region> bucket, resolved from the target 

80 region's own SSM parameter. The bucket is general purpose and usable by 

81 any in-region workload. 

82 

83 Examples: 

84 gco models upload-regional ./data/ --region us-east-1 

85 gco models upload-regional ./file.bin -r eu-west-1 --prefix datasets 

86 """ 

87 from ..models import get_regional_bucket_manager 

88 

89 formatter = get_output_formatter(config) 

90 

91 try: 

92 manager = get_regional_bucket_manager(config) 

93 if config.output_format == "table": 93 ↛ 95line 93 didn't jump to line 95 because the condition on line 93 was always true

94 formatter.print_info(f"Uploading {local_path} to region '{region}'...") 

95 result = manager.upload(local_path, region, prefix=prefix) 

96 

97 if config.output_format == "table": 

98 formatter.print_success( 

99 f"Uploaded {result['files_uploaded']} file(s) to {result['s3_uri']}" 

100 ) 

101 else: 

102 formatter.print(result) 

103 

104 except Exception as e: 

105 formatter.print_error(f"Failed to upload to regional bucket: {e}") 

106 sys.exit(1) 

107 

108 

109@models.command("list") 

110@pass_config 

111def models_list(config: Any) -> None: 

112 """List models in the central S3 bucket. 

113 

114 Examples: 

115 gco models list 

116 """ 

117 from ..models import get_model_manager 

118 

119 formatter = get_output_formatter(config) 

120 

121 try: 

122 manager = get_model_manager(config) 

123 model_list = manager.list_models() 

124 

125 if config.output_format != "table": 125 ↛ 126line 125 didn't jump to line 126 because the condition on line 125 was never true

126 formatter.print(model_list) 

127 return 

128 

129 if not model_list: 

130 formatter.print_info("No models found. Upload with 'gco models upload'") 

131 return 

132 

133 print(f"\n Models ({len(model_list)} found)") 

134 print(" " + "-" * 70) 

135 print(f" {'NAME':<25} {'FILES':>5} {'SIZE (GB)':>10} {'S3 URI'}") 

136 print(" " + "-" * 70) 

137 for m in model_list: 

138 print( 

139 f" {m['model_name']:<25} {m['files']:>5} {m['total_size_gb']:>10.2f} {m['s3_uri']}" 

140 ) 

141 print() 

142 

143 except Exception as e: 

144 formatter.print_error(f"Failed to list models: {e}") 

145 sys.exit(1) 

146 

147 

148@models.command("delete") 

149@click.argument("model_name") 

150@click.option("--yes", "-y", is_flag=True, help="Skip confirmation") 

151@pass_config 

152def models_delete(config: Any, model_name: Any, yes: Any) -> None: 

153 """Permanently delete a model and all object history from the central S3 bucket. 

154 

155 Every current object, historical version, and delete marker beneath the 

156 model prefix is removed. Successful batches cannot be restored if a later 

157 batch fails. 

158 

159 Examples: 

160 gco models delete llama3-8b -y 

161 """ 

162 from ..models import get_model_manager 

163 

164 formatter = get_output_formatter(config) 

165 

166 if not yes: 166 ↛ 167line 166 didn't jump to line 167 because the condition on line 166 was never true

167 click.confirm( 

168 f"Permanently delete model '{model_name}', including all current files " 

169 "and historical S3 versions? This cannot be undone.", 

170 abort=True, 

171 ) 

172 

173 try: 

174 manager = get_model_manager(config) 

175 deleted = manager.delete_model(model_name) 

176 

177 if deleted > 0: 

178 formatter.print_success(f"Deleted {deleted} file(s) for model '{model_name}'") 

179 else: 

180 formatter.print_warning(f"No files found for model '{model_name}'") 

181 

182 except Exception as e: 

183 formatter.print_error(f"Failed to delete model: {e}") 

184 sys.exit(1) 

185 

186 

187@models.command("uri") 

188@click.argument("model_name") 

189@pass_config 

190def models_uri(config: Any, model_name: Any) -> None: 

191 """Get the S3 URI for a model (for use with --model-source). 

192 

193 Examples: 

194 gco models uri llama3-8b 

195 """ 

196 from ..models import get_model_manager 

197 

198 formatter = get_output_formatter(config) 

199 

200 try: 

201 manager = get_model_manager(config) 

202 uri = manager.get_model_uri(model_name) 

203 print(uri) 

204 

205 except Exception as e: 

206 formatter.print_error(f"Failed to get model URI: {e}") 

207 sys.exit(1)