Coverage for cli / models.py: 100.00%
148 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +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:
163 continue
164 identifier = {"Key": key}
165 if version_id is not None:
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 response = s3.delete_objects(Bucket=bucket, Delete={"Objects": batch})
173 errors = response.get("Errors", []) if isinstance(response, dict) else []
174 errors = [error for error in errors if isinstance(error, dict)]
176 for identifier in batch:
177 failed = any(
178 error.get("Key") == identifier["Key"]
179 and (
180 error.get("VersionId") is None
181 or error.get("VersionId") == identifier.get("VersionId")
182 )
183 for error in errors
184 )
185 if not failed:
186 deleted_keys.add(identifier["Key"])
188 for error in errors:
189 key = error.get("Key", "<unknown key>")
190 version_id = error.get("VersionId")
191 target = f"{key} (version {version_id})" if version_id else str(key)
192 code = error.get("Code", "UnknownError")
193 message = error.get("Message", "no error message")
194 deletion_errors.append(f"{target}: {code}: {message}")
196 if deletion_errors:
197 details = "; ".join(deletion_errors)
198 raise RuntimeError(
199 f"Failed to delete {len(deletion_errors)} model object version(s): {details}"
200 )
202 return len(deleted_keys)
205class RegionalBucketManager:
206 """Uploads local files to a region's general-purpose regional bucket.
208 Mirrors :class:`ModelManager` but targets the per-region general-purpose
209 regional bucket (CloudFormation-generated name) instead of the central
210 model bucket. The bucket name is always resolved from the *target
211 region's own* SSM parameter store, never the global region's or any other
212 region's, so an upload only ever writes to the bucket that lives in the
213 region the caller named.
214 """
216 def __init__(self, config: GCOConfig | None = None):
217 self.config = config or get_config()
219 def _get_bucket_name(self, region: str) -> str:
220 """Resolve the regional bucket name from the target region's SSM store.
222 Reads ``/<project_name>/regional-shared-bucket/name`` from the
223 parameter store in ``region``. The regional bucket is always
224 provisioned, so this parameter is present once the region's stack is
225 deployed. A missing parameter means the region has not been deployed
226 yet and is treated as a hard "bucket not found" failure.
227 """
228 from gco.services.aws_ssm import get_ssm_parameter_optional
229 from gco.stacks.constants import regional_shared_ssm_parameter_prefix
231 name = get_ssm_parameter_optional(
232 f"{regional_shared_ssm_parameter_prefix(self.config.project_name)}/name",
233 region=region,
234 )
235 if not name:
236 raise RuntimeError(
237 f"Regional bucket not found in region '{region}'. Deploy that "
238 f"region's stack first with 'gco stacks deploy'."
239 )
240 return name
242 def _get_s3_client(self, region: str) -> Any:
243 """Get an S3 client scoped to the target region."""
244 return boto3.client("s3", region_name=region)
246 def upload(
247 self,
248 local_path: str,
249 region: str,
250 *,
251 prefix: str = "uploads",
252 ) -> dict[str, Any]:
253 """
254 Upload local files or a directory to a region's regional bucket.
256 Args:
257 local_path: Local file or directory path
258 region: Target region whose regional bucket receives the objects
259 prefix: S3 prefix for uploaded objects (default: "uploads")
261 Returns:
262 Upload result with the region, bucket, S3 URI, and file count
264 Raises:
265 RuntimeError: If the target region's bucket cannot be resolved (no
266 objects are written) or if an object fails mid-upload (the
267 upload stops and the offending object is named).
268 FileNotFoundError: If ``local_path`` does not exist.
269 """
270 local = Path(local_path)
271 if not local.exists():
272 raise FileNotFoundError(f"Path not found: {local_path}")
274 # Resolve the bucket before writing anything so an undeployed region
275 # fails fast without partial uploads.
276 bucket = self._get_bucket_name(region)
277 s3 = self._get_s3_client(region)
278 uploaded = 0
280 files: list[tuple[Path, str]]
281 if local.is_file():
282 files = [(local, local.name)]
283 else:
284 files = []
285 for root, _dirs, names in os.walk(local):
286 for fname in names:
287 walk_path = Path(root) / fname
288 rel = walk_path.relative_to(local)
289 files.append((walk_path, str(rel)))
291 for file_path, relative in files:
292 key = f"{prefix}/{relative}"
293 try:
294 s3.upload_file(str(file_path), bucket, key)
295 except Exception as e:
296 raise RuntimeError(
297 f"Upload did not complete: failed to write object "
298 f"'s3://{bucket}/{key}' to region '{region}': {e}"
299 ) from e
300 uploaded += 1
302 s3_uri = f"s3://{bucket}/{prefix}"
303 return {
304 "region": region,
305 "bucket": bucket,
306 "s3_uri": s3_uri,
307 "files_uploaded": uploaded,
308 }
310 def populate_kv_cache(
311 self,
312 local_path: str,
313 region: str,
314 endpoint_name: str,
315 ) -> dict[str, Any]:
316 """Upload data into an endpoint's Mooncake KV-cache cold tier.
318 Writes ``local_path`` to the region's general-purpose bucket under the
319 cold-tier key prefix the per-region monitor reads from for this endpoint
320 (``mooncake-kv/<endpoint_name>/``), so an endpoint deployed with the
321 cold tier enabled warm-starts its prefix cache from the uploaded
322 objects. Resolution and upload mechanics are exactly those of
323 :meth:`upload`; the returned mapping additionally carries the endpoint
324 name.
326 Args:
327 local_path: Local file or directory to upload.
328 region: Region whose general-purpose bucket backs the cold tier.
329 endpoint_name: The endpoint whose cold-tier prefix receives the data.
331 Returns:
332 The :meth:`upload` result with an added ``endpoint`` key.
333 """
334 from gco.stacks.constants import MOONCAKE_COLD_TIER_KEY_PREFIX
336 prefix = f"{MOONCAKE_COLD_TIER_KEY_PREFIX}/{endpoint_name}"
337 result = self.upload(local_path, region, prefix=prefix)
338 result["endpoint"] = endpoint_name
339 return result
342def get_model_manager(config: GCOConfig | None = None) -> ModelManager:
343 """Factory function for ModelManager."""
344 return ModelManager(config)
347def get_regional_bucket_manager(
348 config: GCOConfig | None = None,
349) -> RegionalBucketManager:
350 """Factory function for RegionalBucketManager."""
351 return RegionalBucketManager(config)