"""Private client for the vLLM Semantic Router built-in model runtime. `vllm-sr-runtime serve` (``vllm-sr serve ``) answers one state per request on ``/v1/systemone`` and reports its artifact on ``/v1/models`` and in each response's ``meta``, instead of the ``X-Decision-Artifact-*`` headers of ``vllm-sr decision serve``. This client keeps the same guarantees as ``DirectGateway``: every answer is bound to the pinned repository, revision, manifest and model identity on the ``exact`` profile; the runtime's answers are projected onto the strict SystemOne envelope without changing a value; and a shared-question batch becomes one runtime request per state, reassembled in request order. """ from __future__ import annotations import asyncio import json import math from collections.abc import Mapping from types import MappingProxyType import httpx from direct_contract import ( CONFIDENCE_DEFINITIONS, MARGIN_ORDINAL_V1, validate_response_for_request, ) from direct_gateway import ( MAX_BATCH_REQUEST_BYTES, MAX_INFERENCE_RESPONSE_BYTES, MAX_SINGLE_REQUEST_BYTES, MAX_TIMEOUT_SECONDS, MIN_TIMEOUT_SECONDS, DirectGatewayError, _hex_digest, _loads_strict, _private_origin, _safe_retry_after, ) SINGLE_PATH = "/v1/systemone" HEALTH_PATH = "/health" MODELS_PATH = "/v1/models" MAX_CONTROL_RESPONSE_BYTES = 256 * 1024 MAX_CONTROL_TIMEOUT_SECONDS = 2.0 EXACT = "exact" # Golden statuses of a runtime that verified its package and answers # deterministically; a "failed" golden check never reaches readiness. SERVING_GOLDEN = frozenset({"matched", "unverified"}) INPUT_ERRORS = frozenset({"invalid_question", "max_length_exceeded"}) RESPONSE_FIELDS = frozenset({"model", "answers", "usage", "meta"}) class ModelRuntimeGateway: """Attested client for one or more built-in runtime instances.""" def __init__( self, endpoints: Mapping[str, str], *, expected_artifacts: Mapping[str, Mapping[str, str]], timeout_seconds: float = 45.0, client: httpx.AsyncClient | None = None, max_concurrency: int = 4, ): expected = tuple(expected_artifacts) if not expected or len(set(expected)) != len(expected): raise ValueError("Configure a nonempty set of canonical Decision models") if set(endpoints) != set(expected): raise ValueError("runtime endpoints must cover exactly the configured models") pinned = {} confidence = {} for model in expected: item = expected_artifacts[model] if not isinstance(item, Mapping): raise ValueError("Pin each Decision model to a released artifact") revision = item.get("revision") manifest = item.get("manifest_sha256") content = item.get("content_sha256") if ( not _hex_digest(revision, 40) or not _hex_digest(manifest, 64) or (content is not None and not _hex_digest(content, 64)) ): raise ValueError("Pin each Decision model revision and manifest digest") pinned[model] = MappingProxyType( { "revision": revision, "manifest_sha256": manifest, **({"content_sha256": content} if content is not None else {}), } ) confidence[model] = item.get("confidence_definition", MARGIN_ORDINAL_V1) if confidence[model] not in CONFIDENCE_DEFINITIONS: raise ValueError("Unknown Decision confidence definition") if ( isinstance(timeout_seconds, bool) or not isinstance(timeout_seconds, (int, float)) or not math.isfinite(timeout_seconds) or not MIN_TIMEOUT_SECONDS <= timeout_seconds <= MAX_TIMEOUT_SECONDS ): raise ValueError("runtime timeout must be between 0.1 and 300 seconds") if type(max_concurrency) is not int or not 1 <= max_concurrency <= 64: raise ValueError("Invalid runtime request concurrency") normalized = {model: _private_origin(endpoints[model]) for model in expected} if len(set(normalized.values())) != len(normalized): raise ValueError("Each configured Decision model requires its own runtime origin") self._endpoints = MappingProxyType(normalized) self.expected_artifacts = MappingProxyType(pinned) self.confidence_definitions = MappingProxyType(confidence) self.models = frozenset(expected) self.timeout_seconds = float(timeout_seconds) self._max_concurrency = max_concurrency self._owns_client = client is None self._client = client or httpx.AsyncClient( follow_redirects=False, timeout=self.timeout_seconds, trust_env=False, ) async def aclose(self) -> None: if self._owns_client: await self._client.aclose() # -- readiness ----------------------------------------------------------- async def ready(self, model: str) -> bool: return (await self.probe(model))["loaded"] async def probe(self, model: str) -> dict[str, object]: """Readiness plus the live artifact, with no private topology.""" if model not in self.models: return {"loaded": False, "artifact": None} timeout = min(self.timeout_seconds, MAX_CONTROL_TIMEOUT_SECONDS) try: health = await self._request_json("GET", model, HEALTH_PATH, timeout=timeout) listing = await self._request_json("GET", model, MODELS_PATH, timeout=timeout) except DirectGatewayError: return {"loaded": False, "artifact": None} artifact = self._card_artifact(model, listing) loaded = ( isinstance(health, dict) and health.get("status") == "ready" and health.get("model") == model and artifact is not None and self._artifact_matches(model, artifact) ) return {"loaded": bool(loaded), "artifact": artifact} @staticmethod def _card_artifact(model: str, listing: object) -> dict[str, str] | None: if not isinstance(listing, dict) or not isinstance(listing.get("data"), list): return None if len(listing["data"]) != 1 or not isinstance(listing["data"][0], dict): return None card = listing["data"][0] golden = card.get("golden") if ( card.get("id") != model or card.get("repo") != model or card.get("ready") is not True or card.get("profile") != EXACT or not isinstance(golden, dict) or golden.get("status") not in SERVING_GOLDEN or not _hex_digest(card.get("revision"), 40) or not _hex_digest(card.get("manifest_sha256"), 64) or not _hex_digest(card.get("model_sha256"), 64) ): return None return { "model": card["repo"], "revision": card["revision"], "manifest_sha256": card["manifest_sha256"], "content_sha256": card["model_sha256"], } def _artifact_matches(self, model: str, artifact: Mapping[str, str]) -> bool: pinned = self.expected_artifacts[model] return ( artifact["model"] == model and artifact["revision"] == pinned["revision"] and artifact["manifest_sha256"] == pinned["manifest_sha256"] and ( "content_sha256" not in pinned or artifact["content_sha256"] == pinned["content_sha256"] ) ) # -- inference ----------------------------------------------------------- async def evaluate(self, payload: Mapping[str, object], *, batch: bool) -> dict: model = payload.get("model") if isinstance(payload, Mapping) else None if not isinstance(model, str) or model not in self.models: raise DirectGatewayError(422, "The selected canonical model is not configured") if batch != ("states" in payload) or ("state" in payload and "states" in payload): raise DirectGatewayError(422, "The Decision request does not match its route") try: encoded = json.dumps( payload, ensure_ascii=False, separators=(",", ":"), allow_nan=False ).encode("utf-8") except (TypeError, ValueError, UnicodeError, RecursionError) as exc: raise DirectGatewayError(422, "The Decision request is not valid JSON") from exc if len(encoded) > (MAX_BATCH_REQUEST_BYTES if batch else MAX_SINGLE_REQUEST_BYTES): raise DirectGatewayError(413, "Request body exceeds the endpoint limit") try: async with asyncio.timeout(self.timeout_seconds): if batch: response = await self._batch(model, payload) else: response = await self._single(model, payload["state"], payload["questions"]) except TimeoutError as exc: raise DirectGatewayError( 504, "The selected Decision runtime did not respond in time" ) from exc try: validate_response_for_request( payload, response, batch=batch, confidence=self.confidence_definitions[model], ) except (TypeError, ValueError, KeyError, OverflowError, RecursionError) as exc: raise DirectGatewayError( 502, "The selected Decision runtime returned an invalid response" ) from exc return response async def _batch(self, model: str, payload: Mapping[str, object]) -> dict: states = payload["states"] if not isinstance(states, list) or not states: raise DirectGatewayError(422, "The Decision request does not match its route") gate = asyncio.Semaphore(self._max_concurrency) async def one(item): if not isinstance(item, dict) or set(item) != {"id", "state"}: raise DirectGatewayError(422, "Each context must contain exactly id and state") async with gate: return await self._single(model, item["state"], payload["questions"]) tasks = [asyncio.create_task(one(item)) for item in states] try: rows = await asyncio.gather(*tasks) except BaseException: for task in tasks: task.cancel() await asyncio.gather(*tasks, return_exceptions=True) raise results = [ {"id": item["id"], "answers": row["answers"], "usage": row["usage"]} for item, row in zip(states, rows, strict=True) ] return { "model": model, "results": results, "usage": { "input_tokens": sum(row["usage"]["input_tokens"] for row in rows), "output_tokens": sum(row["usage"]["output_tokens"] for row in rows), }, } async def _single(self, model: str, state: object, questions: object) -> dict: body = { "model": model, "state": state, "questions": questions, "options": {"profile": EXACT, "return_meta": True}, } response = await self._request_json( "POST", model, SINGLE_PATH, content=json.dumps( body, ensure_ascii=False, separators=(",", ":"), allow_nan=False ).encode("utf-8"), limit=MAX_INFERENCE_RESPONSE_BYTES, ) if ( not isinstance(response, dict) or set(response) != RESPONSE_FIELDS or response["model"] != model or not isinstance(response["answers"], dict) or not isinstance(response["usage"], dict) ): raise DirectGatewayError( 502, "The selected Decision runtime returned an invalid response" ) if not self._meta_matches(model, response["meta"]): raise DirectGatewayError( 503, "The selected Decision runtime artifact is unavailable" ) errors = { answer.get("error") for answer in response["answers"].values() if isinstance(answer, dict) and "error" in answer } if errors & INPUT_ERRORS: raise DirectGatewayError(413, "Decision input exceeds the selected model limit") if errors: raise DirectGatewayError(503, "The selected Decision runtime is unavailable") return { "model": response["model"], "answers": response["answers"], "usage": response["usage"], } def _meta_matches(self, model: str, meta: object) -> bool: pinned = self.expected_artifacts[model] return ( isinstance(meta, dict) and meta.get("revision") == pinned["revision"] and meta.get("profile") == EXACT and meta.get("numerics") == EXACT and _hex_digest(meta.get("model_sha256"), 64) and ( "content_sha256" not in pinned or meta["model_sha256"] == pinned["content_sha256"] ) ) async def _request_json( self, method: str, model: str, path: str, *, content: bytes | None = None, limit: int = MAX_CONTROL_RESPONSE_BYTES, timeout: float | None = None, ): endpoint = self._endpoints[model] request_timeout = timeout or self.timeout_seconds headers = {"Accept": "application/json", "Accept-Encoding": "identity"} if content is not None: headers["Content-Type"] = "application/json" try: async with asyncio.timeout(request_timeout): async with self._client.stream( method, endpoint + path, content=content, headers=headers, follow_redirects=False, timeout=request_timeout, ) as response: if 300 <= response.status_code < 400: raise DirectGatewayError( 502, "The selected Decision runtime returned an unsupported redirect" ) if response.status_code >= 400 and not ( path == HEALTH_PATH and response.status_code == 503 ): self._raise_status(response) return await self._read_json(response, limit) except asyncio.CancelledError: raise except DirectGatewayError: raise except (TimeoutError, httpx.TimeoutException) as exc: raise DirectGatewayError( 504, "The selected Decision runtime did not respond in time" ) from exc except httpx.HTTPError as exc: raise DirectGatewayError(503, "The selected Decision runtime is unavailable") from exc @staticmethod def _raise_status(response: httpx.Response) -> None: status = response.status_code if status in {400, 413}: raise DirectGatewayError(413, "Decision input exceeds the selected model limit") if status in {429, 529}: raise DirectGatewayError( status, "The selected Decision runtime is temporarily overloaded", retry_after=_safe_retry_after(response.headers.get("retry-after")), ) if status == 503: raise DirectGatewayError( 503, "The selected Decision runtime is unavailable", retry_after=_safe_retry_after(response.headers.get("retry-after")), ) if status == 504: raise DirectGatewayError(504, "The selected Decision runtime did not respond in time") raise DirectGatewayError(502, "The selected Decision runtime rejected a validated request") @staticmethod async def _read_json(response: httpx.Response, limit: int): content_type = response.headers.get("content-type", "").split(";", 1)[0].strip().lower() if content_type != "application/json": raise DirectGatewayError( 502, "The selected Decision runtime returned an invalid content type" ) encoding = response.headers.get("content-encoding", "identity").strip().lower() if encoding not in {"", "identity"}: raise DirectGatewayError( 502, "The selected Decision runtime returned an unsupported encoding" ) body = bytearray() async for chunk in response.aiter_bytes(chunk_size=64 * 1024): if len(chunk) > limit - len(body): raise DirectGatewayError( 502, "The selected Decision runtime returned an oversized response" ) body.extend(chunk) try: return _loads_strict(bytes(body)) except (TypeError, ValueError, UnicodeError, RecursionError) as exc: raise DirectGatewayError( 502, "The selected Decision runtime returned invalid JSON" ) from exc