Coverage for scripts / live_release_validation / checks / inference_scenarios.py: 100.00%
124 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +0000
1"""Scenario-specific phases of the live inference matrix.
3Two scenarios carry steps the classic baseline/HPA legs do not:
5- The Triton leg serves a model repository the harness generates and uploads
6 to the model registry (``gco models upload``), binds the registry URI to the
7 deploy's ``--model-source``, and after the endpoint is proven absent removes
8 the model and requires ``gco models list`` to no longer name it.
9- The Mooncake store leg reconfigures the store after its first invocation
10 (``configure-store``) and requires the role Deployment to roll onto the
11 re-rendered ``mooncake.json`` (a new ``gco.io/mooncake-config-digest`` on
12 the pod template), the endpoint to be promoted back to ``running``, and a
13 second invocation to be answered by the rolled pod.
15The mixin relies on the lifecycle for command execution, checkpointing and the
16DDB wait, exactly like the runtime and inventory mixins.
17"""
19from __future__ import annotations
21import time
22from collections.abc import Callable
23from typing import Any, cast
25from ..models import (
26 MOONCAKE_CONFIG_DIGEST_ANNOTATION,
27 TRITON_MODEL_NAME,
28 TRITON_MODEL_VERSION,
29 write_triton_model_repository,
30)
31from .inference_commands import (
32 EndpointPlan,
33 _last_json_document,
34 build_configure_store_command,
35 build_invoke_command,
36 build_model_delete_command,
37 build_model_list_command,
38 build_model_upload_command,
39 build_model_uri_command,
40 build_structured_decision_command,
41 extract_generated_text,
42 extract_structured_decision,
43)
44from .inference_common import InferenceCommandFailure as _CommandFailure
45from .inference_common import ManagedInferenceValidationError
48def _mapping(value: Any) -> dict[str, Any]:
49 """A dict payload as itself, anything else as an empty dict."""
50 return cast(dict[str, Any], value) if isinstance(value, dict) else {}
53class InferenceScenarioMixin:
54 """Mixin for the model-registry and Mooncake-store phases."""
56 settings: Any
58 def _persist(self) -> None: # pragma: no cover - implemented by lifecycle
59 raise NotImplementedError
61 def _record_failure(
62 self, record: dict[str, Any], stage: str, exc: BaseException
63 ) -> None: # pragma: no cover - implemented by lifecycle
64 raise NotImplementedError
66 def _run_command(
67 self,
68 record: dict[str, Any],
69 stage: str,
70 command: list[str],
71 *,
72 deadline: float | None = None,
73 ) -> str: # pragma: no cover - implemented by lifecycle
74 raise NotImplementedError
76 def _set_phase(
77 self, record: dict[str, Any], phase: str, **values: Any
78 ) -> None: # pragma: no cover - implemented by lifecycle
79 raise NotImplementedError
81 # Supplied by the sibling mixins the lifecycle composes (inventory:
82 # ``_kubectl_json``; runtime: the rest). Declared as annotations only, never
83 # as values, so no base class carries a competing definition and the
84 # resolution cannot depend on base-class order; the same convention the
85 # runtime mixin uses for ``_kubectl_json``.
86 _kubectl_json: Callable[..., Any | None]
87 keep_cluster_tunnel_alive: Callable[..., float]
88 _diagnostics_due: Callable[..., float]
89 _timeout_error: Callable[..., ManagedInferenceValidationError]
90 wait_for_ddb_running: Callable[..., None]
92 def ensure_model_repository(
93 self,
94 plan: EndpointPlan,
95 record: dict[str, Any],
96 ) -> str | None:
97 """Upload the plan's generated model repository and return its registry URI.
99 Only runtimes served from the registry (Triton) have one. The repository
100 is written into the private report directory, published with
101 ``gco models upload`` under a run-owned name, and its URI read back with
102 ``gco models uri`` so the deploy's ``--model-source`` is exactly what the
103 registry resolves. Re-running after a crash re-uploads the same files,
104 which is idempotent; the recorded URI is the one the stored spec is
105 later checked against.
106 """
107 if not plan.runtime.uses_model_registry:
108 return None
109 repository_root = write_triton_model_repository(
110 self.settings.report_dir / "model-repositories" / plan.name
111 )
112 record["model_repository"] = {
113 "registry_model_name": plan.registry_model_name,
114 "phase": "uploading",
115 "model": TRITON_MODEL_NAME,
116 "version": TRITON_MODEL_VERSION,
117 }
118 self._persist()
119 try:
120 self._run_command(
121 record,
122 "model-upload",
123 build_model_upload_command(plan, repository_root),
124 )
125 uri_output = self._run_command(record, "model-uri", build_model_uri_command(plan))
126 except _CommandFailure as exc:
127 self._record_failure(record, "model-upload", exc)
128 raise ManagedInferenceValidationError(
129 "managed inference model repository upload failed; inspect the private checkpoint"
130 ) from None
131 uri_document = _last_json_document(uri_output)
132 model_source = uri_document.get("s3_uri") if isinstance(uri_document, dict) else None
133 if not isinstance(model_source, str) or not model_source.startswith("s3://"):
134 raise ManagedInferenceValidationError(
135 "managed inference model registry did not resolve an S3 URI for the upload"
136 )
137 record["model_source"] = model_source
138 record["model_repository"]["phase"] = "uploaded"
139 self._persist()
140 return model_source
142 def cleanup_model_repository(
143 self,
144 plan: EndpointPlan,
145 record: dict[str, Any],
146 ) -> dict[str, Any] | None:
147 """Remove the plan's registry model and prove the registry no longer lists it.
149 Runs only after the endpoint itself is proven absent, and only when an
150 upload was attempted, so a leg that failed before uploading never
151 touches the registry. A model that was already removed is still proven
152 absent through the listing.
153 """
154 if not plan.runtime.uses_model_registry or "model_repository" not in record:
155 return None
156 repository = record["model_repository"]
157 if repository.get("cleanup") == "absent":
158 return cast(dict[str, Any], repository)
159 try:
160 self._run_command(record, "model-delete", build_model_delete_command(plan))
161 listing_output = self._run_command(record, "model-list", build_model_list_command())
162 except _CommandFailure as exc:
163 self._record_failure(record, "model-cleanup", exc)
164 raise ManagedInferenceValidationError(
165 "managed inference model repository cleanup failed; inspect the private checkpoint"
166 ) from None
167 listing = _last_json_document(listing_output)
168 names = (
169 {
170 entry.get("model_name")
171 for entry in listing
172 if isinstance(entry, dict) and isinstance(entry.get("model_name"), str)
173 }
174 if isinstance(listing, list)
175 else set()
176 )
177 if plan.registry_model_name in names:
178 repository["cleanup"] = "listed-after-delete"
179 self._persist()
180 raise ManagedInferenceValidationError(
181 "managed inference model registry still lists the deleted repository"
182 )
183 repository["cleanup"] = "absent"
184 self._persist()
185 return cast(dict[str, Any], repository)
187 def _store_deployment_rollout(
188 self,
189 plan: EndpointPlan,
190 record: dict[str, Any],
191 *,
192 deadline: float,
193 ) -> tuple[str | None, bool]:
194 """Return the store Deployment's config digest and whether its rollout is complete."""
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 None, False
208 metadata = _mapping(deployment.get("metadata"))
209 spec = _mapping(deployment.get("spec"))
210 status = _mapping(deployment.get("status"))
211 annotations = _mapping(
212 _mapping(_mapping(spec.get("template")).get("metadata")).get("annotations")
213 )
214 digest = annotations.get(MOONCAKE_CONFIG_DIGEST_ANNOTATION)
215 desired = int(spec.get("replicas") or 0)
216 complete = (
217 desired > 0
218 and int(status.get("observedGeneration") or 0) >= int(metadata.get("generation") or 0)
219 and int(status.get("updatedReplicas") or 0) == desired
220 and int(status.get("readyReplicas") or 0) == desired
221 and int(status.get("availableReplicas") or 0) == desired
222 and int(status.get("replicas") or 0) == desired
223 )
224 return (digest if isinstance(digest, str) else None), complete
226 def verify_structured_decision(
227 self,
228 plan: EndpointPlan,
229 record: dict[str, Any],
230 ) -> None:
231 """Ask vLLM's ``/v1/systemone`` one typed question of each kind and check the answers.
233 The request goes through ``gco inference invoke --path /v1/systemone``,
234 so it proves the authenticated inference proxy forwards the path, the
235 pinned image serves it for the matrix model under ``--logprobs-mode
236 raw_logprobs``, and the answer comes back typed. A structured decision
237 reads label probabilities and changes nothing, so a resume may repeat it.
238 """
239 if not plan.structured_decision:
240 return
241 try:
242 output = self._run_command(
243 record,
244 "structured-decision",
245 build_structured_decision_command(self.settings, plan),
246 )
247 evidence = extract_structured_decision(output)
248 except (ManagedInferenceValidationError, _CommandFailure) as exc:
249 self._record_failure(record, "structured-decision", exc)
250 raise ManagedInferenceValidationError(
251 "managed vLLM structured decision failed its response contract"
252 ) from None
253 record["structured_decision"] = {**evidence, "phase": "verified"}
254 self._set_phase(record, "structured-decision-verified")
256 def verify_store_reconfigure(
257 self,
258 plan: EndpointPlan,
259 record: dict[str, Any],
260 ) -> None:
261 """Reconfigure the store and require the role pods to roll onto the new file.
263 The Mooncake connectors read ``mooncake.json`` once at start-up, so a
264 ``configure-store`` that only re-rendered the ConfigMap would leave the
265 running pod on its old settings. The monitor stamps a digest of the
266 rendered file on the pod template and rolls the Deployment when it
267 changes; this phase proves that: the digest observed before the change
268 must differ from the one after it, the rollout must complete with the
269 same replica count, the endpoint must be promoted back to ``running``
270 (``update_spec`` re-enters ``deploying``), and a second invocation must
271 succeed against the rolled pod.
272 """
273 if not plan.mooncake_store:
274 return
275 deadline = time.monotonic() + self.settings.readiness_timeout_seconds
276 digest_before, _complete = self._store_deployment_rollout(plan, record, deadline=deadline)
277 if digest_before is None:
278 raise ManagedInferenceValidationError(
279 "managed store Deployment carries no mooncake config digest before reconfigure"
280 )
281 evidence: dict[str, Any] = {"digest_before": digest_before, "phase": "configuring"}
282 record["store_reconfigure"] = evidence
283 self._persist()
284 try:
285 self._run_command(record, "configure-store", build_configure_store_command(plan))
286 except _CommandFailure as exc:
287 self._record_failure(record, "configure-store", exc)
288 raise ManagedInferenceValidationError(
289 "managed store reconfiguration command failed; inspect the private checkpoint"
290 ) from None
291 # From here the stored spec carries the new buffer size.
292 record["store_reconfigured"] = True
293 evidence["phase"] = "rolling"
294 self._persist()
295 heartbeat_at = time.monotonic()
296 diagnostics_at = time.monotonic()
297 while True:
298 digest, complete = self._store_deployment_rollout(plan, record, deadline=deadline)
299 evidence["digest_after"] = digest
300 evidence["rollout_complete"] = complete
301 self._persist()
302 if digest is not None and digest != digest_before and complete:
303 break
304 if time.monotonic() >= deadline:
305 raise self._timeout_error(
306 plan,
307 record,
308 "managed store Deployment did not roll onto the re-rendered mooncake.json",
309 reason="store-reconfigure-timeout",
310 )
311 heartbeat_at = self.keep_cluster_tunnel_alive(record, heartbeat_at, deadline=deadline)
312 diagnostics_at = self._diagnostics_due(
313 record, plan, diagnostics_at, "store-reconfigure-wait"
314 )
315 time.sleep(
316 min(
317 float(self.settings.poll_interval_seconds),
318 max(0.0, deadline - time.monotonic()),
319 )
320 )
321 evidence["phase"] = "rolled"
322 self._persist()
323 # The spec update parked the endpoint in ``deploying``; the monitor must
324 # promote it again before the rolled pod is proven to answer.
325 self.wait_for_ddb_running(plan, record)
326 try:
327 output = self._run_command(
328 record, "invoke-after-reconfigure", build_invoke_command(self.settings, plan)
329 )
330 generated_text = extract_generated_text(output, plan.runtime.framework)
331 except (ManagedInferenceValidationError, _CommandFailure) as exc:
332 self._record_failure(record, "invoke-after-reconfigure", exc)
333 raise ManagedInferenceValidationError(
334 "managed store invocation after reconfigure failed its response contract"
335 ) from None
336 evidence.update(
337 {
338 "phase": "verified",
339 "invoked_after_rollout": True,
340 "generated_text_length": len(generated_text),
341 }
342 )
343 self._set_phase(record, "store-reconfigured")