Coverage for scripts / live_release_validation / checks / inference_commands.py: 100.00%
180 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +0000
1"""Plans and argv-only ``gco`` commands for the live inference matrix.
3Everything here is pure: the scenario plan derived from the runtime matrix,
4the exact command arrays the lifecycle runs (deploy, probes, invoke, delete,
5the model-registry upload/lookup/removal the Triton scenario needs, and the
6``configure-store`` the Mooncake store scenario applies), and the response
7parsing that holds each runtime to its own schema. The lifecycle in
8``inference.py`` executes them; nothing here touches a subprocess.
9"""
11from __future__ import annotations
13import json
14import sys
15from dataclasses import dataclass
16from pathlib import Path
17from typing import Any
19from ..models import (
20 INFERENCE_OWNER_LABEL,
21 MOONCAKE_STORE_RECONFIGURED_LOCAL_BUFFER_BYTES,
22 STRUCTURED_DECISION_PATH,
23 STRUCTURED_DECISION_QUESTIONS,
24 TRITON_INPUT_DATA,
25 EndpointRole,
26 InferenceRuntimeSpec,
27 RunSettings,
28)
29from .inference_common import ManagedInferenceValidationError
31OWNER_LABEL = INFERENCE_OWNER_LABEL
34@dataclass(frozen=True)
35class EndpointPlan:
36 """One immutable framework/role scenario in the sequential matrix."""
38 ordinal: int
39 role: EndpointRole
40 name: str
41 replicas: int
42 autoscaling: bool
43 runtime: InferenceRuntimeSpec
45 @property
46 def mooncake_store(self) -> bool:
47 """Whether this scenario is the Mooncake shared KV-cache store."""
48 return self.role == "store"
50 @property
51 def structured_decision(self) -> bool:
52 """Whether this scenario also calls vLLM's ``/v1/systemone``."""
53 return self.role in self.runtime.structured_decision_roles
55 @property
56 def registry_model_name(self) -> str | None:
57 """The model-registry name this scenario uploads its weights under, if any."""
58 if not self.runtime.uses_model_registry:
59 return None
60 return f"{self.name}-repo"
62 def private_dict(self, owner_nonce: str) -> dict[str, Any]:
63 return {
64 "ordinal": self.ordinal,
65 "framework": self.runtime.framework,
66 "role": self.role,
67 "name": self.name,
68 "replicas": self.replicas,
69 "autoscaling": self.autoscaling,
70 "registry_model_name": self.registry_model_name,
71 "owner_nonce": owner_nonce,
72 }
75def build_endpoint_plans(
76 settings: RunSettings,
77 owner_nonce: str,
78) -> tuple[EndpointPlan, ...]:
79 """Build every runtime's scenarios in matrix order with peak concurrency of one.
81 vLLM runs baseline, CPU-HPA and the Mooncake store; SGLang runs baseline
82 and CPU-HPA; Triton runs the baseline, served from the model registry.
83 """
84 token = owner_nonce[:16]
85 plans: list[EndpointPlan] = []
86 for runtime in settings.inference_runtimes:
87 for role in runtime.roles:
88 autoscaling = role == "hpa"
89 plans.append(
90 EndpointPlan(
91 ordinal=len(plans) + 1,
92 role=role,
93 name=f"gco-mi-{token}-{runtime.framework}-{role}",
94 replicas=(
95 settings.autoscale_initial_replicas
96 if autoscaling
97 else settings.baseline_replicas
98 ),
99 autoscaling=autoscaling,
100 runtime=runtime,
101 )
102 )
103 return tuple(plans)
106def _checkout_cli_prefix() -> list[str]:
107 """Bind every side effect to the interpreter and checkout preflight attested."""
108 return [sys.executable, "-m", "cli.main"]
111def build_deploy_command(
112 settings: RunSettings,
113 plan: EndpointPlan,
114 owner_nonce: str,
115 *,
116 model_source: str | None = None,
117) -> list[str]:
118 """Build the real argv-only ``gco inference deploy`` command.
120 ``model_source`` is the registry URI a Triton scenario serves from (the
121 plan's uploaded repository); it is required for such a plan and rejected
122 for every other one, so a deploy can never silently fall back to a
123 launcher that has no weights.
124 """
125 runtime = plan.runtime
126 if runtime.uses_model_registry != bool(model_source):
127 raise ManagedInferenceValidationError(
128 "managed inference deploy needs a model-registry source exactly when the "
129 "runtime serves from the registry"
130 )
131 command = [
132 *_checkout_cli_prefix(),
133 "--output",
134 "json",
135 "inference",
136 "deploy",
137 plan.name,
138 "--image",
139 runtime.image,
140 ]
141 # ``--framework`` is the strict adapter contract the CLI knows for the two
142 # Hugging Face launchers; Triton is recognised from its image.
143 if runtime.framework in ("vllm", "sglang"):
144 command.extend(("--framework", runtime.framework))
145 command.extend(
146 (
147 "--region",
148 settings.selected_region,
149 "--replicas",
150 str(plan.replicas),
151 "--gpu-count",
152 str(settings.gpu_count),
153 "--port",
154 str(runtime.port),
155 "--health-path",
156 settings.runtime_health_path(runtime),
157 "--namespace",
158 settings.namespace,
159 "--label",
160 f"{OWNER_LABEL}={owner_nonce}",
161 "--no-rewrite-image",
162 )
163 )
164 if model_source:
165 command.extend(("--model-source", model_source))
166 for key, value in sorted(settings.framework_env(runtime).items()):
167 command.extend(("--env", f"{key}={value}"))
168 for value in settings.deploy_extra_args(runtime, endpoint_name=plan.name):
169 if value.startswith("-"):
170 command.append(f"--extra-args={value}")
171 else:
172 command.extend(("--extra-args", value))
173 if plan.mooncake_store:
174 # TCP keeps the scenario off the EFA pool; the cold tier is what
175 # exercises the regional bucket discovery grant and its rendering.
176 command.extend(
177 (
178 "--mooncake-mode",
179 "store",
180 "--mooncake-protocol",
181 "tcp",
182 "--mooncake-cold-tier",
183 )
184 )
185 if plan.autoscaling:
186 command.extend(
187 (
188 "--autoscale-metric",
189 f"cpu:{settings.hpa_cpu_target}",
190 "--min-replicas",
191 str(settings.hpa_min_replicas),
192 "--max-replicas",
193 str(settings.hpa_max_replicas),
194 )
195 )
196 return command
199def build_invoke_command(
200 settings: RunSettings,
201 plan: EndpointPlan,
202) -> list[str]:
203 """Build the deterministic, buffered real CLI invocation."""
204 return [
205 *_checkout_cli_prefix(),
206 "--output",
207 "json",
208 "inference",
209 "invoke",
210 plan.name,
211 "--data",
212 json.dumps(
213 settings.request_body(plan.runtime),
214 sort_keys=True,
215 separators=(",", ":"),
216 ),
217 "--path",
218 plan.runtime.request_path,
219 "--region",
220 settings.selected_region,
221 ]
224def build_structured_decision_command(
225 settings: RunSettings,
226 plan: EndpointPlan,
227) -> list[str]:
228 """Build the buffered ``/v1/systemone`` call through the same authenticated path."""
229 return [
230 *_checkout_cli_prefix(),
231 "--output",
232 "json",
233 "inference",
234 "invoke",
235 plan.name,
236 "--data",
237 json.dumps(
238 settings.structured_decision_body(plan.runtime),
239 sort_keys=True,
240 separators=(",", ":"),
241 ),
242 "--path",
243 STRUCTURED_DECISION_PATH,
244 "--region",
245 settings.selected_region,
246 ]
249def build_health_command(settings: RunSettings, plan: EndpointPlan) -> list[str]:
250 """Build an authenticated health probe through the public CLI path."""
251 return [
252 *_checkout_cli_prefix(),
253 "--output",
254 "json",
255 "inference",
256 "health",
257 plan.name,
258 "--region",
259 settings.selected_region,
260 ]
263def build_models_command(settings: RunSettings, plan: EndpointPlan) -> list[str]:
264 """Build the framework-aware authenticated model-identity probe.
266 ``--framework`` names the two launcher contracts the CLI accepts; for
267 Triton the command resolves the KServe statistics path from the image.
268 """
269 command = [
270 *_checkout_cli_prefix(),
271 "--output",
272 "json",
273 "inference",
274 "models",
275 plan.name,
276 ]
277 if plan.runtime.framework in ("vllm", "sglang"):
278 command.extend(("--framework", plan.runtime.framework))
279 command.extend(("--region", settings.selected_region))
280 return command
283def build_model_upload_command(plan: EndpointPlan, repository_root: Path) -> list[str]:
284 """Build the registry upload of the plan's generated model repository."""
285 return [
286 *_checkout_cli_prefix(),
287 "--output",
288 "json",
289 "models",
290 "upload",
291 str(repository_root),
292 "--name",
293 str(plan.registry_model_name),
294 ]
297def build_model_uri_command(plan: EndpointPlan) -> list[str]:
298 """Build the lookup of the plan's registry URI (the deploy's ``--model-source``)."""
299 return [
300 *_checkout_cli_prefix(),
301 "--output",
302 "json",
303 "models",
304 "uri",
305 str(plan.registry_model_name),
306 ]
309def build_model_delete_command(plan: EndpointPlan) -> list[str]:
310 """Build the non-interactive removal of the plan's registry model."""
311 return [
312 *_checkout_cli_prefix(),
313 "models",
314 "delete",
315 str(plan.registry_model_name),
316 "--yes",
317 ]
320def build_model_list_command() -> list[str]:
321 """Build the registry listing used to prove a model's absence."""
322 return [*_checkout_cli_prefix(), "--output", "json", "models", "list"]
325def build_configure_store_command(plan: EndpointPlan) -> list[str]:
326 """Build the store reconfiguration whose re-rendered file must roll the role pods."""
327 return [
328 *_checkout_cli_prefix(),
329 "--output",
330 "json",
331 "inference",
332 "configure-store",
333 plan.name,
334 "--local-buffer-size",
335 str(MOONCAKE_STORE_RECONFIGURED_LOCAL_BUFFER_BYTES),
336 ]
339def build_delete_command(
340 settings: RunSettings,
341 plan: EndpointPlan,
342 owner_nonce: str,
343 lifecycle_id: str,
344) -> list[str]:
345 """Build deletion atomically bound to owner nonce and endpoint incarnation."""
346 del settings # Kept for a parallel command-builder interface.
347 return [
348 *_checkout_cli_prefix(),
349 "inference",
350 "delete",
351 plan.name,
352 "--expected-owner-label",
353 f"{OWNER_LABEL}={owner_nonce}",
354 "--expected-lifecycle-id",
355 lifecycle_id,
356 "--yes",
357 ]
360def _last_json_document(output: str) -> Any:
361 decoder = json.JSONDecoder()
362 candidate: Any = None
363 found = False
364 for index, character in enumerate(output):
365 if character not in "[{":
366 continue
367 try:
368 value, end = decoder.raw_decode(output, index)
369 except json.JSONDecodeError:
370 continue
371 if output[end:].strip():
372 continue
373 candidate = value
374 found = True
375 if not found:
376 raise ManagedInferenceValidationError(
377 "managed inference invoke returned no terminal JSON document"
378 )
379 return candidate
382def extract_generated_text(output: str, framework: str) -> str:
383 """Require the exact response schema for the request adapter in use."""
384 payload = _last_json_document(output)
385 text: Any = None
386 if framework == "vllm":
387 if isinstance(payload, dict):
388 choices = payload.get("choices")
389 if isinstance(choices, list) and choices and isinstance(choices[0], dict):
390 text = choices[0].get("text")
391 elif framework == "sglang":
392 # Native ``/generate`` answers a single prompt with one object; a
393 # batched prompt would answer with a list, which this adapter never
394 # sends and therefore never accepts.
395 if isinstance(payload, dict):
396 text = payload.get("text")
397 elif framework == "triton":
398 # KServe v2: the identity model must echo the request tensor exactly,
399 # so this is an equality check on ``outputs[0].data`` rather than a
400 # non-emptiness check on generated text.
401 outputs = payload.get("outputs") if isinstance(payload, dict) else None
402 first = outputs[0] if isinstance(outputs, list) and outputs else None
403 data = first.get("data") if isinstance(first, dict) else None
404 if (
405 isinstance(first, dict)
406 and first.get("name") == "OUTPUT0"
407 and isinstance(data, list)
408 and len(data) == len(TRITON_INPUT_DATA)
409 and all(
410 isinstance(value, int | float)
411 and not isinstance(value, bool)
412 and float(value) == expected
413 for value, expected in zip(data, TRITON_INPUT_DATA, strict=True)
414 )
415 ):
416 text = json.dumps([float(value) for value in data])
417 else:
418 raise ManagedInferenceValidationError("unknown managed inference response framework")
419 if not isinstance(text, str) or not text.strip():
420 raise ManagedInferenceValidationError(
421 f"managed {framework} inference response did not match its non-empty text schema"
422 )
423 return text.strip()
426def _probability(value: Any) -> float | None:
427 if isinstance(value, bool) or not isinstance(value, int | float):
428 return None
429 number = float(value)
430 return number if 0.0 <= number <= 1.0 else None
433def extract_structured_decision(output: str) -> dict[str, Any]:
434 """Require vLLM's typed ``/v1/systemone`` answer for every question asked.
436 Each answer must carry the type it was asked as, a value of that type in
437 range, label probabilities over exactly the asked labels that sum to one,
438 and a ``confidence``; every question needs a ``label_mass`` diagnostic. The
439 state is written so the answers are unambiguous, so the chosen ``topic``
440 and the ``late`` probability are checked too. Returns sanitized evidence.
441 """
442 payload = _last_json_document(output)
444 def fail(reason: str) -> ManagedInferenceValidationError:
445 return ManagedInferenceValidationError(
446 f"managed vLLM structured decision response did not match its schema: {reason}"
447 )
449 if not isinstance(payload, dict) or payload.get("object") != "structured_decision":
450 raise fail("object is not structured_decision")
451 answers = payload.get("answers")
452 diagnostics = payload.get("diagnostics")
453 if not isinstance(answers, dict) or set(answers) != set(STRUCTURED_DECISION_QUESTIONS):
454 raise fail("answers do not cover exactly the asked questions")
455 if not isinstance(diagnostics, dict) or set(diagnostics) != set(STRUCTURED_DECISION_QUESTIONS):
456 raise fail("diagnostics do not cover exactly the asked questions")
457 label_mass: dict[str, float] = {}
458 for name, question in STRUCTURED_DECISION_QUESTIONS.items():
459 answer = answers[name]
460 kind = question["type"]
461 if not isinstance(answer, dict) or answer.get("type") != kind:
462 raise fail(f"{name} is not a {kind} answer")
463 probabilities = answer.get("probabilities")
464 if kind == "choice":
465 labels = set(question["criteria"])
466 elif kind == "noul":
467 labels = {"yes", "no"}
468 else:
469 labels = {str(index) for index in range(len(question["criteria"]))}
470 if not isinstance(probabilities, dict) or set(probabilities) != labels:
471 raise fail(f"{name} probabilities do not cover exactly its labels")
472 values = [_probability(value) for value in probabilities.values()]
473 if any(value is None for value in values) or abs(sum(values) - 1.0) > 1e-3: # type: ignore[arg-type]
474 raise fail(f"{name} probabilities are not a distribution")
475 if _probability(answer.get("confidence")) is None:
476 raise fail(f"{name} confidence is not a probability")
477 if kind == "choice" and answer.get("choice") not in labels:
478 raise fail(f"{name} choice is not one of its options")
479 if kind == "noul" and _probability(answer.get("noul")) is None:
480 raise fail(f"{name} noul is not a probability")
481 if kind == "score":
482 score = answer.get("score")
483 top = len(question["criteria"]) - 1
484 if (
485 isinstance(score, bool)
486 or not isinstance(score, int | float)
487 or not 0 <= score <= top
488 ):
489 raise fail(f"{name} score is outside its scale")
490 diagnostic = diagnostics[name]
491 mass = _probability(diagnostic.get("label_mass")) if isinstance(diagnostic, dict) else None
492 if mass is None:
493 raise fail(f"{name} label_mass is not a probability")
494 label_mass[name] = mass
495 if answers["topic"].get("choice") != "shipping":
496 raise fail("topic was not answered as shipping")
497 if float(answers["late"]["noul"]) <= 0.5:
498 raise fail("late was not answered yes")
499 return {
500 "object": "structured_decision",
501 "questions": len(answers),
502 "types": sorted(question["type"] for question in STRUCTURED_DECISION_QUESTIONS.values()),
503 "topic": answers["topic"]["choice"],
504 "late_probability": round(float(answers["late"]["noul"]), 4),
505 "urgency_score": round(float(answers["urgency"]["score"]), 4),
506 "min_label_mass": round(min(label_mass.values()), 4),
507 }