Coverage for cli/models.py: 96.97%
150 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-30 21:22 +0000
1"""
2Model weight management for GCO CLI.
4Provides functionality to upload, list, and manage model weights
5in the central S3 model bucket. Models uploaded here are automatically
6available to inference endpoints across all regions via init container sync.
7"""
9from __future__ import annotations
11import logging
12import os
13from pathlib import Path
14from typing import Any
16import boto3
18from .config import GCOConfig, get_config
20logger = logging.getLogger(__name__)
23class ModelManager:
24 """Manages model weights in the central S3 bucket."""
26 def __init__(self, config: GCOConfig | None = None):
27 self.config = config or get_config()
28 self._bucket_name: str | None = None
30 def _get_bucket_name(self) -> str:
31 """Discover the model bucket name from SSM."""
32 if self._bucket_name:
33 return self._bucket_name
35 from gco.services.aws_ssm import get_ssm_parameter
37 try:
38 self._bucket_name = get_ssm_parameter(
39 f"/{self.config.project_name}/model-bucket-name",
40 region=self.config.global_region,
41 )
42 return self._bucket_name
43 except Exception as e:
44 raise RuntimeError(
45 "Model bucket not found. Deploy the global stack first "
46 "with 'gco stacks deploy gco-global'."
47 ) from e
49 def _get_s3_client(self) -> Any:
50 """Get S3 client for the global region."""
51 return boto3.client("s3", region_name=self.config.global_region)
53 def upload(
54 self,
55 local_path: str,
56 model_name: str,
57 prefix: str = "models",
58 ) -> dict[str, Any]:
59 """
60 Upload model weights to S3.
62 Args:
63 local_path: Local file or directory path
64 model_name: Name for the model in the bucket
65 prefix: S3 prefix (default: "models")
67 Returns:
68 Upload result with S3 URI and file count
69 """
70 bucket = self._get_bucket_name()
71 s3 = self._get_s3_client()
72 s3_prefix = f"{prefix}/{model_name}"
74 local = Path(local_path)
75 uploaded = 0
77 if local.is_file():
78 key = f"{s3_prefix}/{local.name}"
79 s3.upload_file(str(local), bucket, key)
80 uploaded = 1
81 elif local.is_dir():
82 for root, _dirs, files in os.walk(local):
83 for fname in files:
84 file_path = Path(root) / fname
85 relative = file_path.relative_to(local)
86 key = f"{s3_prefix}/{relative}"
87 s3.upload_file(str(file_path), bucket, key)
88 uploaded += 1
89 else:
90 raise FileNotFoundError(f"Path not found: {local_path}")
92 s3_uri = f"s3://{bucket}/{s3_prefix}"
93 return {
94 "model_name": model_name,
95 "s3_uri": s3_uri,
96 "bucket": bucket,
97 "prefix": s3_prefix,
98 "files_uploaded": uploaded,
99 }
101 def list_models(self, prefix: str = "models") -> list[dict[str, Any]]:
102 """List all models in the bucket."""
103 bucket = self._get_bucket_name()
104 s3 = self._get_s3_client()
106 # List top-level "directories" under the prefix
107 response = s3.list_objects_v2(
108 Bucket=bucket,
109 Prefix=f"{prefix}/",
110 Delimiter="/",
111 )
113 models = []
114 for cp in response.get("CommonPrefixes", []):
115 model_prefix = cp["Prefix"]
116 model_name = model_prefix.rstrip("/").split("/")[-1]
118 # Get total size and file count
119 total_size = 0
120 file_count = 0
121 paginator = s3.get_paginator("list_objects_v2")
122 for page in paginator.paginate(Bucket=bucket, Prefix=model_prefix):
123 for obj in page.get("Contents", []):
124 total_size += obj.get("Size", 0)
125 file_count += 1
127 models.append(
128 {
129 "model_name": model_name,
130 "s3_uri": f"s3://{bucket}/{model_prefix.rstrip('/')}",
131 "files": file_count,
132 "total_size_gb": round(total_size / (1024**3), 2),
133 }
134 )
136 return models
138 def get_model_uri(self, model_name: str, prefix: str = "models") -> str:
139 """Get the S3 URI for a model."""
140 bucket = self._get_bucket_name()
141 return f"s3://{bucket}/{prefix}/{model_name}"
143 def delete_model(self, model_name: str, prefix: str = "models") -> int:
144 """Delete every version and delete marker for a model prefix.
146 The central model bucket is versioned. Deleting only the current
147 objects creates delete markers and leaves prior versions behind, which
148 can prevent later bucket removal and retain model data unexpectedly.
149 """
150 bucket = self._get_bucket_name()
151 s3 = self._get_s3_client()
152 s3_prefix = f"{prefix}/{model_name}/"
154 deleted_keys: set[str] = set()
155 deletion_errors: list[str] = []
156 paginator = s3.get_paginator("list_object_versions")
157 for page in paginator.paginate(Bucket=bucket, Prefix=s3_prefix):
158 versioned_objects = []
159 for item in [*page.get("Versions", []), *page.get("DeleteMarkers", [])]:
160 key = item.get("Key")
161 version_id = item.get("VersionId")
162 if not key: 162 ↛ 163line 162 didn't jump to line 163 because the condition on line 162 was never true
163 continue
164 identifier = {"Key": key}
165 if version_id is not None: 165 ↛ 167line 165 didn't jump to line 167 because the condition on line 165 was always true
166 identifier["VersionId"] = version_id
167 versioned_objects.append(identifier)
169 # S3 accepts at most 1,000 identifiers per DeleteObjects request.
170 for start in range(0, len(versioned_objects), 1000):
171 batch = versioned_objects[start : start + 1000]
172 if not batch: 172 ↛ 173line 172 didn't jump to line 173 because the condition on line 172 was never true
173 continue
175 response = s3.delete_objects(Bucket=bucket, Delete={"Objects": batch})
176 errors = response.get("Errors", []) if isinstance(response, dict) else []
177 errors = [error for error in errors if isinstance(error, dict)]
179 for identifier in batch:
180 failed = any(
181 error.get("Key") == identifier["Key"]
182 and (
183 error.get("VersionId") is None
184 or error.get("VersionId") == identifier.get("VersionId")
185 )
186 for error in errors
187 )
188 if not failed:
189 deleted_keys.add(identifier["Key"])
191 for error in errors:
192 key = error.get("Key", "<unknown key>")
193 version_id = error.get("VersionId")
194 target = f"{key} (version {version_id})" if version_id else str(key)
195 code = error.get("Code", "UnknownError")
196 message = error.get("Message", "no error message")
197 deletion_errors.append(f"{target}: {code}: {message}")
199 if deletion_errors:
200 details = "; ".join(deletion_errors)
201 raise RuntimeError(
202 f"Failed to delete {len(deletion_errors)} model object version(s): {details}"
203 )
205 return len(deleted_keys)
208class RegionalBucketManager:
209 """Uploads local files to a region's general-purpose regional bucket.
211 Mirrors :class:`ModelManager` but targets the per-region
212 ``gco-regional-shared-<account>-<region>`` bucket instead of the central
213 model bucket. The bucket name is always resolved from the *target
214 region's own* SSM parameter store, never the global region's or any other
215 region's, so an upload only ever writes to the bucket that lives in the
216 region the caller named.
217 """
219 def __init__(self, config: GCOConfig | None = None):
220 self.config = config or get_config()
222 def _get_bucket_name(self, region: str) -> str:
223 """Resolve the regional bucket name from the target region's SSM store.
225 Reads ``/<project_name>/regional-shared-bucket/name`` from the
226 parameter store in ``region``. The regional bucket is always
227 provisioned, so this parameter is present once the region's stack is
228 deployed. A missing parameter means the region has not been deployed
229 yet and is treated as a hard "bucket not found" failure.
230 """
231 from gco.services.aws_ssm import get_ssm_parameter_optional
232 from gco.stacks.constants import regional_shared_ssm_parameter_prefix
234 name = get_ssm_parameter_optional(
235 f"{regional_shared_ssm_parameter_prefix(self.config.project_name)}/name",
236 region=region,
237 )
238 if not name:
239 raise RuntimeError(
240 f"Regional bucket not found in region '{region}'. Deploy that "
241 f"region's stack first with 'gco stacks deploy'."
242 )
243 return name
245 def _get_s3_client(self, region: str) -> Any:
246 """Get an S3 client scoped to the target region."""
247 return boto3.client("s3", region_name=region)
249 def upload(
250 self,
251 local_path: str,
252 region: str,
253 *,
254 prefix: str = "uploads",
255 ) -> dict[str, Any]:
256 """
257 Upload local files or a directory to a region's regional bucket.
259 Args:
260 local_path: Local file or directory path
261 region: Target region whose regional bucket receives the objects
262 prefix: S3 prefix for uploaded objects (default: "uploads")
264 Returns:
265 Upload result with the region, bucket, S3 URI, and file count
267 Raises:
268 RuntimeError: If the target region's bucket cannot be resolved (no
269 objects are written) or if an object fails mid-upload (the
270 upload stops and the offending object is named).
271 FileNotFoundError: If ``local_path`` does not exist.
272 """
273 local = Path(local_path)
274 if not local.exists():
275 raise FileNotFoundError(f"Path not found: {local_path}")
277 # Resolve the bucket before writing anything so an undeployed region
278 # fails fast without partial uploads.
279 bucket = self._get_bucket_name(region)
280 s3 = self._get_s3_client(region)
281 uploaded = 0
283 files: list[tuple[Path, str]]
284 if local.is_file():
285 files = [(local, local.name)]
286 else:
287 files = []
288 for root, _dirs, names in os.walk(local):
289 for fname in names:
290 walk_path = Path(root) / fname
291 rel = walk_path.relative_to(local)
292 files.append((walk_path, str(rel)))
294 for file_path, relative in files:
295 key = f"{prefix}/{relative}"
296 try:
297 s3.upload_file(str(file_path), bucket, key)
298 except Exception as e:
299 raise RuntimeError(
300 f"Upload did not complete: failed to write object "
301 f"'s3://{bucket}/{key}' to region '{region}': {e}"
302 ) from e
303 uploaded += 1
305 s3_uri = f"s3://{bucket}/{prefix}"
306 return {
307 "region": region,
308 "bucket": bucket,
309 "s3_uri": s3_uri,
310 "files_uploaded": uploaded,
311 }
313 def populate_kv_cache(
314 self,
315 local_path: str,
316 region: str,
317 endpoint_name: str,
318 ) -> dict[str, Any]:
319 """Upload data into an endpoint's Mooncake KV-cache cold tier.
321 Writes ``local_path`` to the region's general-purpose bucket under the
322 cold-tier key prefix the per-region monitor reads from for this endpoint
323 (``mooncake-kv/<endpoint_name>/``), so an endpoint deployed with the
324 cold tier enabled warm-starts its prefix cache from the uploaded
325 objects. Resolution and upload mechanics are exactly those of
326 :meth:`upload`; the returned mapping additionally carries the endpoint
327 name.
329 Args:
330 local_path: Local file or directory to upload.
331 region: Region whose general-purpose bucket backs the cold tier.
332 endpoint_name: The endpoint whose cold-tier prefix receives the data.
334 Returns:
335 The :meth:`upload` result with an added ``endpoint`` key.
336 """
337 from gco.stacks.constants import MOONCAKE_COLD_TIER_KEY_PREFIX
339 prefix = f"{MOONCAKE_COLD_TIER_KEY_PREFIX}/{endpoint_name}"
340 result = self.upload(local_path, region, prefix=prefix)
341 result["endpoint"] = endpoint_name
342 return result
345def get_model_manager(config: GCOConfig | None = None) -> ModelManager:
346 """Factory function for ModelManager."""
347 return ModelManager(config)
350def get_regional_bucket_manager(
351 config: GCOConfig | None = None,
352) -> RegionalBucketManager:
353 """Factory function for RegionalBucketManager."""
354 return RegionalBucketManager(config)