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

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

2 

3import sys 

4from typing import Any 

5 

6import click 

7 

8from ..config import GCOConfig 

9from ..output import confirm, 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 

19 

20@models.command("upload") 

21@click.argument("local_path") 

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

23@pass_config 

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

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

26 

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

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

29 

30 Examples: 

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

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

33 """ 

34 from ..models import get_model_manager 

35 

36 formatter = get_output_formatter(config) 

37 

38 try: 

39 manager = get_model_manager(config) 

40 if config.output_format == "table": 

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

42 result = manager.upload(local_path, name) 

43 

44 if config.output_format == "table": 

45 formatter.print_success( 

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

47 ) 

48 formatter.print_info( 

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

50 ) 

51 else: 

52 formatter.print(result) 

53 

54 except Exception as e: 

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

56 sys.exit(1) 

57 

58 

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

60@click.argument("local_path") 

61@click.option( 

62 "--region", 

63 "-r", 

64 required=True, 

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

66) 

67@click.option( 

68 "--prefix", 

69 default="uploads", 

70 show_default=True, 

71 help="S3 prefix for uploaded objects", 

72) 

73@pass_config 

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

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

76 

77 Objects are written to that region's general-purpose regional bucket, 

78 resolved from the target region's own /<project>/regional-shared-bucket/name 

79 SSM parameter (the bucket's name is CloudFormation-generated). The bucket is general purpose and usable by 

80 any in-region workload. 

81 

82 Examples: 

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

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

85 """ 

86 from ..models import get_regional_bucket_manager 

87 

88 formatter = get_output_formatter(config) 

89 

90 try: 

91 manager = get_regional_bucket_manager(config) 

92 if config.output_format == "table": 

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

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

95 

96 if config.output_format == "table": 

97 formatter.print_success( 

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

99 ) 

100 else: 

101 formatter.print(result) 

102 

103 except Exception as e: 

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

105 sys.exit(1) 

106 

107 

108@models.command("list") 

109@pass_config 

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

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

112 

113 Examples: 

114 gco models list 

115 """ 

116 from ..models import get_model_manager 

117 

118 formatter = get_output_formatter(config) 

119 

120 try: 

121 manager = get_model_manager(config) 

122 model_list = manager.list_models() 

123 

124 if config.output_format != "table": 

125 formatter.print(model_list) 

126 return 

127 

128 if not model_list: 

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

130 return 

131 

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

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

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

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

136 for m in model_list: 

137 print( 

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

139 ) 

140 print() 

141 

142 except Exception as e: 

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

144 sys.exit(1) 

145 

146 

147@models.command("delete") 

148@click.argument("model_name") 

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

150@pass_config 

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

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

153 

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

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

156 batch fails. 

157 

158 Examples: 

159 gco models delete llama3-8b -y 

160 """ 

161 from ..models import get_model_manager 

162 

163 formatter = get_output_formatter(config) 

164 

165 if not yes: 

166 confirm( 

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

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

169 abort=True, 

170 ) 

171 

172 try: 

173 manager = get_model_manager(config) 

174 deleted = manager.delete_model(model_name) 

175 

176 if deleted > 0: 

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

178 else: 

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

180 

181 except Exception as e: 

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

183 sys.exit(1) 

184 

185 

186@models.command("uri") 

187@click.argument("model_name") 

188@pass_config 

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

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

191 

192 Examples: 

193 gco models uri llama3-8b 

194 """ 

195 from ..models import get_model_manager 

196 

197 formatter = get_output_formatter(config) 

198 

199 try: 

200 manager = get_model_manager(config) 

201 uri = manager.get_model_uri(model_name) 

202 if config.output_format == "table": 

203 print(uri) 

204 else: 

205 formatter.print({"model_name": model_name, "s3_uri": uri}) 

206 

207 except Exception as e: 

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

209 sys.exit(1)