Coverage for cli / commands / storage_cmd.py: 100.00%
108 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""Commands for discovering and syncing GCO S3 buckets."""
3import signal
4import sys
5import threading
6from collections.abc import Iterator
7from contextlib import contextmanager
8from types import FrameType
9from typing import Any
11import click
13from ..config import GCOConfig
14from ..output import get_output_formatter
16pass_config = click.make_pass_decorator(GCOConfig, ensure=True)
19class _StorageSyncTerminated(RuntimeError):
20 """Raised on SIGTERM so managed S3 transfers can unwind cooperatively."""
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
30 previous_handler = signal.getsignal(signal.SIGTERM)
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 )
37 signal.signal(signal.SIGTERM, terminate_handler)
38 try:
39 yield
40 finally:
41 signal.signal(signal.SIGTERM, previous_handler)
44@click.group()
45@pass_config
46def storage(config: Any) -> None:
47 """Discover and sync user-facing GCO S3 buckets."""
48 pass
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`.
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
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)
86@storage.command("s3-inventory")
87@click.option(
88 "--region",
89 "-r",
90 help="Limit regional entries to this region (global entries are always included)",
91)
92@pass_config
93def storage_s3_inventory(config: Any, region: str | None) -> None:
94 """Describe every S3 bucket this deployment creates, as JSON.
96 Complements `storage list`, which reports only the four user-facing buckets
97 addressable by `storage sync`. This is the full set — the always-on central
98 and per-region shared buckets, the model-weights bucket, the cost-report
99 bucket, the optional analytics Studio bucket, and every server-access-log
100 sink — each with the deployment-contract facts:
102 \b
103 - which stack owns it, and in which region
104 - what it is for, and which object-key prefixes are already reserved
105 - whether job pods can read/write it, and how they discover its name
106 - what teardown does to it (removal policy)
107 - whether it is currently deployed
109 A bucket whose stack is not deployed is listed with status "not-deployed"
110 rather than omitted, so the inventory is complete before a region is rolled
111 out.
113 This inventories buckets and their deployment contract. It is unrelated to
114 the AWS "S3 Inventory" feature, which reports the objects inside a bucket.
116 Examples:
117 gco -o json storage s3-inventory
118 gco -o json storage s3-inventory --region us-east-1
119 gco -o json storage s3-inventory | jq '.summary.pod_writable'
120 gco -o json storage s3-inventory | jq '.buckets[] | select(.pod_access=="read-write")'
121 """
122 from ..storage import get_storage_manager
124 formatter = get_output_formatter(config)
125 try:
126 result = get_storage_manager(config).s3_inventory(region=region)
128 if config.output_format != "table":
129 formatter.print(result)
130 return
132 summary = result["summary"]
133 print(f"\n S3 inventory — project {result['project_name']}, account {result['account']}")
134 print(
135 f" {summary['deployed']}/{summary['total']} deployed"
136 f" ({summary['not_deployed']} not deployed)"
137 )
138 # Nested values render as "<dict>" in table mode, so flatten to the
139 # columns that matter and leave the full shape to -o json.
140 formatter.print(
141 result["buckets"],
142 columns=["id", "role", "region", "bucket", "pod_access", "status"],
143 )
144 if summary["pod_writable"]:
145 print("\n Pod-writable buckets:")
146 for name in summary["pod_writable"]:
147 print(f" {name}")
148 print()
150 except Exception as exc:
151 formatter.print_error(f"Failed to build the GCO S3 inventory: {exc}")
152 sys.exit(1)
155@storage.command("sync")
156@click.argument("bucket_alias")
157@click.argument("local_dir", metavar="LOCAL_PATH")
158@click.option(
159 "--direction",
160 type=click.Choice(["download", "upload"], case_sensitive=False),
161 default="download",
162 show_default=True,
163 help="Transfer direction: S3 to local or local to S3",
164)
165@click.option(
166 "--region",
167 "-r",
168 help="Region for the unqualified regional-shared alias",
169)
170@click.option(
171 "--prefix",
172 default="",
173 help="Remote S3 key prefix to download from or upload into",
174)
175@click.option(
176 "--dry-run",
177 is_flag=True,
178 help="Show a summary of the planned transfer without writing local files or S3 objects",
179)
180@click.option(
181 "--force",
182 is_flag=True,
183 help="Transfer every file even when the destination appears current",
184)
185@click.option(
186 "--_gco-storage-root",
187 "confinement_root",
188 hidden=True,
189)
190@click.option(
191 "--_gco-storage-root-device",
192 "confinement_device",
193 type=int,
194 hidden=True,
195)
196@click.option(
197 "--_gco-storage-root-inode",
198 "confinement_inode",
199 type=int,
200 hidden=True,
201)
202@pass_config
203def storage_sync(
204 config: Any,
205 bucket_alias: str,
206 local_dir: str,
207 direction: str,
208 region: str | None,
209 prefix: str,
210 dry_run: bool,
211 force: bool,
212 confinement_root: str | None,
213 confinement_device: int | None,
214 confinement_inode: int | None,
215) -> None:
216 """Sync between a GCO S3 bucket and LOCAL_PATH in one direction.
218 BUCKET_ALIAS is one of cluster-shared, model-weights,
219 analytics-studio, or regional-shared:REGION. The unqualified
220 regional-shared alias can be paired with --region.
222 The default direction downloads from S3. Use --direction upload to send a
223 local file or directory to S3. Neither direction deletes destination-only
224 files or objects; there is no automatic two-way conflict resolution.
226 Examples:
227 gco storage sync cluster-shared ./cluster-data
228 gco storage sync regional-shared:us-east-1 ./regional-data
229 gco storage sync regional-shared ./regional-data -r us-east-1
230 gco storage sync model-weights ./models --prefix models/llama3
231 gco storage sync cluster-shared ./results --direction upload --prefix results
232 """
233 from ..storage import get_storage_manager
235 formatter = get_output_formatter(config)
236 try:
237 if config.output_format == "table":
238 if dry_run:
239 action = f"Planning {direction}"
240 else:
241 action = "Downloading" if direction == "download" else "Uploading"
242 formatter.print_info(f"{action} for '{bucket_alias}'...")
244 with _cooperative_storage_sigterm():
245 result = get_storage_manager(config).sync(
246 bucket_alias,
247 local_dir,
248 region=region,
249 prefix=prefix,
250 direction=direction,
251 dry_run=dry_run,
252 force=force,
253 confinement_root=confinement_root,
254 confinement_device=confinement_device,
255 confinement_inode=confinement_inode,
256 )
258 if config.output_format != "table":
259 formatter.print(result)
260 return
262 result_direction = result["direction"]
263 transfer_verb = "downloaded" if result_direction == "download" else "uploaded"
264 if dry_run:
265 formatter.print_success(
266 f"Dry run: {result['files_planned']} file(s) "
267 f"({result['bytes_planned']} bytes) would be {transfer_verb}"
268 )
269 else:
270 files_key = "files_downloaded" if result_direction == "download" else "files_uploaded"
271 bytes_key = "bytes_downloaded" if result_direction == "download" else "bytes_uploaded"
272 formatter.print_success(
273 f"{transfer_verb.capitalize()} {result[files_key]} file(s) "
274 f"({result[bytes_key]} bytes)"
275 )
276 formatter.print_info(f"Source: {result['source']}")
277 formatter.print_info(f"Destination: {result['destination']}")
278 if result["files_skipped"]:
279 formatter.print_info(f"Skipped {result['files_skipped']} current file(s)")
280 except Exception as exc:
281 formatter.print_error(f"Failed to sync GCO S3 bucket: {exc}")
282 sys.exit(1)