Coverage for scripts / live_release_validation / checks / inference_runtime.py: 100.00%
223 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""Readiness and HPA stability checks for live inference endpoints."""
3from __future__ import annotations
5import subprocess
6import time
7from collections.abc import Callable
8from typing import Any, cast
10from .inference_common import ManagedInferenceValidationError
12_TUNNEL_HEARTBEAT_INTERVAL_SECONDS = 240.0
15class InferenceRuntimeMixin:
16 """Mixin for DDB readiness, Kubernetes readiness, and HPA stability."""
18 settings: Any
20 def _persist(self) -> None: # pragma: no cover - implemented by lifecycle
21 raise NotImplementedError
23 def _strong_get(
24 self, record: dict[str, Any]
25 ) -> dict[str, Any] | None: # pragma: no cover - implemented by lifecycle
26 raise NotImplementedError
28 def _is_owned(
29 self, item: dict[str, Any]
30 ) -> bool: # pragma: no cover - implemented by lifecycle
31 raise NotImplementedError
33 def _verify_item_contract(
34 self, plan: Any, item: dict[str, Any], record: dict[str, Any]
35 ) -> None: # pragma: no cover - implemented by lifecycle
36 raise NotImplementedError
38 def _set_phase(
39 self, record: dict[str, Any], phase: str, **values: Any
40 ) -> None: # pragma: no cover - implemented by lifecycle
41 raise NotImplementedError
43 _kubectl_json: Callable[..., Any | None]
44 kubectl: Callable[..., tuple[int, str, str]]
46 def keep_cluster_tunnel_alive(
47 self,
48 record: dict[str, Any],
49 last_heartbeat: float,
50 *,
51 deadline: float | None = None,
52 ) -> float:
53 """Send bounded Kubernetes traffic before SSM's idle-session timeout."""
54 now = time.monotonic()
55 if now - last_heartbeat < _TUNNEL_HEARTBEAT_INTERVAL_SECONDS:
56 return last_heartbeat
57 process_timeout = 8.0
58 if deadline is not None:
59 remaining = deadline - now
60 if remaining <= 0:
61 raise ManagedInferenceValidationError(
62 "managed inference tunnel heartbeat deadline expired"
63 )
64 process_timeout = min(process_timeout, remaining)
65 observation: dict[str, Any] = {"started_at_monotonic": now}
66 try:
67 returncode, stdout, stderr = self.kubectl(
68 "--request-timeout=5s",
69 "get",
70 "--raw=/readyz",
71 timeout=process_timeout,
72 )
73 except (OSError, subprocess.TimeoutExpired) as exc:
74 observation.update(
75 {
76 "healthy": False,
77 "error": f"{type(exc).__name__}: {exc}",
78 }
79 )
80 else:
81 observation.update(
82 {
83 "healthy": returncode == 0 and stdout.strip() == "ok",
84 "returncode": returncode,
85 "stderr": stderr[-1000:],
86 }
87 )
88 history = record.setdefault("tunnel_heartbeats", [])
89 if not isinstance(history, list):
90 raise ManagedInferenceValidationError("managed tunnel heartbeat history is invalid")
91 history.append(observation)
92 self._persist()
93 if observation["healthy"] is not True:
94 raise ManagedInferenceValidationError(
95 "managed inference Kubernetes tunnel heartbeat failed"
96 )
97 return time.monotonic()
99 def _wait_for_owned_record(
100 self,
101 plan: Any,
102 record: dict[str, Any],
103 ) -> dict[str, Any]:
104 """Wait for the run-owned DDB record while keeping the tunnel active."""
105 deadline = time.monotonic() + self.settings.readiness_timeout_seconds
106 heartbeat_at = float("-inf")
107 while True:
108 if time.monotonic() >= deadline:
109 raise ManagedInferenceValidationError(
110 "managed inference endpoint ownership did not appear before timeout"
111 )
112 item = self._strong_get(record)
113 if item is not None:
114 if not self._is_owned(item):
115 raise ManagedInferenceValidationError(
116 "managed inference endpoint collision detected; refusing ownership"
117 )
118 self._verify_item_contract(plan, item, record)
119 record["owned"] = True
120 self._set_phase(record, "ownership-confirmed")
121 return item
122 heartbeat_at = self.keep_cluster_tunnel_alive(
123 record,
124 heartbeat_at,
125 deadline=deadline,
126 )
127 remaining = deadline - time.monotonic()
128 if remaining <= 0:
129 raise ManagedInferenceValidationError(
130 "managed inference endpoint ownership did not appear before timeout"
131 )
132 time.sleep(min(float(self.settings.poll_interval_seconds), remaining))
134 def wait_for_ddb_running(self, plan: Any, record: dict[str, Any]) -> None:
135 """Require this run's exact DDB record and running regional observation."""
136 deadline = time.monotonic() + self.settings.readiness_timeout_seconds
137 heartbeat_at = float("-inf")
138 while True:
139 if time.monotonic() >= deadline:
140 raise ManagedInferenceValidationError(
141 "managed inference DDB running state was not observed before timeout"
142 )
143 item = self._strong_get(record)
144 if item is not None:
145 if not self._is_owned(item):
146 raise ManagedInferenceValidationError(
147 "managed inference endpoint ownership changed while waiting"
148 )
149 self._verify_item_contract(plan, item, record)
150 statuses = item.get("region_status")
151 regional = (
152 statuses.get(self.settings.selected_region)
153 if isinstance(statuses, dict)
154 else None
155 )
156 if (
157 item.get("desired_state") == "running"
158 and isinstance(regional, dict)
159 and regional.get("state") == "running"
160 ):
161 self._set_phase(record, "ddb-running")
162 return
163 heartbeat_at = self.keep_cluster_tunnel_alive(
164 record,
165 heartbeat_at,
166 deadline=deadline,
167 )
168 remaining = deadline - time.monotonic()
169 if remaining <= 0:
170 raise ManagedInferenceValidationError(
171 "managed inference DDB running state was not observed before timeout"
172 )
173 time.sleep(min(float(self.settings.poll_interval_seconds), remaining))
175 @staticmethod
176 def _ready_condition(container_statuses: Any) -> bool:
177 return (
178 isinstance(container_statuses, list)
179 and bool(container_statuses)
180 and all(
181 isinstance(status, dict) and status.get("ready") is True
182 for status in container_statuses
183 )
184 )
186 def _deployment_ready_snapshot(
187 self,
188 plan: Any,
189 record: dict[str, Any],
190 expected_replicas: int,
191 *,
192 exact: bool,
193 deadline: float | None = None,
194 ) -> tuple[bool, dict[str, int]]:
195 deployment = self._kubectl_json(
196 record,
197 "get",
198 "deployment",
199 plan.name,
200 "--namespace",
201 self.settings.namespace,
202 "--output",
203 "json",
204 deadline=deadline,
205 )
206 if not isinstance(deployment, dict):
207 return False, {}
208 metadata = deployment.get("metadata")
209 spec = deployment.get("spec")
210 status = deployment.get("status")
211 if (
212 not isinstance(metadata, dict)
213 or not isinstance(spec, dict)
214 or not isinstance(status, dict)
215 ):
216 return False, {}
217 desired = int(spec.get("replicas") or 0)
218 ready = int(status.get("readyReplicas") or 0)
219 available = int(status.get("availableReplicas") or 0)
220 updated = int(status.get("updatedReplicas") or 0)
221 generation = int(metadata.get("generation") or 0)
222 observed = int(status.get("observedGeneration") or 0)
223 replica_match = desired == expected_replicas if exact else desired >= expected_replicas
224 deployment_ready = (
225 replica_match
226 and ready >= expected_replicas
227 and available >= expected_replicas
228 and updated >= expected_replicas
229 and observed >= generation
230 )
232 pods_payload = self._kubectl_json(
233 record,
234 "get",
235 "pods",
236 "--namespace",
237 self.settings.namespace,
238 "--selector",
239 f"app={plan.name}",
240 "--output",
241 "json",
242 deadline=deadline,
243 )
244 items = pods_payload.get("items", []) if isinstance(pods_payload, dict) else []
245 ready_pods = 0
246 if isinstance(items, list):
247 for item in items:
248 pod_status = item.get("status") if isinstance(item, dict) else None
249 if (
250 isinstance(pod_status, dict)
251 and pod_status.get("phase") == "Running"
252 and self._ready_condition(pod_status.get("containerStatuses"))
253 ):
254 ready_pods += 1
255 evidence = {
256 "desired": desired,
257 "ready": ready,
258 "available": available,
259 "updated": updated,
260 "ready_pods": ready_pods,
261 }
262 return deployment_ready and ready_pods >= expected_replicas, evidence
264 def wait_for_kubernetes_ready(self, plan: Any, record: dict[str, Any]) -> None:
265 """Require Deployment convergence and ready Running pods."""
266 deadline = time.monotonic() + self.settings.readiness_timeout_seconds
267 while True:
268 ready, evidence = self._deployment_ready_snapshot(
269 plan,
270 record,
271 plan.replicas,
272 exact=not plan.autoscaling,
273 deadline=deadline,
274 )
275 record["last_readiness"] = evidence
276 self._persist()
277 if ready:
278 self._set_phase(record, "kubernetes-ready")
279 return
280 if time.monotonic() >= deadline:
281 raise ManagedInferenceValidationError(
282 "managed inference Kubernetes readiness was not observed before timeout"
283 )
284 time.sleep(
285 min(
286 float(self.settings.poll_interval_seconds),
287 max(0.0, deadline - time.monotonic()),
288 )
289 )
291 def _hpa_matches(
292 self,
293 plan: Any,
294 record: dict[str, Any],
295 *,
296 deadline: float | None = None,
297 ) -> bool:
298 hpa = self._kubectl_json(
299 record,
300 "get",
301 "horizontalpodautoscaler.autoscaling",
302 plan.name,
303 "--namespace",
304 self.settings.namespace,
305 "--output",
306 "json",
307 deadline=deadline,
308 )
309 if not isinstance(hpa, dict):
310 return False
311 spec = hpa.get("spec")
312 if not isinstance(spec, dict):
313 return False
314 target = spec.get("scaleTargetRef")
315 if not isinstance(target, dict) or target != {
316 "apiVersion": "apps/v1",
317 "kind": "Deployment",
318 "name": plan.name,
319 }:
320 return False
321 if spec.get("minReplicas") != self.settings.hpa_min_replicas:
322 return False
323 if spec.get("maxReplicas") != self.settings.hpa_max_replicas:
324 return False
325 metrics = spec.get("metrics")
326 if not isinstance(metrics, list):
327 return False
328 return any(
329 isinstance(metric, dict)
330 and metric.get("type") == "Resource"
331 and isinstance(metric.get("resource"), dict)
332 and metric["resource"].get("name") == "cpu"
333 and isinstance(metric["resource"].get("target"), dict)
334 and metric["resource"]["target"].get("type") == "Utilization"
335 and metric["resource"]["target"].get("averageUtilization")
336 == self.settings.hpa_cpu_target
337 for metric in metrics
338 )
340 def verify_shared_proxy_autoscaling(self, state: dict[str, Any]) -> None:
341 """Prove the deployed TLS sidecar request and active ContainerResource HPA."""
342 record = state.setdefault(
343 "shared_proxy_autoscaling",
344 {
345 "namespace": "gco-system",
346 "deployment": "inference-proxy",
347 "hpa": "inference-proxy-hpa",
348 "phase": "waiting",
349 "commands": [],
350 },
351 )
352 if not isinstance(record, dict):
353 raise ManagedInferenceValidationError("shared proxy checkpoint evidence is invalid")
354 record["phase"] = "waiting"
355 self._persist()
356 deadline = time.monotonic() + self.settings.hpa_timeout_seconds
357 while True:
358 deployment = self._kubectl_json(
359 record,
360 "get",
361 "deployment",
362 "inference-proxy",
363 "--namespace",
364 "gco-system",
365 "--output",
366 "json",
367 deadline=deadline,
368 )
369 hpa = self._kubectl_json(
370 record,
371 "get",
372 "horizontalpodautoscaler.autoscaling",
373 "inference-proxy-hpa",
374 "--namespace",
375 "gco-system",
376 "--output",
377 "json",
378 deadline=deadline,
379 )
380 observed: dict[str, Any] = {}
381 if isinstance(deployment, dict):
382 deployment_spec = deployment.get("spec")
383 template = (
384 deployment_spec.get("template") if isinstance(deployment_spec, dict) else None
385 )
386 pod_spec = template.get("spec") if isinstance(template, dict) else None
387 containers = pod_spec.get("containers") if isinstance(pod_spec, dict) else None
388 tls_containers = (
389 [
390 item
391 for item in containers
392 if isinstance(item, dict) and item.get("name") == "api-tls-proxy"
393 ]
394 if isinstance(containers, list)
395 else []
396 )
397 if len(tls_containers) == 1:
398 resources = tls_containers[0].get("resources")
399 requests = resources.get("requests") if isinstance(resources, dict) else None
400 if isinstance(requests, dict):
401 observed["tls_cpu_request"] = requests.get("cpu")
402 if isinstance(hpa, dict):
403 metadata_value = hpa.get("metadata")
404 spec_value = hpa.get("spec")
405 status_value = hpa.get("status")
406 metadata = (
407 cast(dict[str, Any], metadata_value) if isinstance(metadata_value, dict) else {}
408 )
409 spec = cast(dict[str, Any], spec_value) if isinstance(spec_value, dict) else {}
410 status = (
411 cast(dict[str, Any], status_value) if isinstance(status_value, dict) else {}
412 )
413 target = spec.get("scaleTargetRef")
414 metrics = spec.get("metrics")
416 def matching_tls_metric(metric: object, *, current: bool) -> bool:
417 if not isinstance(metric, dict) or metric.get("type") != "ContainerResource":
418 return False
419 source = metric.get("containerResource")
420 if not isinstance(source, dict):
421 return False
422 value = source.get("current" if current else "target")
423 return (
424 source.get("name") == "cpu"
425 and source.get("container") == "api-tls-proxy"
426 and isinstance(value, dict)
427 and (current or value.get("type") == "Utilization")
428 )
430 tls_metrics = (
431 [metric for metric in metrics if matching_tls_metric(metric, current=False)]
432 if isinstance(metrics, list)
433 else []
434 )
435 current_metrics = status.get("currentMetrics")
436 active_metrics = (
437 [
438 metric
439 for metric in current_metrics
440 if matching_tls_metric(metric, current=True)
441 ]
442 if isinstance(current_metrics, list)
443 else []
444 )
445 conditions = status.get("conditions")
446 active_conditions = (
447 [
448 condition
449 for condition in conditions
450 if isinstance(condition, dict)
451 and condition.get("type") == "ScalingActive"
452 and condition.get("status") == "True"
453 ]
454 if isinstance(conditions, list)
455 else []
456 )
457 tls_target: object = None
458 if len(tls_metrics) == 1:
459 # matching_tls_metric admitted this metric only after proving
460 # containerResource.target is a Utilization object.
461 tls_target = tls_metrics[0]["containerResource"]["target"].get(
462 "averageUtilization"
463 )
464 observed.update(
465 {
466 "target_matches": target
467 == {
468 "apiVersion": "apps/v1",
469 "kind": "Deployment",
470 "name": "inference-proxy",
471 },
472 "tls_metric_count": len(tls_metrics),
473 "tls_cpu_target": tls_target,
474 "active_tls_metric_count": len(active_metrics),
475 "scaling_active": bool(active_conditions),
476 "scaling_active_reason": (
477 active_conditions[0].get("reason") if active_conditions else None
478 ),
479 "observed_generation_current": int(status.get("observedGeneration") or 0)
480 >= int(metadata.get("generation") or 0),
481 }
482 )
483 record["last_observed"] = observed
484 self._persist()
485 if (
486 observed.get("tls_cpu_request") == self.settings.proxy_tls_cpu_request
487 and observed.get("target_matches") is True
488 and observed.get("tls_metric_count") == 1
489 and observed.get("tls_cpu_target") == self.settings.proxy_tls_cpu_target
490 and observed.get("active_tls_metric_count") == 1
491 and observed.get("scaling_active") is True
492 and observed.get("observed_generation_current") is True
493 ):
494 record["phase"] = "verified"
495 record["expected"] = {
496 "tls_cpu_request": self.settings.proxy_tls_cpu_request,
497 "tls_cpu_target": self.settings.proxy_tls_cpu_target,
498 }
499 self._persist()
500 return
501 if time.monotonic() >= deadline:
502 raise ManagedInferenceValidationError(
503 "shared inference-proxy TLS autoscaling contract was not active before timeout"
504 )
505 time.sleep(
506 min(
507 float(self.settings.poll_interval_seconds),
508 max(0.0, deadline - time.monotonic()),
509 )
510 )
512 def verify_hpa_stability(self, plan: Any, record: dict[str, Any]) -> None:
513 """Prove HPA target/bounds and two full monitor intervals at two replicas."""
514 deadline = time.monotonic() + self.settings.hpa_timeout_seconds
515 while not self._hpa_matches(plan, record, deadline=deadline):
516 if time.monotonic() >= deadline:
517 raise ManagedInferenceValidationError(
518 "managed inference HPA contract was not observed before timeout"
519 )
520 time.sleep(
521 min(
522 float(self.settings.poll_interval_seconds),
523 max(0.0, deadline - time.monotonic()),
524 )
525 )
526 self._set_phase(record, "hpa-verified")
528 while True:
529 ready, evidence = self._deployment_ready_snapshot(
530 plan,
531 record,
532 self.settings.hpa_min_replicas,
533 exact=True,
534 deadline=deadline,
535 )
536 record["last_hpa_replica_observation"] = evidence
537 self._persist()
538 if ready:
539 break
540 if time.monotonic() >= deadline:
541 raise ManagedInferenceValidationError(
542 "managed inference HPA did not reach two ready replicas before timeout"
543 )
544 time.sleep(
545 min(
546 float(self.settings.poll_interval_seconds),
547 max(0.0, deadline - time.monotonic()),
548 )
549 )
551 observations = [record["last_hpa_replica_observation"]]
552 monitor_interval = float(self.settings.monitor_interval_seconds)
553 for _ in range(self.settings.hpa_stability_intervals):
554 if deadline - time.monotonic() < monitor_interval:
555 raise ManagedInferenceValidationError(
556 "managed inference HPA stability exceeded its phase deadline"
557 )
558 time.sleep(monitor_interval)
559 ready, evidence = self._deployment_ready_snapshot(
560 plan,
561 record,
562 self.settings.hpa_min_replicas,
563 exact=True,
564 deadline=deadline,
565 )
566 observations.append(evidence)
567 if not ready:
568 record["hpa_stability_observations"] = observations
569 self._persist()
570 raise ManagedInferenceValidationError(
571 "managed inference HPA replicas did not remain stable"
572 )
573 record["hpa_stability_observations"] = observations
574 self._set_phase(record, "hpa-stable")