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

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 formatter.print_info(f"Uploading {local_path} as '{name}'...") 

42 result = manager.upload(local_path, name) 

43 

44 formatter.print_success( 

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

46 ) 

47 formatter.print_info( 

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

49 ) 

50 

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

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("list") 

60@pass_config 

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

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

63 

64 Examples: 

65 gco models list 

66 """ 

67 from ..models import get_model_manager 

68 

69 formatter = get_output_formatter(config) 

70 

71 try: 

72 manager = get_model_manager(config) 

73 model_list = manager.list_models() 

74 

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

76 formatter.print(model_list) 

77 return 

78 

79 if not model_list: 

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

81 return 

82 

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

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

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

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

87 for m in model_list: 

88 print( 

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

90 ) 

91 print() 

92 

93 except Exception as e: 

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

95 sys.exit(1) 

96 

97 

98@models.command("delete") 

99@click.argument("model_name") 

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

101@pass_config 

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

103 """Delete a model from the central S3 bucket. 

104 

105 Examples: 

106 gco models delete llama3-8b -y 

107 """ 

108 from ..models import get_model_manager 

109 

110 formatter = get_output_formatter(config) 

111 

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

113 click.confirm(f"Delete model '{model_name}' and all its files?", abort=True) 

114 

115 try: 

116 manager = get_model_manager(config) 

117 deleted = manager.delete_model(model_name) 

118 

119 if deleted > 0: 

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

121 else: 

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

123 

124 except Exception as e: 

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

126 sys.exit(1) 

127 

128 

129@models.command("uri") 

130@click.argument("model_name") 

131@pass_config 

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

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

134 

135 Examples: 

136 gco models uri llama3-8b 

137 """ 

138 from ..models import get_model_manager 

139 

140 formatter = get_output_formatter(config) 

141 

142 try: 

143 manager = get_model_manager(config) 

144 uri = manager.get_model_uri(model_name) 

145 print(uri) 

146 

147 except Exception as e: 

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

149 sys.exit(1)