From 6be72c657549d5728011d70e9349fa7ce818b704 Mon Sep 17 00:00:00 2001 From: Yash-Chindam Date: Sun, 30 Aug 2026 15:35:56 +0530 Subject: [PATCH 1/5] feat: add a vLLM OpenAI-compatible inference backend Dispatches to a real vLLM server, serving a selected LoRA adapter by name over the shared base model, parses usage and finish reasons, streams content deltas, and reports engine health. Unreachable engines and unusable bodies raise an explicit backend error. --- pyproject.toml | 2 +- src/llm_router/backends.py | 137 ++++++++++++++++++++++++++++- tests/unit/test_backends.py | 168 ++++++++++++++++++++++++++++++++++++ 3 files changed, 304 insertions(+), 3 deletions(-) create mode 100644 tests/unit/test_backends.py diff --git a/pyproject.toml b/pyproject.toml index 587a0d0..69c7245 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,6 +10,7 @@ readme = "README.md" requires-python = ">=3.11" dependencies = [ "fastapi>=0.141.1,<1", + "httpx>=0.28.1,<1", "prometheus-client>=0.26.0,<1", "pyyaml>=6.0.2,<7", "pydantic-settings>=2.15.0,<3", @@ -18,7 +19,6 @@ dependencies = [ [project.optional-dependencies] dev = [ - "httpx>=0.28.1,<1", "mypy>=2.3.1,<3", "pytest>=9.1.1,<10", "pytest-asyncio>=1.4.0,<2", diff --git a/src/llm_router/backends.py b/src/llm_router/backends.py index a092f78..5a4a59e 100644 --- a/src/llm_router/backends.py +++ b/src/llm_router/backends.py @@ -1,9 +1,24 @@ +"""Inference backends behind the routing decision (sections 7.4 and 10). + +`MockInferenceBackend` keeps continuous integration deterministic and GPU-free. +`VLLMBackend` talks to a real vLLM OpenAI-compatible server, serving a LoRA +adapter by name when the router selected one. +""" + +import json +from collections.abc import AsyncIterator from dataclasses import dataclass -from typing import Protocol +from typing import Any, Protocol + +import httpx from llm_router.models import ChatCompletionRequest, RouteDecision +class BackendUnavailableError(RuntimeError): + """Raised when an inference engine is unreachable or returns an error.""" + + @dataclass(frozen=True) class BackendResult: text: str @@ -17,6 +32,18 @@ async def generate( self, request: ChatCompletionRequest, decision: RouteDecision ) -> BackendResult: ... + def stream( + self, request: ChatCompletionRequest, decision: RouteDecision + ) -> AsyncIterator[str]: ... + + async def healthy(self) -> bool: ... + + +def served_model_name(decision: RouteDecision) -> str: + """vLLM serves an adapter under its own name over the shared base model.""" + + return decision.adapter_id or decision.profile.id + class MockInferenceBackend: """Deterministic backend used until vLLM deployments are configured.""" @@ -24,7 +51,7 @@ class MockInferenceBackend: async def generate( self, request: ChatCompletionRequest, decision: RouteDecision ) -> BackendResult: - response = f"[{decision.profile.id}] accepted {decision.task.value} request" + response = f"[{served_model_name(decision)}] accepted {decision.task.value} request" prompt_tokens = max(1, len(request.prompt) // 4) completion_tokens = max(1, len(response) // 4) return BackendResult( @@ -32,3 +59,109 @@ async def generate( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, ) + + async def stream( + self, request: ChatCompletionRequest, decision: RouteDecision + ) -> AsyncIterator[str]: + result = await self.generate(request, decision) + for token in result.text.split(" "): + yield f"{token} " + + async def healthy(self) -> bool: + return True + + +@dataclass +class VLLMBackend: + """Client for a vLLM OpenAI-compatible server, optionally behind Ray Serve.""" + + base_url: str + client: httpx.AsyncClient + request_timeout_seconds: float = 60.0 + + def _payload( + self, request: ChatCompletionRequest, decision: RouteDecision, *, stream: bool + ) -> dict[str, Any]: + return { + "model": served_model_name(decision), + "messages": [message.model_dump() for message in request.messages], + "max_tokens": request.max_tokens, + "temperature": request.temperature, + "stream": stream, + } + + async def generate( + self, request: ChatCompletionRequest, decision: RouteDecision + ) -> BackendResult: + try: + response = await self.client.post( + f"{self.base_url}/v1/chat/completions", + json=self._payload(request, decision, stream=False), + timeout=self.request_timeout_seconds, + ) + except httpx.HTTPError as error: + raise BackendUnavailableError(f"inference engine unreachable: {error}") from error + + if response.status_code >= 400: + raise BackendUnavailableError( + f"inference engine returned {response.status_code} for " + f"{served_model_name(decision)}" + ) + + body = response.json() + try: + choice = body["choices"][0] + usage = body.get("usage", {}) + return BackendResult( + text=choice["message"]["content"], + prompt_tokens=int(usage.get("prompt_tokens", 0)), + completion_tokens=int(usage.get("completion_tokens", 0)), + finish_reason=str(choice.get("finish_reason", "stop")), + ) + except (KeyError, IndexError, TypeError) as error: + raise BackendUnavailableError("inference engine returned an unusable body") from error + + async def stream( + self, request: ChatCompletionRequest, decision: RouteDecision + ) -> AsyncIterator[str]: + payload = self._payload(request, decision, stream=True) + try: + async with self.client.stream( + "POST", + f"{self.base_url}/v1/chat/completions", + json=payload, + timeout=self.request_timeout_seconds, + ) as response: + if response.status_code >= 400: + raise BackendUnavailableError( + f"inference engine returned {response.status_code} while streaming" + ) + async for line in response.aiter_lines(): + delta = _parse_stream_line(line) + if delta: + yield delta + except httpx.HTTPError as error: + raise BackendUnavailableError(f"inference engine unreachable: {error}") from error + + async def healthy(self) -> bool: + try: + response = await self.client.get(f"{self.base_url}/health", timeout=2.0) + except httpx.HTTPError: + return False + return response.status_code < 400 + + +def _parse_stream_line(line: str) -> str: + """Extract the content delta from one server-sent-event line.""" + + if not line.startswith("data:"): + return "" + payload = line.removeprefix("data:").strip() + if not payload or payload == "[DONE]": + return "" + try: + document = json.loads(payload) + content = document["choices"][0]["delta"].get("content") + except (ValueError, KeyError, IndexError, TypeError): + return "" + return str(content) if content else "" diff --git a/tests/unit/test_backends.py b/tests/unit/test_backends.py new file mode 100644 index 0000000..b537dbf --- /dev/null +++ b/tests/unit/test_backends.py @@ -0,0 +1,168 @@ +import json + +import httpx +import pytest + +from llm_router.backends import ( + BackendUnavailableError, + MockInferenceBackend, + VLLMBackend, + served_model_name, +) +from llm_router.models import ( + ChatCompletionRequest, + ChatMessage, + ModelProfile, + RouteDecision, + TaskClass, +) + +COMPLETION_BODY = { + "choices": [ + {"message": {"role": "assistant", "content": "extracted"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 11, "completion_tokens": 3}, +} + + +def build_decision(adapter: str | None = None) -> RouteDecision: + profile = ModelProfile( + id="small-specialist", + revision="mock-small@sha256:dev", + local=True, + context_limit=8192, + supported_tasks=frozenset({TaskClass.EXTRACTION}), + quality=0.82, + ) + return RouteDecision( + profile=profile, + task=TaskClass.EXTRACTION, + reason="policy", + score=1.0, + candidate_count=1, + adapter_id=adapter, + adapter_revision=None if adapter is None else f"{adapter}@1", + ) + + +def build_request() -> ChatCompletionRequest: + return ChatCompletionRequest( + messages=[ChatMessage(role="user", content="Extract the claim fields")], max_tokens=32 + ) + + +def backend_with(handler: object) -> VLLMBackend: + transport = httpx.MockTransport(handler) # type: ignore[arg-type] + return VLLMBackend(base_url="http://engine:8000", client=httpx.AsyncClient(transport=transport)) + + +def test_served_model_name_prefers_the_selected_adapter() -> None: + assert served_model_name(build_decision()) == "small-specialist" + assert served_model_name(build_decision("claims-lora")) == "claims-lora" + + +@pytest.mark.asyncio +async def test_mock_backend_reports_the_served_name_and_streams_tokens() -> None: + backend = MockInferenceBackend() + decision = build_decision("claims-lora") + + result = await backend.generate(build_request(), decision) + chunks = [chunk async for chunk in backend.stream(build_request(), decision)] + + assert "claims-lora" in result.text + assert await backend.healthy() is True + assert "".join(chunks).strip() == result.text + + +@pytest.mark.asyncio +async def test_vllm_backend_sends_the_adapter_name_and_parses_usage() -> None: + seen: dict[str, object] = {} + + def handler(request: httpx.Request) -> httpx.Response: + seen.update(json.loads(request.content)) + return httpx.Response(200, json=COMPLETION_BODY) + + backend = backend_with(handler) + result = await backend.generate(build_request(), build_decision("claims-lora")) + + assert seen["model"] == "claims-lora" + assert seen["stream"] is False + assert seen["max_tokens"] == 32 + assert result.text == "extracted" + assert (result.prompt_tokens, result.completion_tokens) == (11, 3) + assert result.finish_reason == "stop" + + +@pytest.mark.asyncio +async def test_vllm_backend_raises_on_engine_error_status() -> None: + backend = backend_with(lambda request: httpx.Response(503, json={"error": "overloaded"})) + + with pytest.raises(BackendUnavailableError, match="returned 503"): + await backend.generate(build_request(), build_decision()) + + +@pytest.mark.asyncio +async def test_vllm_backend_raises_when_the_engine_is_unreachable() -> None: + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("refused", request=request) + + backend = backend_with(handler) + + with pytest.raises(BackendUnavailableError, match="unreachable"): + await backend.generate(build_request(), build_decision()) + + +@pytest.mark.asyncio +async def test_vllm_backend_rejects_an_unusable_response_body() -> None: + backend = backend_with(lambda request: httpx.Response(200, json={"choices": []})) + + with pytest.raises(BackendUnavailableError, match="unusable body"): + await backend.generate(build_request(), build_decision()) + + +@pytest.mark.asyncio +async def test_vllm_backend_streams_content_deltas_and_ignores_control_lines() -> None: + events = ( + 'data: {"choices":[{"delta":{"role":"assistant"}}]}\n\n' + 'data: {"choices":[{"delta":{"content":"ex"}}]}\n\n' + "\n" + 'data: {"choices":[{"delta":{"content":"tracted"}}]}\n\n' + "data: [DONE]\n\n" + ) + backend = backend_with(lambda request: httpx.Response(200, text=events)) + + chunks = [chunk async for chunk in backend.stream(build_request(), build_decision())] + + assert "".join(chunks) == "extracted" + + +@pytest.mark.asyncio +async def test_vllm_stream_reports_engine_errors() -> None: + backend = backend_with(lambda request: httpx.Response(500, text="")) + + with pytest.raises(BackendUnavailableError, match="while streaming"): + [chunk async for chunk in backend.stream(build_request(), build_decision())] + + +@pytest.mark.asyncio +async def test_vllm_stream_reports_transport_failures() -> None: + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ReadTimeout("timed out", request=request) + + backend = backend_with(handler) + + with pytest.raises(BackendUnavailableError, match="unreachable"): + [chunk async for chunk in backend.stream(build_request(), build_decision())] + + +@pytest.mark.asyncio +async def test_health_probe_reflects_engine_availability() -> None: + healthy = backend_with(lambda request: httpx.Response(200, text="ok")) + degraded = backend_with(lambda request: httpx.Response(500, text="down")) + + def unreachable(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("refused", request=request) + + assert await healthy.healthy() is True + assert await degraded.healthy() is False + assert await backend_with(unreachable).healthy() is False From ea4014023e6e64c15ccaa362782138e10fb7d161 Mon Sep 17 00:00:00 2001 From: Yash-Chindam Date: Sun, 30 Aug 2026 15:38:01 +0530 Subject: [PATCH 2/5] feat: stream chat completions and surface engine health Emits OpenAI-compatible SSE chunks, caches and replays streamed results under the same eligibility rules, fails an unreachable engine with 502 and retry guidance, and makes readiness depend on backend health. --- src/llm_router/app.py | 141 +++++++++++++++++++++++- src/llm_router/config.py | 4 + src/llm_router/models.py | 8 +- tests/integration/test_api.py | 2 +- tests/integration/test_streaming_api.py | 121 ++++++++++++++++++++ 5 files changed, 263 insertions(+), 13 deletions(-) create mode 100644 tests/integration/test_streaming_api.py diff --git a/src/llm_router/app.py b/src/llm_router/app.py index a084f79..85091ac 100644 --- a/src/llm_router/app.py +++ b/src/llm_router/app.py @@ -1,4 +1,5 @@ import hashlib +import json import secrets import time import uuid @@ -6,8 +7,9 @@ from contextlib import asynccontextmanager from pathlib import Path +import httpx from fastapi import Depends, FastAPI, Header, HTTPException, Request, Response, status -from fastapi.responses import JSONResponse +from fastapi.responses import JSONResponse, StreamingResponse from llm_router.admission import ( AdmissionController, @@ -15,7 +17,13 @@ QuotaExceededError, SlidingWindowQuota, ) -from llm_router.backends import BackendResult, InferenceBackend, MockInferenceBackend +from llm_router.backends import ( + BackendResult, + BackendUnavailableError, + InferenceBackend, + MockInferenceBackend, + VLLMBackend, +) from llm_router.caching import ( CachedCompletion, CacheStore, @@ -74,7 +82,18 @@ def create_app( runtime_settings.admission_timeout_seconds, ) quota = SlidingWindowQuota(runtime_settings.quota_requests_per_minute) - inference_backend = backend or MockInferenceBackend() + engine_client = ( + httpx.AsyncClient() if backend is None and runtime_settings.backend == "vllm" else None + ) + inference_backend: InferenceBackend = backend or ( + VLLMBackend( + base_url=runtime_settings.vllm_base_url, + client=engine_client, + request_timeout_seconds=runtime_settings.backend_timeout_seconds, + ) + if engine_client is not None + else MockInferenceBackend() + ) telemetry = metrics if metrics is not None else Metrics() exact_cache: CacheStore = cache_store or InMemoryCacheStore( max_entries=runtime_settings.cache_max_entries, @@ -95,6 +114,8 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: app.state.ready = True yield app.state.ready = False + if engine_client is not None: + await engine_client.aclose() app = FastAPI( title="Local LLM Inference Router", @@ -136,6 +157,15 @@ async def admission_handler(_: Request, error: AdmissionRejectedError) -> JSONRe content={"error": {"message": str(error), "type": "overloaded"}}, ) + @app.exception_handler(BackendUnavailableError) + async def backend_handler(_: Request, error: BackendUnavailableError) -> JSONResponse: + telemetry.record_rejection("backend_unavailable") + return JSONResponse( + status_code=502, + headers={"Retry-After": "5"}, + content={"error": {"message": str(error), "type": "backend_unavailable"}}, + ) + @app.exception_handler(QuotaExceededError) async def quota_handler(_: Request, error: QuotaExceededError) -> JSONResponse: telemetry.record_rejection("quota_exceeded") @@ -153,6 +183,8 @@ async def health() -> dict[str, str]: async def readiness(request: Request) -> dict[str, str]: if not getattr(request.app.state, "ready", False): raise HTTPException(status_code=503, detail="not ready") + if not await inference_backend.healthy(): + raise HTTPException(status_code=503, detail="inference backend is unhealthy") return {"status": "ready"} @app.get("/metrics") @@ -258,6 +290,26 @@ def _completion_response( routing=routing, ) + def _cached_stream(entry: CachedCompletion, hit_name: str) -> StreamingResponse: + completion_id = f"chatcmpl-{uuid.uuid4().hex}" + created = int(time.time()) + + async def iterator() -> AsyncIterator[str]: + yield _chunk(completion_id, created, entry.model_id, delta={"role": "assistant"}) + yield _chunk(completion_id, created, entry.model_id, delta={"content": entry.text}) + yield _chunk(completion_id, created, entry.model_id, delta={}, finish_reason="stop") + yield "data: [DONE]\n\n" + + return StreamingResponse( + iterator(), + media_type="text/event-stream", + headers={ + "X-Cache": hit_name, + "X-Route-Model": entry.model_id, + "X-Route-Revision": entry.model_revision, + }, + ) + async def _lookup_cache( payload: ChatCompletionRequest, prompt: str, cache_key: str, tenant: str ) -> tuple[str, CachedCompletion] | None: @@ -305,12 +357,75 @@ async def _store_cache( scope = semantic_cache.scope(payload, tenant=tenant, model_revision=catalog_version) semantic_cache.store(scope, prompt, entry) - @app.post("/v1/chat/completions", response_model=ChatCompletionResponse) + def _route_headers(decision: RouteDecision, cache_state: str) -> dict[str, str]: + headers = { + "X-Cache": cache_state, + "X-Route-Model": decision.profile.id, + "X-Route-Revision": decision.profile.revision, + "X-Route-Reason": decision.reason, + } + if decision.adapter_id is not None: + headers["X-Route-Adapter"] = decision.adapter_id + return headers + + def _chunk(completion_id: str, created: int, model_id: str, **choice: object) -> str: + document = { + "id": completion_id, + "object": "chat.completion.chunk", + "created": created, + "model": model_id, + "choices": [{"index": 0, **choice}], + } + return f"data: {json.dumps(document)}\n\n" + + async def _stream_completion( + payload: ChatCompletionRequest, + prompt: str, + cache_key: str, + subject: str, + decision: RouteDecision, + started: float, + queue_seconds: float, + ) -> AsyncIterator[str]: + completion_id = f"chatcmpl-{uuid.uuid4().hex}" + created = int(time.time()) + model_id = decision.profile.id + collected: list[str] = [] + + yield _chunk(completion_id, created, model_id, delta={"role": "assistant"}) + async for delta in inference_backend.stream(payload, decision): + collected.append(delta) + yield _chunk(completion_id, created, model_id, delta={"content": delta}) + yield _chunk(completion_id, created, model_id, delta={}, finish_reason="stop") + yield "data: [DONE]\n\n" + + text = "".join(collected) + prompt_tokens = max(1, len(prompt) // 4) + completion_tokens = max(1, len(text) // 4) + telemetry.record_completion( + decision, + latency_seconds=time.perf_counter() - started, + queue_seconds=queue_seconds, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + ) + await _store_cache( + payload, + prompt, + cache_key, + subject, + decision, + BackendResult( + text=text, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens + ), + ) + + @app.post("/v1/chat/completions", response_model=None) async def chat_completions( payload: ChatCompletionRequest, response: Response, subject: str = Depends(authenticate), - ) -> ChatCompletionResponse: + ) -> ChatCompletionResponse | StreamingResponse: started = time.perf_counter() await quota.consume(subject) prompt = payload.prompt @@ -321,6 +436,8 @@ async def chat_completions( cached = await _lookup_cache(payload, prompt, cache_key, subject) if cached is not None: hit_name, entry = cached + if payload.stream: + return _cached_stream(entry, hit_name) response.headers["X-Cache"] = hit_name response.headers["X-Route-Model"] = entry.model_id response.headers["X-Route-Revision"] = entry.model_revision @@ -356,6 +473,20 @@ async def chat_completions( queue_seconds = time.perf_counter() - started telemetry.inflight_requests.inc() try: + if payload.stream: + return StreamingResponse( + _stream_completion( + payload, + prompt, + cache_key, + subject, + decision, + started, + queue_seconds, + ), + media_type="text/event-stream", + headers=_route_headers(decision, "miss"), + ) result = await inference_backend.generate(payload, decision) finally: telemetry.inflight_requests.dec() diff --git a/src/llm_router/config.py b/src/llm_router/config.py index 6d49aa2..4644ff9 100644 --- a/src/llm_router/config.py +++ b/src/llm_router/config.py @@ -1,4 +1,5 @@ from functools import lru_cache +from typing import Literal from pydantic import Field, model_validator from pydantic_settings import BaseSettings, SettingsConfigDict @@ -15,6 +16,9 @@ class Settings(BaseSettings): admission_timeout_seconds: float = Field(default=0.25, gt=0) quota_requests_per_minute: int = Field(default=120, ge=1) external_fallback_enabled: bool = False + backend: Literal["mock", "vllm"] = "mock" + vllm_base_url: str = "http://127.0.0.1:8001" + backend_timeout_seconds: float = Field(default=60.0, gt=0) registry_path: str = "config/registry.yaml" routing_policy_version: str = "v1" cache_enabled: bool = True diff --git a/src/llm_router/models.py b/src/llm_router/models.py index aa981a8..c7b521e 100644 --- a/src/llm_router/models.py +++ b/src/llm_router/models.py @@ -1,7 +1,7 @@ from enum import StrEnum from typing import Any, Literal -from pydantic import BaseModel, Field, model_validator +from pydantic import BaseModel, Field class TaskClass(StrEnum): @@ -42,12 +42,6 @@ class ChatCompletionRequest(BaseModel): stream: bool = False routing: RoutingOptions = Field(default_factory=RoutingOptions) - @model_validator(mode="after") - def reject_streaming_for_initial_slice(self) -> "ChatCompletionRequest": - if self.stream: - raise ValueError("streaming is not available in the initial control-plane slice") - return self - @property def prompt(self) -> str: return "\n".join(message.content for message in self.messages) diff --git a/tests/integration/test_api.py b/tests/integration/test_api.py index 3ee1fd6..0c9538f 100644 --- a/tests/integration/test_api.py +++ b/tests/integration/test_api.py @@ -53,7 +53,7 @@ def test_validation_and_quota_errors_are_explicit() -> None: invalid = client.post( "/v1/chat/completions", headers=headers, - json={"messages": [{"role": "user", "content": "hello"}], "stream": True}, + json={"messages": [{"role": "user", "content": "hello"}], "max_tokens": 0}, ) assert invalid.status_code == 422 first = client.post( diff --git a/tests/integration/test_streaming_api.py b/tests/integration/test_streaming_api.py new file mode 100644 index 0000000..ff7dd2d --- /dev/null +++ b/tests/integration/test_streaming_api.py @@ -0,0 +1,121 @@ +import json +from collections.abc import AsyncIterator + +from fastapi.testclient import TestClient + +from llm_router.app import create_app +from llm_router.backends import BackendResult, BackendUnavailableError +from llm_router.config import Settings +from llm_router.models import ChatCompletionRequest, RouteDecision + +HEADERS = {"Authorization": "Bearer stream-key"} +BODY = { + "model": "auto", + "messages": [{"role": "user", "content": "Summarize the quarterly report"}], + "stream": True, + "routing": {"privacy": "public"}, +} + + +class FailingBackend: + async def generate( + self, request: ChatCompletionRequest, decision: RouteDecision + ) -> BackendResult: + raise BackendUnavailableError("inference engine unreachable: connection refused") + + async def stream( + self, request: ChatCompletionRequest, decision: RouteDecision + ) -> AsyncIterator[str]: + raise BackendUnavailableError("inference engine unreachable while streaming") + yield "" # pragma: no cover - unreachable, keeps this an async generator + + async def healthy(self) -> bool: + return False + + +def build_client(**overrides: object) -> TestClient: + settings = Settings(api_keys="stream-key", **overrides) # type: ignore[arg-type] + return TestClient(create_app(settings)) + + +def parse_events(text: str) -> list[dict[str, object]]: + return [ + json.loads(line.removeprefix("data:").strip()) + for line in text.splitlines() + if line.startswith("data:") and line.removeprefix("data:").strip() != "[DONE]" + ] + + +def test_streaming_response_emits_openai_compatible_chunks() -> None: + with build_client() as client: + response = client.post("/v1/chat/completions", headers=HEADERS, json=BODY) + + assert response.status_code == 200 + assert response.headers["content-type"].startswith("text/event-stream") + assert response.headers["x-route-model"] + assert response.text.rstrip().endswith("data: [DONE]") + + events = parse_events(response.text) + assert events[0]["choices"][0]["delta"] == {"role": "assistant"} + assert events[-1]["choices"][0]["finish_reason"] == "stop" + assert all(event["object"] == "chat.completion.chunk" for event in events) + text = "".join( + str(event["choices"][0]["delta"].get("content", "")) + for event in events # type: ignore[union-attr] + ) + assert "accepted" in text + + +def test_streamed_result_is_cached_and_replayed() -> None: + with build_client() as client: + client.post("/v1/chat/completions", headers=HEADERS, json=BODY) + replay = client.post("/v1/chat/completions", headers=HEADERS, json=BODY) + + assert replay.headers["x-cache"] == "exact" + assert replay.headers["content-type"].startswith("text/event-stream") + events = parse_events(replay.text) + assert events[1]["choices"][0]["delta"]["content"] # type: ignore[index] + + +def test_streaming_records_completion_metrics() -> None: + with build_client() as client: + client.post("/v1/chat/completions", headers=HEADERS, json=BODY) + metrics = client.get("/metrics").text + + assert "router_requests_total" in metrics + assert "router_tokens_total" in metrics + + +def test_unreachable_backend_returns_bad_gateway_with_retry_guidance() -> None: + settings = Settings(api_keys="stream-key") + with TestClient(create_app(settings, backend=FailingBackend())) as client: + response = client.post( + "/v1/chat/completions", + headers=HEADERS, + json={"model": "auto", "messages": [{"role": "user", "content": "hello"}]}, + ) + metrics = client.get("/metrics").text + + assert response.status_code == 502 + assert response.headers["retry-after"] == "5" + assert response.json()["error"]["type"] == "backend_unavailable" + assert 'router_rejections_total{type="backend_unavailable"} 1.0' in metrics + + +def test_readiness_fails_while_the_backend_is_unhealthy() -> None: + settings = Settings(api_keys="stream-key") + with TestClient(create_app(settings, backend=FailingBackend())) as client: + assert client.get("/healthz").status_code == 200 + assert client.get("/readyz").status_code == 503 + + +def test_vllm_backend_is_selected_by_configuration() -> None: + with build_client(backend="vllm", vllm_base_url="http://127.0.0.1:9") as client: + response = client.post( + "/v1/chat/completions", + headers=HEADERS, + json={"model": "auto", "messages": [{"role": "user", "content": "hello"}]}, + ) + + assert response.status_code == 502 + assert response.json()["error"]["type"] == "backend_unavailable" From 5c7ac9c07c75719c6d4fee6687fc64199f76244e Mon Sep 17 00:00:00 2001 From: Yash-Chindam Date: Sun, 30 Aug 2026 15:39:16 +0530 Subject: [PATCH 3/5] feat: generate Ray Serve deployment configuration from the catalog Renders per-tier autoscaling, GPU pool placement, tensor parallelism, prefix caching, quantization, and Multi-LoRA settings from model cards, plus a canary description naming its rollback target. The committed config/ray-serve.yaml is verified against the catalog by a test. --- config/ray-serve.yaml | 99 +++++++++++++++++++++++ src/llm_router/serving.py | 155 +++++++++++++++++++++++++++++++++++++ tests/unit/test_serving.py | 116 +++++++++++++++++++++++++++ 3 files changed, 370 insertions(+) create mode 100644 config/ray-serve.yaml create mode 100644 src/llm_router/serving.py create mode 100644 tests/unit/test_serving.py diff --git a/config/ray-serve.yaml b/config/ray-serve.yaml new file mode 100644 index 0000000..d049a09 --- /dev/null +++ b/config/ray-serve.yaml @@ -0,0 +1,99 @@ +policy_version: v1 +applications: +- model_id: small-specialist + model_revision: mock-small@sha256:dev + stage: production + model_loading_config: + model_id: small-specialist + model_source: registry://small-specialist@mock-small@sha256:dev + tokenizer: mock-small-tokenizer + accelerator_type: nvidia-l4 + deployment_config: + autoscaling_config: + min_replicas: 1 + max_replicas: 4 + target_ongoing_requests: 8 + upscale_delay_s: 10 + downscale_delay_s: 300 + max_ongoing_requests: 16 + ray_actor_options: + num_gpus: 1 + resources: + gpu_pool_nvidia-l4: 0.001 + engine_kwargs: + max_model_len: 8192 + tensor_parallel_size: 1 + enable_prefix_caching: true + enable_chunked_prefill: true + quantization: awq + enable_lora: true + max_loras: 3 + max_lora_rank: 32 + lora_config: + dynamic_lora_loading_path: registry://adapters/small-specialist + adapters: + - id: claims-extraction-lora + revision: claims-lora@sha256:dev + domain: claims + quantized: true + stage: production + - id: claims-extraction-lora-next + revision: claims-lora@sha256:next + domain: claims + quantized: false + stage: staging + - id: support-classification-lora + revision: support-lora@sha256:dev + domain: support + quantized: false + stage: production +- model_id: general-local + model_revision: mock-general@sha256:dev + stage: production + model_loading_config: + model_id: general-local + model_source: registry://general-local@mock-general@sha256:dev + tokenizer: mock-general-tokenizer + accelerator_type: nvidia-a10g + deployment_config: + autoscaling_config: + min_replicas: 1 + max_replicas: 4 + target_ongoing_requests: 8 + upscale_delay_s: 10 + downscale_delay_s: 300 + max_ongoing_requests: 16 + ray_actor_options: + num_gpus: 1 + resources: + gpu_pool_nvidia-a10g: 0.001 + engine_kwargs: + max_model_len: 32768 + tensor_parallel_size: 1 + enable_prefix_caching: true + enable_chunked_prefill: true +- model_id: high-capability + model_revision: mock-high@sha256:dev + stage: production + model_loading_config: + model_id: high-capability + model_source: registry://high-capability@mock-high@sha256:dev + tokenizer: mock-high-tokenizer + accelerator_type: nvidia-a100 + deployment_config: + autoscaling_config: + min_replicas: 0 + max_replicas: 2 + target_ongoing_requests: 8 + upscale_delay_s: 10 + downscale_delay_s: 60 + max_ongoing_requests: 16 + ray_actor_options: + num_gpus: 2 + resources: + gpu_pool_nvidia-a100: 0.001 + engine_kwargs: + max_model_len: 65536 + tensor_parallel_size: 2 + enable_prefix_caching: true + enable_chunked_prefill: true diff --git a/src/llm_router/serving.py b/src/llm_router/serving.py new file mode 100644 index 0000000..d14f351 --- /dev/null +++ b/src/llm_router/serving.py @@ -0,0 +1,155 @@ +"""Ray Serve deployment configuration derived from the registry (section 7.3). + +The control plane does not invent deployment topology at runtime. This module +renders a declarative Ray Serve LLM configuration from the governed catalog so +what is deployed is traceable to a model card, an adapter, and a GPU pool. +""" + +from typing import Any + +import yaml + +from llm_router.registry import LifecycleStage, ModelCard, Quantization, Registry + +DEFAULT_MIN_REPLICAS = 1 +DEFAULT_MAX_REPLICAS = 4 +DEFAULT_MAX_ONGOING_REQUESTS = 16 +DEFAULT_TARGET_ONGOING_REQUESTS = 8 + + +class ServingConfigError(RuntimeError): + """Raised when the catalog cannot produce a deployable configuration.""" + + +def _autoscaling_config(card: ModelCard) -> dict[str, Any]: + """Latency-sensitive tiers keep a warm replica; heavier tiers scale wider.""" + + warm = card.tier.value in {"small-specialist", "general-local"} + return { + "min_replicas": DEFAULT_MIN_REPLICAS if warm else 0, + "max_replicas": DEFAULT_MAX_REPLICAS if warm else 2, + "target_ongoing_requests": DEFAULT_TARGET_ONGOING_REQUESTS, + "upscale_delay_s": 10, + "downscale_delay_s": 300 if warm else 60, + } + + +def _engine_kwargs(card: ModelCard, adapter_count: int) -> dict[str, Any]: + engine: dict[str, Any] = { + "max_model_len": card.context_limit, + "tensor_parallel_size": card.hardware.tensor_parallel_size, + "enable_prefix_caching": True, + "enable_chunked_prefill": True, + } + if card.quantization is not Quantization.NONE: + engine["quantization"] = card.quantization.value + if adapter_count: + engine["enable_lora"] = True + engine["max_loras"] = adapter_count + engine["max_lora_rank"] = 32 + return engine + + +def build_serving_config(registry: Registry) -> dict[str, Any]: + """Render one Ray Serve LLM application entry per servable local model.""" + + applications: list[dict[str, Any]] = [] + for card in registry.servable_models(): + if not card.local: + continue + adapters = [ + adapter + for adapter in registry.servable_adapters() + if adapter.base_model_id == card.id and adapter.base_revision == card.revision + ] + entry: dict[str, Any] = { + "model_id": card.id, + "model_revision": card.revision, + "stage": card.stage.value, + "model_loading_config": { + "model_id": card.id, + "model_source": f"registry://{card.id}@{card.revision}", + "tokenizer": card.tokenizer, + }, + "accelerator_type": card.hardware.accelerator, + "deployment_config": { + "autoscaling_config": _autoscaling_config(card), + "max_ongoing_requests": DEFAULT_MAX_ONGOING_REQUESTS, + "ray_actor_options": { + "num_gpus": card.hardware.count, + "resources": {f"gpu_pool_{card.hardware.accelerator}": 0.001}, + }, + }, + "engine_kwargs": _engine_kwargs(card, len(adapters)), + } + if adapters: + entry["lora_config"] = { + "dynamic_lora_loading_path": f"registry://adapters/{card.id}", + "adapters": [ + { + "id": adapter.id, + "revision": adapter.adapter_revision, + "domain": adapter.domain, + "quantized": adapter.quantized, + "stage": adapter.stage.value, + } + for adapter in sorted(adapters, key=lambda item: item.id) + ], + } + applications.append(entry) + + if not applications: + raise ServingConfigError("catalog contains no servable local models") + + return { + "policy_version": registry.policy.version, + "applications": applications, + } + + +def canary_config(registry: Registry, deployment_id: str) -> dict[str, Any]: + """Describe a canary and the revision it rolls back to (section 13).""" + + current = next((item for item in registry.deployments if item.id == deployment_id), None) + if current is None: + raise ServingConfigError(f"unknown deployment {deployment_id}") + target = registry.rollback_target(deployment_id) + staged_adapters = [ + adapter.id + for adapter in registry.servable_adapters() + if adapter.stage is LifecycleStage.STAGING + ] + return { + "deployment_id": current.id, + "container_digest": current.container_digest, + "gpu_pool": current.gpu_pool, + "canary_traffic_percent": 10, + "promote_after_successful_requests": 500, + "rollback_to": None if target is None else target.id, + "rollback_triggers": [ + "readiness probe failure", + "p95 latency above the tier objective", + "quality below the benchmark floor", + "error rate above one percent", + ], + "staged_adapters": sorted(staged_adapters), + } + + +def render_serving_config(registry: Registry) -> str: + return yaml.safe_dump(build_serving_config(registry), sort_keys=False) + + +def main() -> None: # pragma: no cover - thin command-line wrapper + """Render the deployment configuration for the committed catalog.""" + + import sys + + from llm_router.registry import load_registry + + catalog = sys.argv[1] if len(sys.argv) > 1 else "config/registry.yaml" + sys.stdout.write(render_serving_config(load_registry(catalog))) + + +if __name__ == "__main__": # pragma: no cover - command-line entry point + main() diff --git a/tests/unit/test_serving.py b/tests/unit/test_serving.py new file mode 100644 index 0000000..faa6533 --- /dev/null +++ b/tests/unit/test_serving.py @@ -0,0 +1,116 @@ +from pathlib import Path + +import pytest +import yaml + +from llm_router.registry import Registry, load_registry +from llm_router.serving import ( + ServingConfigError, + build_serving_config, + canary_config, + render_serving_config, +) + +CATALOG = "config/registry.yaml" + + +@pytest.fixture(scope="module") +def registry() -> Registry: + return load_registry(CATALOG) + + +def application(config: dict[str, object], model_id: str) -> dict[str, object]: + applications = config["applications"] + assert isinstance(applications, list) + return next(entry for entry in applications if entry["model_id"] == model_id) + + +def test_serving_config_covers_local_models_only(registry: Registry) -> None: + config = build_serving_config(registry) + + identifiers = {entry["model_id"] for entry in config["applications"]} + assert identifiers == {"small-specialist", "general-local", "high-capability"} + assert config["policy_version"] == "v1" + + +def test_engine_settings_follow_the_model_card(registry: Registry) -> None: + config = build_serving_config(registry) + + small = application(config, "small-specialist") + high = application(config, "high-capability") + + assert small["engine_kwargs"]["quantization"] == "awq" + assert small["engine_kwargs"]["max_model_len"] == 8192 + assert small["engine_kwargs"]["enable_prefix_caching"] is True + assert high["engine_kwargs"]["tensor_parallel_size"] == 2 + assert "quantization" not in high["engine_kwargs"] + assert high["deployment_config"]["ray_actor_options"]["num_gpus"] == 2 + + +def test_latency_sensitive_tiers_keep_a_warm_replica(registry: Registry) -> None: + config = build_serving_config(registry) + + small = application(config, "small-specialist")["deployment_config"]["autoscaling_config"] + high = application(config, "high-capability")["deployment_config"]["autoscaling_config"] + + assert small["min_replicas"] == 1 + assert small["max_replicas"] == 4 + assert high["min_replicas"] == 0 + assert high["downscale_delay_s"] < small["downscale_delay_s"] + + +def test_multi_lora_serving_is_enabled_only_where_adapters_exist(registry: Registry) -> None: + config = build_serving_config(registry) + + small = application(config, "small-specialist") + general = application(config, "general-local") + + assert small["engine_kwargs"]["enable_lora"] is True + assert small["engine_kwargs"]["max_loras"] == len(small["lora_config"]["adapters"]) + assert [adapter["id"] for adapter in small["lora_config"]["adapters"]] == [ + "claims-extraction-lora", + "claims-extraction-lora-next", + "support-classification-lora", + ] + assert "lora_config" not in general + assert "enable_lora" not in general["engine_kwargs"] + + +def test_rendered_configuration_is_valid_yaml(registry: Registry) -> None: + document = yaml.safe_load(render_serving_config(registry)) + + assert document["applications"][0]["model_loading_config"]["model_source"].startswith( + "registry://" + ) + + +def test_empty_catalog_is_rejected() -> None: + with pytest.raises(ServingConfigError, match="no servable local models"): + build_serving_config(Registry(models=())) + + +def test_canary_config_names_its_rollback_target(registry: Registry) -> None: + canary = canary_config(registry, "deploy-0002") + + assert canary["rollback_to"] == "deploy-0001" + assert canary["canary_traffic_percent"] == 10 + assert canary["staged_adapters"] == ["claims-extraction-lora-next"] + assert "readiness probe failure" in canary["rollback_triggers"] + + +def test_canary_config_reports_an_unknown_deployment(registry: Registry) -> None: + with pytest.raises(ServingConfigError, match="unknown deployment"): + canary_config(registry, "deploy-9999") + + +def test_first_deployment_has_no_rollback_target(registry: Registry) -> None: + assert canary_config(registry, "deploy-0001")["rollback_to"] is None + + +def test_committed_deployment_configuration_matches_the_catalog(registry: Registry) -> None: + committed = Path("config/ray-serve.yaml").read_text(encoding="utf-8") + + assert committed == render_serving_config(registry), ( + "config/ray-serve.yaml is stale; regenerate with " + "`python -m llm_router.serving > config/ray-serve.yaml`" + ) From d235a6ba5c3b5b665fe0c591b42b16ee82cacb8b Mon Sep 17 00:00:00 2001 From: Yash-Chindam Date: Sun, 30 Aug 2026 15:39:47 +0530 Subject: [PATCH 4/5] test: verify streamed chunk sequence end to end --- tests/e2e/api.spec.ts | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/tests/e2e/api.spec.ts b/tests/e2e/api.spec.ts index 75268d1..ffa8588 100644 --- a/tests/e2e/api.spec.ts +++ b/tests/e2e/api.spec.ts @@ -126,3 +126,33 @@ test("publishes model cards and deployment rollback targets", async ({ request } ); expect(current.rollback_target).toBe("deploy-0001"); }); + +test("streams an OpenAI-compatible chunk sequence", async ({ request }) => { + const response = await request.post("/v1/chat/completions", { + data: { + model: "auto", + messages: [{ role: "user", content: "Summarize this streaming probe" }], + stream: true, + routing: { privacy: "public" }, + }, + }); + + expect(response.status()).toBe(200); + expect(response.headers()["content-type"]).toContain("text/event-stream"); + + const body = await response.text(); + expect(body.trimEnd().endsWith("data: [DONE]")).toBeTruthy(); + + const events = body + .split("\n") + .filter((line) => line.startsWith("data:")) + .map((line) => line.replace(/^data:\s*/, "")) + .filter((payload) => payload !== "[DONE]") + .map((payload) => JSON.parse(payload)); + + expect(events[0].object).toBe("chat.completion.chunk"); + expect(events[0].choices[0].delta.role).toBe("assistant"); + expect(events.at(-1).choices[0].finish_reason).toBe("stop"); + const text = events.map((event) => event.choices[0].delta.content ?? "").join(""); + expect(text.length).toBeGreaterThan(0); +}); From 6bea4a7c5d5bdf6808d66de89ddf812fe99ba0b9 Mon Sep 17 00:00:00 2001 From: Yash-Chindam Date: Sun, 30 Aug 2026 15:39:48 +0530 Subject: [PATCH 5/5] docs: document serving backends, streaming, and Ray Serve config --- README.md | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/README.md b/README.md index bb6b59c..4956d51 100644 --- a/README.md +++ b/README.md @@ -61,6 +61,28 @@ Successful PR CI runs are merged automatically only for trusted same-repository and Dependabot. Forks, drafts, and untrusted author associations are deliberately skipped; repository branch-protection and review requirements continue to apply. +## Serving backends + +`ROUTER_BACKEND=mock` (the default) keeps CI deterministic and GPU-free. +`ROUTER_BACKEND=vllm` dispatches to a vLLM OpenAI-compatible server; a selected LoRA +adapter is served by name over the shared base model. An unreachable or failing engine +returns `502` with retry guidance, and `/readyz` fails while the engine is unhealthy. + +Set `"stream": true` to receive OpenAI-compatible `text/event-stream` chunks. Streamed +results are cached under the same eligibility rules and replayed as chunks on a hit. + +### Ray Serve deployment + +[`config/ray-serve.yaml`](config/ray-serve.yaml) is generated from the catalog, never +hand-edited, and verified by a test: + +```bash +python -m llm_router.serving > config/ray-serve.yaml +``` + +It carries per-tier autoscaling (latency-sensitive tiers keep a warm replica), GPU pool +placement, tensor parallelism, prefix caching, quantization, and Multi-LoRA settings. + ## Model registry [`config/registry.yaml`](config/registry.yaml) is the governed source of truth for what may @@ -127,6 +149,9 @@ All settings use the `ROUTER_` prefix. | `ROUTER_ADMISSION_TIMEOUT_SECONDS` | `0.25` | Time allowed to wait for capacity. | | `ROUTER_QUOTA_REQUESTS_PER_MINUTE` | `120` | Per-token sliding-window quota. | | `ROUTER_EXTERNAL_FALLBACK_ENABLED` | `false` | Operator gate for external fallback. | +| `ROUTER_BACKEND` | `mock` | `mock` or `vllm`. | +| `ROUTER_VLLM_BASE_URL` | `http://127.0.0.1:8001` | vLLM OpenAI-compatible endpoint. | +| `ROUTER_BACKEND_TIMEOUT_SECONDS` | `60` | Per-request engine timeout. | | `ROUTER_REGISTRY_PATH` | `config/registry.yaml` | Governed model catalog; built-in profiles are used if absent. | | `ROUTER_ROUTING_POLICY_VERSION` | `v1` | Invalidates router and response caches when changed. | | `ROUTER_CACHE_ENABLED` | `true` | Master switch for all cache tiers. |