Coverage for gco/services/inference_store.py: 93.63%
117 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-30 21:22 +0000
1"""
2DynamoDB-backed store for inference endpoint state.
4Provides CRUD operations for inference endpoints. The inference_monitor
5in each regional cluster polls this table to reconcile desired state
6with actual Kubernetes resources.
7"""
9from __future__ import annotations
11import logging
12import os
13from datetime import UTC, datetime
14from typing import Any
16import boto3
17from botocore.exceptions import ClientError
19logger = logging.getLogger(__name__)
21DEFAULT_TABLE_NAME = "gco-inference-endpoints"
24def _utc_now_iso() -> str:
25 return datetime.now(UTC).isoformat()
28def _validate_endpoint_spec(spec: dict[str, Any]) -> None:
29 """Reject endpoint shapes the reconciler cannot safely materialize."""
30 if not isinstance(spec, dict): 30 ↛ 31line 30 didn't jump to line 31 because the condition on line 30 was never true
31 raise ValueError("Endpoint spec must be a mapping")
32 if "mooncake" in spec and "canary" in spec:
33 raise ValueError("Endpoint spec cannot combine 'mooncake' and 'canary' blocks")
36class InferenceEndpointStore:
37 """DynamoDB store for inference endpoint desired state."""
39 def __init__(self, table_name: str | None = None, region: str | None = None):
40 self.table_name = table_name or os.getenv(
41 "INFERENCE_ENDPOINTS_TABLE_NAME", DEFAULT_TABLE_NAME
42 )
43 self._region = region or os.getenv("DYNAMODB_REGION") or os.getenv("REGION", "us-east-1")
44 self._dynamodb = boto3.resource("dynamodb", region_name=self._region)
45 self._table = self._dynamodb.Table(self.table_name)
47 def create_endpoint(
48 self,
49 endpoint_name: str,
50 spec: dict[str, Any],
51 target_regions: list[str],
52 namespace: str = "gco-inference",
53 labels: dict[str, str] | None = None,
54 created_by: str | None = None,
55 ) -> dict[str, Any]:
56 """Create a new inference endpoint entry."""
57 _validate_endpoint_spec(spec)
58 now = _utc_now_iso()
59 ingress_path = f"/inference/{endpoint_name}"
61 item: dict[str, Any] = {
62 "endpoint_name": endpoint_name,
63 "desired_state": "deploying",
64 "target_regions": target_regions,
65 "namespace": namespace,
66 "spec": _serialize_for_dynamo(spec),
67 "ingress_path": ingress_path,
68 "created_at": now,
69 "updated_at": now,
70 "region_status": {},
71 }
72 if labels:
73 item["labels"] = labels
74 if created_by:
75 item["created_by"] = created_by
77 try:
78 self._table.put_item(
79 Item=item,
80 ConditionExpression="attribute_not_exists(endpoint_name)",
81 )
82 except ClientError as e:
83 if e.response["Error"]["Code"] == "ConditionalCheckFailedException":
84 raise ValueError(f"Endpoint '{endpoint_name}' already exists") from e
85 raise
87 return item
89 def get_endpoint(self, endpoint_name: str) -> dict[str, Any] | None:
90 """Get an endpoint by name."""
91 response = self._table.get_item(Key={"endpoint_name": endpoint_name})
92 item = response.get("Item")
93 if item:
94 return _deserialize_from_dynamo(item)
95 return None
97 def list_endpoints(
98 self,
99 desired_state: str | None = None,
100 target_region: str | None = None,
101 ) -> list[dict[str, Any]]:
102 """List all endpoints, optionally filtered."""
103 response = self._table.scan()
104 items = [_deserialize_from_dynamo(i) for i in response.get("Items", [])]
106 if desired_state:
107 items = [i for i in items if i.get("desired_state") == desired_state]
108 if target_region:
109 items = [i for i in items if target_region in i.get("target_regions", [])]
111 return sorted(items, key=lambda x: x.get("created_at", ""), reverse=True)
113 def update_desired_state(self, endpoint_name: str, desired_state: str) -> dict[str, Any] | None:
114 """Update the desired state of an endpoint."""
115 try:
116 response = self._table.update_item(
117 Key={"endpoint_name": endpoint_name},
118 UpdateExpression="SET desired_state = :s, updated_at = :u",
119 ExpressionAttributeValues={
120 ":s": desired_state,
121 ":u": _utc_now_iso(),
122 },
123 ConditionExpression="attribute_exists(endpoint_name)",
124 ReturnValues="ALL_NEW",
125 )
126 return _deserialize_from_dynamo(response.get("Attributes", {}))
127 except ClientError as e:
128 if e.response["Error"]["Code"] == "ConditionalCheckFailedException": 128 ↛ 130line 128 didn't jump to line 130 because the condition on line 128 was always true
129 return None
130 raise
132 def update_spec(self, endpoint_name: str, spec: dict[str, Any]) -> dict[str, Any] | None:
133 """Update the spec of an endpoint (triggers re-reconciliation)."""
134 _validate_endpoint_spec(spec)
135 try:
136 response = self._table.update_item(
137 Key={"endpoint_name": endpoint_name},
138 UpdateExpression="SET spec = :s, updated_at = :u, desired_state = :ds",
139 ExpressionAttributeValues={
140 ":s": _serialize_for_dynamo(spec),
141 ":u": _utc_now_iso(),
142 ":ds": "deploying",
143 },
144 ConditionExpression="attribute_exists(endpoint_name)",
145 ReturnValues="ALL_NEW",
146 )
147 return _deserialize_from_dynamo(response.get("Attributes", {}))
148 except ClientError as e:
149 if e.response["Error"]["Code"] == "ConditionalCheckFailedException": 149 ↛ 151line 149 didn't jump to line 151 because the condition on line 149 was always true
150 return None
151 raise
153 def update_region_status(
154 self,
155 endpoint_name: str,
156 region: str,
157 state: str,
158 replicas_ready: int = 0,
159 replicas_desired: int = 0,
160 error: str | None = None,
161 extra: dict[str, Any] | None = None,
162 ) -> None:
163 """Update the sync status for a specific region.
165 ``extra`` carries optional, additive sub-status that a richer endpoint
166 shape needs — for example a role-keyed breakdown of a split topology
167 (``{"roles": {...}, "store": {...}}``). Its keys are merged into the
168 stored status alongside the flat fields, so a consumer that only reads
169 the flat shape is unaffected.
170 """
171 status_value: dict[str, Any] = {
172 "state": state,
173 "replicas_ready": replicas_ready,
174 "replicas_desired": replicas_desired,
175 "last_sync": _utc_now_iso(),
176 }
177 if error:
178 status_value["error"] = error
179 if extra:
180 status_value.update(extra)
182 try:
183 self._table.update_item(
184 Key={"endpoint_name": endpoint_name},
185 UpdateExpression="SET region_status.#r = :s, updated_at = :u",
186 ExpressionAttributeNames={"#r": region},
187 ExpressionAttributeValues={
188 ":s": status_value,
189 ":u": _utc_now_iso(),
190 },
191 )
192 except ClientError as e:
193 logger.error(
194 "Failed to update region status for %s/%s: %s",
195 endpoint_name,
196 region,
197 e,
198 )
200 def delete_endpoint(self, endpoint_name: str) -> bool:
201 """Delete an endpoint record entirely."""
202 try:
203 self._table.delete_item(
204 Key={"endpoint_name": endpoint_name},
205 ConditionExpression="attribute_exists(endpoint_name)",
206 )
207 return True
208 except ClientError as e:
209 if e.response["Error"]["Code"] == "ConditionalCheckFailedException": 209 ↛ 211line 209 didn't jump to line 211 because the condition on line 209 was always true
210 return False
211 raise
213 def scale_endpoint(self, endpoint_name: str, replicas: int) -> dict[str, Any] | None:
214 """Update the replica count in the spec."""
215 try:
216 response = self._table.update_item(
217 Key={"endpoint_name": endpoint_name},
218 UpdateExpression="SET spec.replicas = :r, updated_at = :u",
219 ExpressionAttributeValues={
220 ":r": replicas,
221 ":u": _utc_now_iso(),
222 },
223 ConditionExpression="attribute_exists(endpoint_name)",
224 ReturnValues="ALL_NEW",
225 )
226 return _deserialize_from_dynamo(response.get("Attributes", {}))
227 except ClientError as e:
228 if e.response["Error"]["Code"] == "ConditionalCheckFailedException": 228 ↛ 230line 228 didn't jump to line 230 because the condition on line 228 was always true
229 return None
230 raise
233def _serialize_for_dynamo(obj: Any) -> Any:
234 """Convert Python objects to DynamoDB-compatible types.
236 Recurses through nested dicts and lists, so an arbitrarily deep
237 configuration block carried on an endpoint spec (for example a nested
238 topology/store/transfer block, including list-valued sub-fields) is
239 converted in place rather than only at the top level.
241 Type handling:
242 - Integers are kept as integers, so whole-number counts are preserved.
243 - Floats are rendered as their decimal-string form, avoiding binary-float
244 rounding on store/reload.
245 - Strings, booleans, and None pass through unchanged. Byte-size values
246 authored as base-10 integer decimal strings therefore stay strings and
247 are never routed through a float, so they reload exactly as written
248 without float-to-Decimal coercion.
249 """
250 if isinstance(obj, dict):
251 return {k: _serialize_for_dynamo(v) for k, v in obj.items()}
252 if isinstance(obj, list):
253 return [_serialize_for_dynamo(i) for i in obj]
254 if isinstance(obj, (int, float)):
255 return str(obj) if isinstance(obj, float) else obj
256 return obj
259def _deserialize_from_dynamo(item: dict[str, Any]) -> dict[str, Any]:
260 """Convert a DynamoDB item back to plain Python types.
262 Recurses through nested dicts and lists to mirror the nested structure
263 produced by :func:`_serialize_for_dynamo`. DynamoDB returns numbers as
264 Decimal: whole values become int (so integer counts survive the
265 round-trip) and fractional values become float. Strings are left
266 untouched, so a byte-size value stored as a decimal string returns as the
267 same string.
268 """
269 from decimal import Decimal
271 def convert(v: Any) -> Any:
272 if isinstance(v, Decimal):
273 return int(v) if v == int(v) else float(v)
274 if isinstance(v, dict):
275 return {k: convert(val) for k, val in v.items()}
276 if isinstance(v, list):
277 return [convert(i) for i in v]
278 return v
280 result: dict[str, Any] = convert(item)
281 return result
284def get_inference_endpoint_store() -> InferenceEndpointStore:
285 """Factory function for InferenceEndpointStore."""
286 return InferenceEndpointStore()