Coverage for gco / services / inference_store.py: 100.00%
272 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""
2DynamoDB-backed store for inference endpoint state.
4Provides lifecycle-fenced CRUD operations for inference endpoints. Regional
5monitors use the immutable ``lifecycle_id`` and deletion generation to ensure
6that stale writers cannot mutate a replacement endpoint or recreate a record
7after terminal deletion.
8"""
10from __future__ import annotations
12import logging
13import os
14import secrets
15from datetime import UTC, datetime
16from typing import Any
18import boto3
19from botocore.exceptions import ClientError
21logger = logging.getLogger(__name__)
23DEFAULT_TABLE_NAME = "gco-inference-endpoints"
26def _utc_now_iso() -> str:
27 return datetime.now(UTC).isoformat()
30def _new_lifecycle_token() -> str:
31 """Return a cryptographically random immutable lifecycle token."""
32 return secrets.token_hex(32)
35def _validate_endpoint_spec(spec: dict[str, Any]) -> None:
36 """Reject endpoint shapes the reconciler cannot safely materialize."""
37 if not isinstance(spec, dict):
38 raise ValueError("Endpoint spec must be a mapping")
39 if "mooncake" in spec and "canary" in spec:
40 raise ValueError("Endpoint spec cannot combine 'mooncake' and 'canary' blocks")
43class InferenceEndpointStore:
44 """DynamoDB store for inference endpoint desired state."""
46 def __init__(self, table_name: str | None = None, region: str | None = None):
47 self.table_name = table_name or os.getenv(
48 "INFERENCE_ENDPOINTS_TABLE_NAME", DEFAULT_TABLE_NAME
49 )
50 self._region = region or os.getenv("DYNAMODB_REGION") or os.getenv("REGION", "us-east-1")
51 self._dynamodb = boto3.resource("dynamodb", region_name=self._region)
52 self._table = self._dynamodb.Table(self.table_name)
54 def create_endpoint(
55 self,
56 endpoint_name: str,
57 spec: dict[str, Any],
58 target_regions: list[str],
59 namespace: str = "gco-inference",
60 labels: dict[str, str] | None = None,
61 created_by: str | None = None,
62 ) -> dict[str, Any]:
63 """Create one endpoint incarnation with immutable lifecycle identity."""
64 _validate_endpoint_spec(spec)
65 now = _utc_now_iso()
66 regions = list(dict.fromkeys(target_regions))
67 lifecycle_id = _new_lifecycle_token()
68 region_generations = {region: _new_lifecycle_token() for region in regions}
69 item: dict[str, Any] = {
70 "endpoint_name": endpoint_name,
71 "lifecycle_id": lifecycle_id,
72 "desired_state": "deploying",
73 "target_regions": regions,
74 # Append-only membership for this lifecycle. Region removal changes
75 # target_regions but never this authoritative cleanup set.
76 "cleanup_regions": list(regions),
77 # Membership changes rotate only the affected Region's token. A
78 # terminal acknowledgement from an earlier remove/re-add cycle can
79 # therefore never suppress cleanup for the current membership.
80 "region_generations": region_generations,
81 "namespace": namespace,
82 "spec": _serialize_for_dynamo(spec),
83 "ingress_path": f"/inference/{endpoint_name}",
84 "created_at": now,
85 "updated_at": now,
86 "region_status": {},
87 }
88 if labels:
89 item["labels"] = labels
90 if created_by:
91 item["created_by"] = created_by
93 try:
94 self._table.put_item(
95 Item=item,
96 ConditionExpression="attribute_not_exists(endpoint_name)",
97 )
98 except ClientError as e:
99 if e.response["Error"]["Code"] == "ConditionalCheckFailedException":
100 raise ValueError(f"Endpoint '{endpoint_name}' already exists") from e
101 raise
102 return item
104 def get_endpoint(
105 self,
106 endpoint_name: str,
107 *,
108 consistent_read: bool = False,
109 ) -> dict[str, Any] | None:
110 """Get an endpoint by name, optionally using a strong read."""
111 response = self._table.get_item(
112 Key={"endpoint_name": endpoint_name},
113 ConsistentRead=consistent_read,
114 )
115 item = response.get("Item")
116 return _deserialize_from_dynamo(item) if isinstance(item, dict) else None
118 def list_endpoints(
119 self,
120 desired_state: str | None = None,
121 target_region: str | None = None,
122 ) -> list[dict[str, Any]]:
123 """List all endpoints, optionally filtered."""
124 response = self._table.scan()
125 items = [_deserialize_from_dynamo(i) for i in response.get("Items", [])]
126 if desired_state:
127 items = [i for i in items if i.get("desired_state") == desired_state]
128 if target_region:
129 items = [i for i in items if target_region in i.get("target_regions", [])]
130 return sorted(items, key=lambda x: x.get("created_at", ""), reverse=True)
132 def ensure_lifecycle_metadata(self, endpoint: dict[str, Any]) -> dict[str, Any] | None:
133 """Conditionally backfill immutable lifecycle metadata on a legacy record.
135 The migration is derived from one strong snapshot and conditioned on
136 its ``updated_at`` value. Existing lifecycle and Region tokens are
137 preserved, while current targets, prior cleanup members, and regions
138 with historical status are unioned into the authoritative cleanup set.
139 A concurrent mutation wins and makes this call return ``None`` so the
140 caller retries from a fresh snapshot instead of overwriting it.
141 """
142 endpoint_name = endpoint.get("endpoint_name")
143 updated_at = endpoint.get("updated_at")
144 if not isinstance(endpoint_name, str) or not endpoint_name:
145 raise ValueError("Endpoint lifecycle migration requires an endpoint name")
146 if not isinstance(updated_at, str) or not updated_at:
147 raise ValueError(f"Endpoint '{endpoint_name}' has no conditional migration timestamp")
149 raw_status = endpoint.get("region_status")
150 status_regions = list(raw_status) if isinstance(raw_status, dict) else []
151 regions = list(
152 dict.fromkeys(
153 region
154 for source in (
155 endpoint.get("cleanup_regions"),
156 endpoint.get("target_regions"),
157 status_regions,
158 )
159 if isinstance(source, list)
160 for region in source
161 if isinstance(region, str) and region
162 )
163 )
164 lifecycle_value = endpoint.get("lifecycle_id")
165 lifecycle_id = (
166 lifecycle_value
167 if isinstance(lifecycle_value, str) and lifecycle_value
168 else _new_lifecycle_token()
169 )
170 raw_generations = endpoint.get("region_generations")
171 existing_generations = raw_generations if isinstance(raw_generations, dict) else {}
172 region_generations = {
173 region: (
174 existing_generations[region]
175 if isinstance(existing_generations.get(region), str)
176 and existing_generations[region]
177 else _new_lifecycle_token()
178 )
179 for region in regions
180 }
182 metadata_complete = (
183 endpoint.get("lifecycle_id") == lifecycle_id
184 and endpoint.get("cleanup_regions") == regions
185 and endpoint.get("region_generations") == region_generations
186 )
187 if metadata_complete:
188 return endpoint
190 try:
191 response = self._table.update_item(
192 Key={"endpoint_name": endpoint_name},
193 UpdateExpression=(
194 "SET lifecycle_id = if_not_exists(lifecycle_id, :lifecycle_id), "
195 "cleanup_regions = :cleanup_regions, "
196 "region_generations = :region_generations, updated_at = :updated_at"
197 ),
198 ExpressionAttributeValues={
199 ":lifecycle_id": lifecycle_id,
200 ":cleanup_regions": regions,
201 ":region_generations": region_generations,
202 ":updated_at": _utc_now_iso(),
203 ":expected_updated_at": updated_at,
204 },
205 ConditionExpression=(
206 "attribute_exists(endpoint_name) AND updated_at = :expected_updated_at"
207 ),
208 ReturnValues="ALL_NEW",
209 )
210 return _deserialize_from_dynamo(response.get("Attributes", {}))
211 except ClientError as error:
212 if error.response["Error"]["Code"] == "ConditionalCheckFailedException":
213 return None
214 raise
216 @staticmethod
217 def _conditioned_identity(
218 expected_label: tuple[str, str] | None,
219 expected_lifecycle_id: str | None,
220 ) -> tuple[str, dict[str, str], dict[str, Any]]:
221 condition = "attribute_exists(endpoint_name)"
222 names: dict[str, str] = {}
223 values: dict[str, Any] = {}
224 if expected_label is not None:
225 label_name, label_value = expected_label
226 if not label_name or not label_value:
227 raise ValueError("Expected endpoint label name and value must be non-empty")
228 condition += " AND labels.#expected_label = :expected_label_value"
229 names["#expected_label"] = label_name
230 values[":expected_label_value"] = label_value
231 if expected_lifecycle_id is not None:
232 if not expected_lifecycle_id:
233 raise ValueError("Expected lifecycle id must be non-empty")
234 condition += " AND lifecycle_id = :expected_lifecycle_id"
235 values[":expected_lifecycle_id"] = expected_lifecycle_id
236 return condition, names, values
238 def update_desired_state(
239 self,
240 endpoint_name: str,
241 desired_state: str,
242 *,
243 expected_label: tuple[str, str] | None = None,
244 expected_lifecycle_id: str | None = None,
245 expected_desired_state: str | None = None,
246 ) -> dict[str, Any] | None:
247 """Conditionally update desired state without reviving deletion.
249 Ordinary transitions require an immutable lifecycle identity and may
250 only mutate a non-deleted record. The first transition to ``deleted``
251 atomically creates an immutable deletion generation and snapshots the
252 lifecycle's append-only cleanup regions; repeated deletes retain both.
253 """
254 if desired_state != "deleted" and not expected_lifecycle_id:
255 raise ValueError("Ordinary state updates require an expected lifecycle id")
256 condition, names, identity_values = self._conditioned_identity(
257 expected_label, expected_lifecycle_id
258 )
259 values: dict[str, Any] = {
260 ":s": desired_state,
261 ":u": _utc_now_iso(),
262 **identity_values,
263 }
264 update_expression = "SET desired_state = :s, updated_at = :u"
265 if desired_state == "deleted":
266 values[":deletion_generation"] = _new_lifecycle_token()
267 update_expression += (
268 ", deletion_generation = if_not_exists("
269 "deletion_generation, :deletion_generation), "
270 "deletion_regions = if_not_exists(deletion_regions, cleanup_regions)"
271 )
272 else:
273 condition += " AND desired_state <> :deleted"
274 values[":deleted"] = "deleted"
275 if expected_desired_state is not None:
276 condition += " AND desired_state = :expected_desired_state"
277 values[":expected_desired_state"] = expected_desired_state
278 kwargs: dict[str, Any] = {}
279 if names:
280 kwargs["ExpressionAttributeNames"] = names
281 try:
282 response = self._table.update_item(
283 Key={"endpoint_name": endpoint_name},
284 UpdateExpression=update_expression,
285 ExpressionAttributeValues=values,
286 ConditionExpression=condition,
287 ReturnValues="ALL_NEW",
288 **kwargs,
289 )
290 return _deserialize_from_dynamo(response.get("Attributes", {}))
291 except ClientError as e:
292 if e.response["Error"]["Code"] == "ConditionalCheckFailedException":
293 return None
294 raise
296 def start_endpoint(self, endpoint_name: str) -> dict[str, Any] | None:
297 """Start only a stopped endpoint; deletion is terminal for this lifecycle."""
298 current = self.get_endpoint(endpoint_name, consistent_read=True)
299 if current is None:
300 return None
301 migrated = self.ensure_lifecycle_metadata(current)
302 if migrated is None:
303 return None
304 current = migrated
305 lifecycle_id = current.get("lifecycle_id")
306 if not isinstance(lifecycle_id, str) or not lifecycle_id:
307 raise ValueError(f"Endpoint '{endpoint_name}' has no lifecycle identity")
308 state = current.get("desired_state")
309 if state == "deleted":
310 raise ValueError(
311 f"Endpoint '{endpoint_name}' is deleted; wait for purge and redeploy it with "
312 "'gco inference deploy'."
313 )
314 if state != "stopped":
315 raise ValueError(
316 f"Endpoint '{endpoint_name}' is in '{state}' state; only stopped endpoints "
317 "can be started."
318 )
319 return self.update_desired_state(
320 endpoint_name,
321 "running",
322 expected_lifecycle_id=lifecycle_id,
323 expected_desired_state="stopped",
324 )
326 def update_spec(
327 self,
328 endpoint_name: str,
329 spec: dict[str, Any],
330 *,
331 expected_lifecycle_id: str,
332 expected_updated_at: str | None = None,
333 ) -> dict[str, Any] | None:
334 """Conditionally update a live lifecycle's spec and trigger reconciliation."""
335 _validate_endpoint_spec(spec)
336 if not expected_lifecycle_id:
337 raise ValueError("Spec updates require an expected lifecycle id")
338 condition = (
339 "attribute_exists(endpoint_name) AND lifecycle_id = :expected_lifecycle_id "
340 "AND desired_state <> :deleted"
341 )
342 values: dict[str, Any] = {
343 ":s": _serialize_for_dynamo(spec),
344 ":u": _utc_now_iso(),
345 ":ds": "deploying",
346 ":deleted": "deleted",
347 ":expected_lifecycle_id": expected_lifecycle_id,
348 }
349 if expected_updated_at is not None:
350 condition += " AND updated_at = :expected_updated_at"
351 values[":expected_updated_at"] = expected_updated_at
352 try:
353 response = self._table.update_item(
354 Key={"endpoint_name": endpoint_name},
355 UpdateExpression="SET spec = :s, updated_at = :u, desired_state = :ds",
356 ExpressionAttributeValues=values,
357 ConditionExpression=condition,
358 ReturnValues="ALL_NEW",
359 )
360 return _deserialize_from_dynamo(response.get("Attributes", {}))
361 except ClientError as e:
362 if e.response["Error"]["Code"] == "ConditionalCheckFailedException":
363 return None
364 raise
366 def update_target_regions(
367 self,
368 endpoint_name: str,
369 target_regions: list[str],
370 cleanup_regions: list[str],
371 region_generations: dict[str, str],
372 *,
373 expected_lifecycle_id: str,
374 expected_updated_at: str,
375 ) -> dict[str, Any] | None:
376 """Conditionally update membership without dropping cleanup authority."""
377 normalized_targets = list(dict.fromkeys(target_regions))
378 requested_cleanup = list(dict.fromkeys(cleanup_regions))
379 current = self.get_endpoint(endpoint_name, consistent_read=True)
380 if (
381 current is None
382 or current.get("lifecycle_id") != expected_lifecycle_id
383 or current.get("updated_at") != expected_updated_at
384 or current.get("desired_state") == "deleted"
385 ):
386 return None
388 historical_cleanup = list(
389 dict.fromkeys(current.get("cleanup_regions") or current.get("target_regions") or [])
390 )
391 normalized_cleanup = list(
392 dict.fromkeys([*historical_cleanup, *requested_cleanup, *normalized_targets])
393 )
394 current_generations = current.get("region_generations")
395 merged_generations = (
396 dict(current_generations) if isinstance(current_generations, dict) else {}
397 )
398 merged_generations.update(region_generations)
399 for region in normalized_cleanup:
400 token = merged_generations.get(region)
401 if not isinstance(token, str) or not token:
402 merged_generations[region] = _new_lifecycle_token()
403 merged_generations = {region: merged_generations[region] for region in normalized_cleanup}
404 try:
405 response = self._table.update_item(
406 Key={"endpoint_name": endpoint_name},
407 UpdateExpression=(
408 "SET target_regions = :targets, cleanup_regions = :cleanup, "
409 "region_generations = :region_generations, updated_at = :u"
410 ),
411 ExpressionAttributeValues={
412 ":targets": normalized_targets,
413 ":cleanup": normalized_cleanup,
414 ":region_generations": merged_generations,
415 ":u": _utc_now_iso(),
416 ":expected_lifecycle_id": expected_lifecycle_id,
417 ":expected_updated_at": expected_updated_at,
418 ":deleted": "deleted",
419 },
420 ConditionExpression=(
421 "attribute_exists(endpoint_name) "
422 "AND lifecycle_id = :expected_lifecycle_id "
423 "AND updated_at = :expected_updated_at "
424 "AND desired_state <> :deleted"
425 ),
426 ReturnValues="ALL_NEW",
427 )
428 return _deserialize_from_dynamo(response.get("Attributes", {}))
429 except ClientError as e:
430 if e.response["Error"]["Code"] == "ConditionalCheckFailedException":
431 return None
432 raise
434 def update_region_status(
435 self,
436 endpoint_name: str,
437 region: str,
438 state: str,
439 replicas_ready: int = 0,
440 replicas_desired: int = 0,
441 error: str | None = None,
442 extra: dict[str, Any] | None = None,
443 *,
444 expected_lifecycle_id: str | None = None,
445 expected_region_generation: str | None = None,
446 expected_deletion_generation: str | None = None,
447 ) -> bool:
448 """Conditionally write one regional observation without upsert risk."""
449 status_value: dict[str, Any] = {
450 "state": state,
451 "replicas_ready": replicas_ready,
452 "replicas_desired": replicas_desired,
453 "last_sync": _utc_now_iso(),
454 }
455 if expected_lifecycle_id is not None:
456 status_value["lifecycle_id"] = expected_lifecycle_id
457 if expected_region_generation is not None:
458 status_value["region_generation"] = expected_region_generation
459 if expected_deletion_generation is not None:
460 status_value["deletion_generation"] = expected_deletion_generation
461 if error:
462 status_value["error"] = error
463 if extra:
464 status_value.update(extra)
466 condition = "attribute_exists(endpoint_name)"
467 names: dict[str, str] = {"#r": region}
468 values: dict[str, Any] = {":s": status_value, ":u": _utc_now_iso()}
469 if expected_lifecycle_id is not None:
470 condition += " AND lifecycle_id = :expected_lifecycle_id"
471 values[":expected_lifecycle_id"] = expected_lifecycle_id
472 if expected_region_generation is not None:
473 condition += " AND region_generations.#r = :expected_region_generation"
474 values[":expected_region_generation"] = expected_region_generation
475 if expected_deletion_generation is not None:
476 condition += (
477 " AND desired_state = :deleted "
478 "AND deletion_generation = :expected_deletion_generation "
479 "AND (attribute_not_exists(region_status.#r.#state) "
480 "OR region_status.#r.#state <> :terminal_deleted)"
481 )
482 names["#state"] = "state"
483 values[":deleted"] = "deleted"
484 values[":terminal_deleted"] = "deleted"
485 values[":expected_deletion_generation"] = expected_deletion_generation
486 else:
487 # Ordinary observations are never allowed to overwrite a terminal
488 # deletion record. Lifecycle/Region tokens fence replacement and
489 # membership races; this state predicate closes the remaining
490 # same-lifecycle window between the first delete transition and
491 # the final generation-scoped cleanup acknowledgement.
492 condition += " AND desired_state <> :deleted"
493 values[":deleted"] = "deleted"
494 try:
495 self._table.update_item(
496 Key={"endpoint_name": endpoint_name},
497 UpdateExpression="SET region_status.#r = :s, updated_at = :u",
498 ExpressionAttributeNames=names,
499 ExpressionAttributeValues=values,
500 ConditionExpression=condition,
501 )
502 return True
503 except ClientError as e:
504 if e.response["Error"]["Code"] == "ConditionalCheckFailedException":
505 logger.info(
506 "Skipped stale regional status for %s/%s after lifecycle change",
507 endpoint_name,
508 region,
509 )
510 return False
511 logger.error(
512 "Failed to update region status for %s/%s: %s",
513 endpoint_name,
514 region,
515 e,
516 )
517 return False
519 def delete_endpoint(
520 self,
521 endpoint_name: str,
522 *,
523 expected_updated_at: str | None = None,
524 expected_lifecycle_id: str | None = None,
525 expected_deletion_generation: str | None = None,
526 ) -> bool:
527 """Delete only the freshly verified endpoint deletion generation."""
528 condition = "attribute_exists(endpoint_name)"
529 values: dict[str, Any] = {}
530 if expected_updated_at is not None:
531 condition += " AND desired_state = :deleted AND updated_at = :expected_updated_at"
532 values.update({":deleted": "deleted", ":expected_updated_at": expected_updated_at})
533 if expected_lifecycle_id is not None:
534 condition += " AND lifecycle_id = :expected_lifecycle_id"
535 values[":expected_lifecycle_id"] = expected_lifecycle_id
536 if expected_deletion_generation is not None:
537 condition += " AND deletion_generation = :expected_deletion_generation"
538 values[":expected_deletion_generation"] = expected_deletion_generation
539 kwargs: dict[str, Any] = {
540 "Key": {"endpoint_name": endpoint_name},
541 "ConditionExpression": condition,
542 }
543 if values:
544 kwargs["ExpressionAttributeValues"] = values
545 try:
546 self._table.delete_item(**kwargs)
547 return True
548 except ClientError as e:
549 if e.response["Error"]["Code"] == "ConditionalCheckFailedException":
550 return False
551 raise
553 def scale_endpoint(
554 self,
555 endpoint_name: str,
556 replicas: int,
557 *,
558 expected_lifecycle_id: str,
559 ) -> dict[str, Any] | None:
560 """Update classic static replicas only on a live, non-Mooncake lifecycle."""
561 if not expected_lifecycle_id:
562 raise ValueError("Scaling requires an expected lifecycle id")
563 condition = (
564 "attribute_exists(endpoint_name) "
565 "AND lifecycle_id = :expected_lifecycle_id "
566 "AND desired_state <> :deleted "
567 "AND attribute_not_exists(#spec.#mooncake) "
568 "AND (attribute_not_exists(#spec.#autoscaling.#enabled) "
569 "OR #spec.#autoscaling.#enabled = :false)"
570 )
571 values: dict[str, Any] = {
572 ":r": replicas,
573 ":u": _utc_now_iso(),
574 ":false": False,
575 ":deleted": "deleted",
576 ":expected_lifecycle_id": expected_lifecycle_id,
577 }
578 try:
579 response = self._table.update_item(
580 Key={"endpoint_name": endpoint_name},
581 UpdateExpression="SET #spec.replicas = :r, updated_at = :u",
582 ExpressionAttributeNames={
583 "#spec": "spec",
584 "#autoscaling": "autoscaling",
585 "#enabled": "enabled",
586 "#mooncake": "mooncake",
587 },
588 ExpressionAttributeValues=values,
589 ConditionExpression=condition,
590 ReturnValues="ALL_NEW",
591 )
592 return _deserialize_from_dynamo(response.get("Attributes", {}))
593 except ClientError as e:
594 if e.response["Error"]["Code"] == "ConditionalCheckFailedException":
595 return None
596 raise
599def _serialize_for_dynamo(obj: Any) -> Any:
600 """Convert Python objects to DynamoDB-compatible types recursively."""
601 if isinstance(obj, dict):
602 return {k: _serialize_for_dynamo(v) for k, v in obj.items()}
603 if isinstance(obj, list):
604 return [_serialize_for_dynamo(i) for i in obj]
605 if isinstance(obj, (int, float)):
606 return str(obj) if isinstance(obj, float) else obj
607 return obj
610def _deserialize_from_dynamo(item: dict[str, Any]) -> dict[str, Any]:
611 """Convert a DynamoDB item back to plain Python types recursively."""
612 from decimal import Decimal
614 def convert(v: Any) -> Any:
615 if isinstance(v, Decimal):
616 return int(v) if v == int(v) else float(v)
617 if isinstance(v, dict):
618 return {k: convert(val) for k, val in v.items()}
619 if isinstance(v, list):
620 return [convert(i) for i in v]
621 return v
623 result: dict[str, Any] = convert(item)
624 return result
627def get_inference_endpoint_store() -> InferenceEndpointStore:
628 """Factory function for InferenceEndpointStore."""
629 return InferenceEndpointStore()