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

1"""Commands for discovering and syncing GCO S3 buckets.""" 

2 

3import signal 

4import sys 

5import threading 

6from collections.abc import Iterator 

7from contextlib import contextmanager 

8from types import FrameType 

9from typing import Any 

10 

11import click 

12 

13from ..config import GCOConfig 

14from ..output import get_output_formatter 

15 

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

17 

18 

19class _StorageSyncTerminated(RuntimeError): 

20 """Raised on SIGTERM so managed S3 transfers can unwind cooperatively.""" 

21 

22 

23@contextmanager 

24def _cooperative_storage_sigterm() -> Iterator[None]: 

25 """Turn SIGTERM into an exception while a storage transfer is active.""" 

26 if threading.current_thread() is not threading.main_thread(): 

27 yield 

28 return 

29 

30 previous_handler = signal.getsignal(signal.SIGTERM) 

31 

32 def terminate_handler(signum: int, frame: FrameType | None) -> None: 

33 raise _StorageSyncTerminated( 

34 "Storage sync was terminated; in-progress managed transfers were cancelled" 

35 ) 

36 

37 signal.signal(signal.SIGTERM, terminate_handler) 

38 try: 

39 yield 

40 finally: 

41 signal.signal(signal.SIGTERM, previous_handler) 

42 

43 

44@click.group() 

45@pass_config 

46def storage(config: Any) -> None: 

47 """Discover and sync user-facing GCO S3 buckets.""" 

48 pass 

49 

50 

51@storage.command("list") 

52@click.option( 

53 "--region", 

54 "-r", 

55 help="Limit regional-bucket discovery to this region", 

56) 

57@pass_config 

58def storage_list(config: Any, region: str | None) -> None: 

59 """List deployed buckets and the aliases accepted by `storage sync`. 

60 

61 Examples: 

62 gco storage list 

63 gco storage list --region us-east-1 

64 gco --output json storage list 

65 """ 

66 from ..storage import get_storage_manager 

67 

68 formatter = get_output_formatter(config) 

69 try: 

70 buckets = get_storage_manager(config).list_buckets(region=region) 

71 if not buckets: 

72 if config.output_format == "table": 

73 formatter.print_info("No user-facing GCO S3 buckets were discovered") 

74 else: 

75 formatter.print([]) 

76 return 

77 formatter.print( 

78 buckets, 

79 columns=["alias", "scope", "region", "bucket", "purpose", "s3_uri"], 

80 ) 

81 except Exception as exc: 

82 formatter.print_error(f"Failed to discover GCO S3 buckets: {exc}") 

83 sys.exit(1) 

84 

85 

86@storage.command("sync") 

87@click.argument("bucket_alias") 

88@click.argument("local_dir", metavar="LOCAL_PATH") 

89@click.option( 

90 "--direction", 

91 type=click.Choice(["download", "upload"], case_sensitive=False), 

92 default="download", 

93 show_default=True, 

94 help="Transfer direction: S3 to local or local to S3", 

95) 

96@click.option( 

97 "--region", 

98 "-r", 

99 help="Region for the unqualified regional-shared alias", 

100) 

101@click.option( 

102 "--prefix", 

103 default="", 

104 help="Remote S3 key prefix to download from or upload into", 

105) 

106@click.option( 

107 "--dry-run", 

108 is_flag=True, 

109 help="Show a summary of the planned transfer without writing local files or S3 objects", 

110) 

111@click.option( 

112 "--force", 

113 is_flag=True, 

114 help="Transfer every file even when the destination appears current", 

115) 

116@click.option( 

117 "--_gco-storage-root", 

118 "confinement_root", 

119 hidden=True, 

120) 

121@click.option( 

122 "--_gco-storage-root-device", 

123 "confinement_device", 

124 type=int, 

125 hidden=True, 

126) 

127@click.option( 

128 "--_gco-storage-root-inode", 

129 "confinement_inode", 

130 type=int, 

131 hidden=True, 

132) 

133@pass_config 

134def storage_sync( 

135 config: Any, 

136 bucket_alias: str, 

137 local_dir: str, 

138 direction: str, 

139 region: str | None, 

140 prefix: str, 

141 dry_run: bool, 

142 force: bool, 

143 confinement_root: str | None, 

144 confinement_device: int | None, 

145 confinement_inode: int | None, 

146) -> None: 

147 """Sync between a GCO S3 bucket and LOCAL_PATH in one direction. 

148 

149 BUCKET_ALIAS is one of cluster-shared, model-weights, 

150 analytics-studio, or regional-shared:REGION. The unqualified 

151 regional-shared alias can be paired with --region. 

152 

153 The default direction downloads from S3. Use --direction upload to send a 

154 local file or directory to S3. Neither direction deletes destination-only 

155 files or objects; there is no automatic two-way conflict resolution. 

156 

157 Examples: 

158 gco storage sync cluster-shared ./cluster-data 

159 gco storage sync regional-shared:us-east-1 ./regional-data 

160 gco storage sync regional-shared ./regional-data -r us-east-1 

161 gco storage sync model-weights ./models --prefix models/llama3 

162 gco storage sync cluster-shared ./results --direction upload --prefix results 

163 """ 

164 from ..storage import get_storage_manager 

165 

166 formatter = get_output_formatter(config) 

167 try: 

168 if config.output_format == "table": 

169 if dry_run: 

170 action = f"Planning {direction}" 

171 else: 

172 action = "Downloading" if direction == "download" else "Uploading" 

173 formatter.print_info(f"{action} for '{bucket_alias}'...") 

174 

175 with _cooperative_storage_sigterm(): 

176 result = get_storage_manager(config).sync( 

177 bucket_alias, 

178 local_dir, 

179 region=region, 

180 prefix=prefix, 

181 direction=direction, 

182 dry_run=dry_run, 

183 force=force, 

184 confinement_root=confinement_root, 

185 confinement_device=confinement_device, 

186 confinement_inode=confinement_inode, 

187 ) 

188 

189 if config.output_format != "table": 

190 formatter.print(result) 

191 return 

192 

193 result_direction = result["direction"] 

194 transfer_verb = "downloaded" if result_direction == "download" else "uploaded" 

195 if dry_run: 

196 formatter.print_success( 

197 f"Dry run: {result['files_planned']} file(s) " 

198 f"({result['bytes_planned']} bytes) would be {transfer_verb}" 

199 ) 

200 else: 

201 files_key = "files_downloaded" if result_direction == "download" else "files_uploaded" 

202 bytes_key = "bytes_downloaded" if result_direction == "download" else "bytes_uploaded" 

203 formatter.print_success( 

204 f"{transfer_verb.capitalize()} {result[files_key]} file(s) " 

205 f"({result[bytes_key]} bytes)" 

206 ) 

207 formatter.print_info(f"Source: {result['source']}") 

208 formatter.print_info(f"Destination: {result['destination']}") 

209 if result["files_skipped"]: 

210 formatter.print_info(f"Skipped {result['files_skipped']} current file(s)") 

211 except Exception as exc: 

212 formatter.print_error(f"Failed to sync GCO S3 bucket: {exc}") 

213 sys.exit(1)