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. | 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/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/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/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/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/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/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); +}); 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" 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 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`" + )