Coverage for cli / commands / storage_cmd.py: 100.00%
107 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +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."""
50@storage.command("list")
51@click.option(
52 "--region",
53 "-r",
54 help="Limit regional-bucket discovery to this region",
55)
56@pass_config
57def storage_list(config: Any, region: str | None) -> None:
58 """List deployed buckets and the aliases accepted by `storage sync`.
60 Examples:
61 gco storage list
62 gco storage list --region us-east-1
63 gco --output json storage list
64 """
65 from ..storage import get_storage_manager
67 formatter = get_output_formatter(config)
68 try:
69 buckets = get_storage_manager(config).list_buckets(region=region)
70 if not buckets:
71 if config.output_format == "table":
72 formatter.print_info("No user-facing GCO S3 buckets were discovered")
73 else:
74 formatter.print([])
75 return
76 formatter.print(
77 buckets,
78 columns=["alias", "scope", "region", "bucket", "purpose", "s3_uri"],
79 )
80 except Exception as exc:
81 formatter.print_error(f"Failed to discover GCO S3 buckets: {exc}")
82 sys.exit(1)
85@storage.command("s3-inventory")
86@click.option(
87 "--region",
88 "-r",
89 help="Limit regional entries to this region (global entries are always included)",
90)
91@pass_config
92def storage_s3_inventory(config: Any, region: str | None) -> None:
93 """Describe every S3 bucket this deployment creates, as JSON.
95 Complements `storage list`, which reports only the four user-facing buckets
96 addressable by `storage sync`. This is the full set — the always-on central
97 and per-region shared buckets, the model-weights bucket, the cost-report
98 bucket, the optional analytics Studio bucket, and every server-access-log
99 sink — each with the deployment-contract facts:
101 \b
102 - which stack owns it, and in which region
103 - what it is for, and which object-key prefixes are already reserved
104 - whether job pods can read/write it, and how they discover its name
105 - what teardown does to it (removal policy)
106 - whether it is currently deployed
108 A bucket whose stack is not deployed is listed with status "not-deployed"
109 rather than omitted, so the inventory is complete before a region is rolled
110 out.
112 This inventories buckets and their deployment contract. It is unrelated to
113 the AWS "S3 Inventory" feature, which reports the objects inside a bucket.
115 Examples:
116 gco -o json storage s3-inventory
117 gco -o json storage s3-inventory --region us-east-1
118 gco -o json storage s3-inventory | jq '.summary.pod_writable'
119 gco -o json storage s3-inventory | jq '.buckets[] | select(.pod_access=="read-write")'
120 """
121 from ..storage import get_storage_manager
123 formatter = get_output_formatter(config)
124 try:
125 result = get_storage_manager(config).s3_inventory(region=region)
127 if config.output_format != "table":
128 formatter.print(result)
129 return
131 summary = result["summary"]
132 print(f"\n S3 inventory — project {result['project_name']}, account {result['account']}")
133 print(
134 f" {summary['deployed']}/{summary['total']} deployed"
135 f" ({summary['not_deployed']} not deployed)"
136 )
137 # Nested values render as "<dict>" in table mode, so flatten to the
138 # columns that matter and leave the full shape to -o json.
139 formatter.print(
140 result["buckets"],
141 columns=["id", "role", "region", "bucket", "pod_access", "status"],
142 )
143 if summary["pod_writable"]:
144 print("\n Pod-writable buckets:")
145 for name in summary["pod_writable"]:
146 print(f" {name}")
147 print()
149 except Exception as exc:
150 formatter.print_error(f"Failed to build the GCO S3 inventory: {exc}")
151 sys.exit(1)
154@storage.command("sync")
155@click.argument("bucket_alias")
156@click.argument("local_dir", metavar="LOCAL_PATH")
157@click.option(
158 "--direction",
159 type=click.Choice(["download", "upload"], case_sensitive=False),
160 default="download",
161 show_default=True,
162 help="Transfer direction: S3 to local or local to S3",
163)
164@click.option(
165 "--region",
166 "-r",
167 help="Region for the unqualified regional-shared alias",
168)
169@click.option(
170 "--prefix",
171 default="",
172 help="Remote S3 key prefix to download from or upload into",
173)
174@click.option(
175 "--dry-run",
176 is_flag=True,
177 help="Show a summary of the planned transfer without writing local files or S3 objects",
178)
179@click.option(
180 "--force",
181 is_flag=True,
182 help="Transfer every file even when the destination appears current",
183)
184@click.option(
185 "--_gco-storage-root",
186 "confinement_root",
187 hidden=True,
188)
189@click.option(
190 "--_gco-storage-root-device",
191 "confinement_device",
192 type=int,
193 hidden=True,
194)
195@click.option(
196 "--_gco-storage-root-inode",
197 "confinement_inode",
198 type=int,
199 hidden=True,
200)
201@pass_config
202def storage_sync(
203 config: Any,
204 bucket_alias: str,
205 local_dir: str,
206 direction: str,
207 region: str | None,
208 prefix: str,
209 dry_run: bool,
210 force: bool,
211 confinement_root: str | None,
212 confinement_device: int | None,
213 confinement_inode: int | None,
214) -> None:
215 """Sync between a GCO S3 bucket and LOCAL_PATH in one direction.
217 BUCKET_ALIAS is one of cluster-shared, model-weights,
218 analytics-studio, or regional-shared:REGION. The unqualified
219 regional-shared alias can be paired with --region.
221 The default direction downloads from S3. Use --direction upload to send a
222 local file or directory to S3. Neither direction deletes destination-only
223 files or objects; there is no automatic two-way conflict resolution.
225 Examples:
226 gco storage sync cluster-shared ./cluster-data
227 gco storage sync regional-shared:us-east-1 ./regional-data
228 gco storage sync regional-shared ./regional-data -r us-east-1
229 gco storage sync model-weights ./models --prefix models/llama3
230 gco storage sync cluster-shared ./results --direction upload --prefix results
231 """
232 from ..storage import get_storage_manager
234 formatter = get_output_formatter(config)
235 try:
236 if config.output_format == "table":
237 if dry_run:
238 action = f"Planning {direction}"
239 else:
240 action = "Downloading" if direction == "download" else "Uploading"
241 formatter.print_info(f"{action} for '{bucket_alias}'...")
243 with _cooperative_storage_sigterm():
244 result = get_storage_manager(config).sync(
245 bucket_alias,
246 local_dir,
247 region=region,
248 prefix=prefix,
249 direction=direction,
250 dry_run=dry_run,
251 force=force,
252 confinement_root=confinement_root,
253 confinement_device=confinement_device,
254 confinement_inode=confinement_inode,
255 )
257 if config.output_format != "table":
258 formatter.print(result)
259 return
261 result_direction = result["direction"]
262 transfer_verb = "downloaded" if result_direction == "download" else "uploaded"
263 if dry_run:
264 formatter.print_success(
265 f"Dry run: {result['files_planned']} file(s) "
266 f"({result['bytes_planned']} bytes) would be {transfer_verb}"
267 )
268 else:
269 files_key = "files_downloaded" if result_direction == "download" else "files_uploaded"
270 bytes_key = "bytes_downloaded" if result_direction == "download" else "bytes_uploaded"
271 formatter.print_success(
272 f"{transfer_verb.capitalize()} {result[files_key]} file(s) "
273 f"({result[bytes_key]} bytes)"
274 )
275 formatter.print_info(f"Source: {result['source']}")
276 formatter.print_info(f"Destination: {result['destination']}")
277 if result["files_skipped"]:
278 formatter.print_info(f"Skipped {result['files_skipped']} current file(s)")
279 except Exception as exc:
280 formatter.print_error(f"Failed to sync GCO S3 bucket: {exc}")
281 sys.exit(1)