diff --git a/Dockerfile b/Dockerfile index 35f1a0c..6bd0c24 100644 --- a/Dockerfile +++ b/Dockerfile @@ -15,8 +15,11 @@ RUN useradd --create-home --uid 10001 appuser COPY --from=builder /wheels /wheels RUN python -m pip install --no-cache-dir /wheels/* && rm -rf /wheels -USER appuser WORKDIR /app +COPY config ./config +RUN chown -R appuser:appuser /app + +USER appuser EXPOSE 8000 HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \ CMD python -c "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8000/healthz', timeout=2)" diff --git a/README.md b/README.md index e9751b2..bb6b59c 100644 --- a/README.md +++ b/README.md @@ -61,6 +61,29 @@ 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. +## Model registry + +[`config/registry.yaml`](config/registry.yaml) is the governed source of truth for what may +be served. A request can never introduce a model path, revision, or adapter. + +- Model cards record license, tokenizer, revision, context limit, quantization, hardware + requirement, intended tasks, limitations, and evaluation evidence. Promotion to + `production` is rejected without evaluation references. +- Adapters bind to one immutable base revision, declare their dataset version and measured + quality delta, and cannot be promoted with unresolved regressions. +- Deployment revisions record container digest, model and adapter checksums, Ray and vLLM + configuration, GPU pool, and the previous revision used for rollback. + +| Endpoint | Purpose | +|---|---| +| `GET /v1/models` | OpenAI-compatible catalog enriched with tier, stage, license, and quantization. | +| `GET /v1/registry/models/{id}` | Full model card with its benchmark evidence. | +| `GET /v1/registry/adapters` | Promoted LoRA and QLoRA adapters. | +| `GET /v1/registry/deployments` | Deployment revisions and rollback targets. | + +Send `routing.domain` to request a domain adapter; the router applies the promoted adapter +with the largest measured quality gain for that base revision and task, or none at all. + ## Observability `GET /metrics` returns Prometheus exposition text and is intentionally unauthenticated so @@ -104,6 +127,7 @@ 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_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. | | `ROUTER_CACHE_TTL_SECONDS` | `300` | Exact-response entry lifetime. | diff --git a/config/registry.yaml b/config/registry.yaml new file mode 100644 index 0000000..2d62ebc --- /dev/null +++ b/config/registry.yaml @@ -0,0 +1,181 @@ +# Governed model catalog. Requests may never introduce a model path or revision; +# only records in this document are servable. +policy: + version: v1 + restricted_privacy_is_local_only: true + quality_floor: 0.0 + resource_ceiling: 2.0 + fallback_order: + - small-specialist + - general-local + - high-capability + - approved-external-fallback + +models: + - id: small-specialist + revision: mock-small@sha256:dev + tier: small-specialist + license: apache-2.0 + tokenizer: mock-small-tokenizer + context_limit: 8192 + quantization: awq + hardware: + accelerator: nvidia-l4 + count: 1 + minimum_memory_gb: 24 + supported_tasks: [extraction, classification] + quality: 0.82 + estimated_queue_ms: 12 + cost_weight: 0.1 + stage: production + intended_tasks: Field extraction and single-label classification. + limitations: Not evaluated for open-ended reasoning or multi-turn dialogue. + evaluation_references: + - benchmark:extraction-v3 + - benchmark:classification-v2 + + - id: general-local + revision: mock-general@sha256:dev + tier: general-local + license: apache-2.0 + tokenizer: mock-general-tokenizer + context_limit: 32768 + quantization: none + hardware: + accelerator: nvidia-a10g + count: 1 + minimum_memory_gb: 48 + supported_tasks: [extraction, classification, rag, summarization, general] + quality: 0.89 + estimated_queue_ms: 35 + cost_weight: 0.35 + stage: production + intended_tasks: Summarization, retrieval-augmented answers, and moderate reasoning. + limitations: Weaker than the high-capability tier on multi-step reasoning. + evaluation_references: + - benchmark:rag-v4 + + - id: high-capability + revision: mock-high@sha256:dev + tier: high-capability + license: apache-2.0 + tokenizer: mock-high-tokenizer + context_limit: 65536 + quantization: none + hardware: + accelerator: nvidia-a100 + count: 2 + minimum_memory_gb: 160 + tensor_parallel_size: 2 + supported_tasks: [extraction, classification, rag, summarization, reasoning, critique, general] + quality: 0.96 + estimated_queue_ms: 90 + cost_weight: 0.9 + stage: production + intended_tasks: Complex reasoning, critique, and difficult fallback. + limitations: Highest GPU cost per request; reserve for genuinely hard work. + evaluation_references: + - benchmark:reasoning-v5 + + - id: approved-external-fallback + revision: external-policy-v1 + tier: external-fallback + local: false + license: provider-terms + tokenizer: provider-managed + context_limit: 128000 + supported_tasks: [extraction, classification, rag, summarization, reasoning, critique, general] + quality: 0.98 + estimated_queue_ms: 45 + cost_weight: 1.5 + hardware: + accelerator: provider-managed + minimum_memory_gb: 1 + stage: production + intended_tasks: Unsupported capability or temporary local saturation only. + limitations: Never eligible for private or restricted data. + evaluation_references: + - benchmark:external-parity-v1 + +adapters: + - id: claims-extraction-lora + base_model_id: small-specialist + base_revision: mock-small@sha256:dev + adapter_revision: claims-lora@sha256:dev + domain: claims + intended_tasks: [extraction] + dataset_version: claims-2026-05 + quantized: true + stage: production + benchmark: + quality_delta: 0.06 + regressions: [] + + - id: support-classification-lora + base_model_id: small-specialist + base_revision: mock-small@sha256:dev + adapter_revision: support-lora@sha256:dev + domain: support + intended_tasks: [classification] + dataset_version: support-2026-04 + stage: production + benchmark: + quality_delta: 0.04 + regressions: [] + + - id: claims-extraction-lora-next + base_model_id: small-specialist + base_revision: mock-small@sha256:dev + adapter_revision: claims-lora@sha256:next + domain: claims + intended_tasks: [extraction] + dataset_version: claims-2026-06 + stage: staging + benchmark: + quality_delta: 0.02 + regressions: [] + +benchmarks: + - id: extraction-v3 + dataset_version: extraction-2026-05 + workload_version: steady-32 + hardware: nvidia-l4 + driver: "570.86" + container_digest: sha256:mock-container + engine_revision: mock-engine-0.1.0 + model_revision: mock-small@sha256:dev + concurrency: 32 + prompt_tokens_p50: 420 + prompt_tokens_p95: 1100 + engine_settings: + max_num_seqs: 64 + enable_prefix_caching: true + quality_score: 0.82 + latency_p95_ms: 480 + throughput_rps: 41.5 + gpu_seconds_per_request: 0.11 + +deployments: + - id: deploy-0002 + container_digest: sha256:mock-container + gpu_pool: l4-pool + model_checksums: + small-specialist: sha256:mock-small + general-local: sha256:mock-general + adapter_checksums: + claims-extraction-lora: sha256:mock-claims + ray_config: + num_replicas: 2 + max_ongoing_requests: 16 + vllm_config: + enable_prefix_caching: true + max_model_len: 8192 + stage: production + previous_revision_id: deploy-0001 + + - id: deploy-0001 + container_digest: sha256:mock-container-previous + gpu_pool: l4-pool + model_checksums: + small-specialist: sha256:mock-small-previous + stage: deprecated diff --git a/pyproject.toml b/pyproject.toml index 3bbdc88..587a0d0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,6 +11,7 @@ requires-python = ">=3.11" dependencies = [ "fastapi>=0.141.1,<1", "prometheus-client>=0.26.0,<1", + "pyyaml>=6.0.2,<7", "pydantic-settings>=2.15.0,<3", "uvicorn[standard]>=0.52.4,<1", ] @@ -23,6 +24,7 @@ dev = [ "pytest-asyncio>=1.4.0,<2", "pytest-cov>=7.1.0,<8", "ruff>=0.16.4,<1", + "types-PyYAML>=6.0.12,<7", ] [tool.hatch.build.targets.wheel] diff --git a/src/llm_router/app.py b/src/llm_router/app.py index fd257b6..a084f79 100644 --- a/src/llm_router/app.py +++ b/src/llm_router/app.py @@ -4,6 +4,7 @@ import uuid from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from pathlib import Path from fastapi import Depends, FastAPI, Header, HTTPException, Request, Response, status from fastapi.responses import JSONResponse @@ -37,20 +38,36 @@ Usage, ) from llm_router.observability import Metrics +from llm_router.registry import Registry, RegistryError, catalog_revisions, load_registry from llm_router.routing import NoEligibleModelError, Router, default_model_profiles +def _load_catalog(path: str) -> Registry | None: + """Load the governed catalog, falling back to built-in profiles when absent.""" + + if not Path(path).exists(): + return None + return load_registry(path) + + def create_app( settings: Settings | None = None, *, backend: InferenceBackend | None = None, metrics: Metrics | None = None, cache_store: CacheStore | None = None, + registry: Registry | None = None, ) -> FastAPI: runtime_settings = settings or get_settings() + catalog = registry if registry is not None else _load_catalog(runtime_settings.registry_path) + profiles = catalog.profiles() if catalog is not None else default_model_profiles() + policy_version = ( + catalog.policy.version if catalog is not None else runtime_settings.routing_policy_version + ) router = Router( - profiles=default_model_profiles(), + profiles=profiles, external_fallback_enabled=runtime_settings.external_fallback_enabled, + registry=catalog, ) admission = AdmissionController( runtime_settings.max_concurrency, @@ -64,12 +81,14 @@ def create_app( ttl_seconds=runtime_settings.cache_ttl_seconds, ) semantic_cache = SemanticCache(threshold=runtime_settings.semantic_similarity_threshold) - decision_cache = RouterDecisionCache(policy_version=runtime_settings.routing_policy_version) + decision_cache = RouterDecisionCache(policy_version=policy_version) prefix_tracker = PrefixTracker() - catalog_version = catalog_fingerprint( - (profile.revision for profile in router.profiles), - runtime_settings.routing_policy_version, + revisions = ( + catalog_revisions(catalog) + if catalog is not None + else (profile.revision for profile in router.profiles) ) + catalog_version = catalog_fingerprint(revisions, policy_version) @asynccontextmanager async def lifespan(app: FastAPI) -> AsyncIterator[None]: @@ -143,19 +162,75 @@ async def prometheus_metrics() -> Response: @app.get("/v1/models", dependencies=[Depends(authenticate)]) async def models() -> dict[str, object]: - visible = [ - { + cards = {card.id: card for card in catalog.servable_models()} if catalog else {} + visible: list[dict[str, object]] = [] + for profile in router.profiles: + if not profile.local and not runtime_settings.external_fallback_enabled: + continue + entry: dict[str, object] = { "id": profile.id, "object": "model", "owned_by": "local" if profile.local else "external-policy", "revision": profile.revision, "healthy": profile.healthy, + "context_limit": profile.context_limit, } - for profile in router.profiles - if profile.local or runtime_settings.external_fallback_enabled - ] + card = cards.get(profile.id) + if card is not None: + entry.update( + { + "tier": card.tier.value, + "license": card.license, + "tokenizer": card.tokenizer, + "quantization": card.quantization.value, + "stage": card.stage.value, + } + ) + visible.append(entry) return {"object": "list", "data": visible} + @app.get("/v1/registry/models/{model_id}", dependencies=[Depends(authenticate)]) + async def model_card(model_id: str) -> dict[str, object]: + if catalog is None: + raise HTTPException(status_code=404, detail="no catalog is configured") + try: + card = catalog.model_card(model_id) + except RegistryError as error: + raise HTTPException(status_code=404, detail=str(error)) from error + payload = card.model_dump(mode="json") + payload["benchmarks"] = [ + run.model_dump(mode="json") for run in catalog.benchmarks_for(card.revision) + ] + return payload + + @app.get("/v1/registry/adapters", dependencies=[Depends(authenticate)]) + async def adapters() -> dict[str, object]: + if catalog is None: + return {"object": "list", "data": []} + return { + "object": "list", + "data": [adapter.model_dump(mode="json") for adapter in catalog.servable_adapters()], + } + + @app.get("/v1/registry/deployments", dependencies=[Depends(authenticate)]) + async def deployments() -> dict[str, object]: + if catalog is None: + return {"object": "list", "data": []} + return { + "object": "list", + "data": [ + { + **revision.model_dump(mode="json"), + "rollback_target": ( + target.id + if (target := catalog.rollback_target(revision.id)) is not None + else None + ), + } + for revision in catalog.deployments + ], + } + def _completion_response( *, model_id: str, @@ -299,6 +374,8 @@ async def chat_completions( await _store_cache(payload, prompt, cache_key, subject, decision, result) response.headers["X-Cache"] = "miss" + if decision.adapter_id is not None: + response.headers["X-Route-Adapter"] = decision.adapter_id response.headers["X-Route-Model"] = decision.profile.id response.headers["X-Route-Revision"] = decision.profile.revision response.headers["X-Route-Reason"] = decision.reason @@ -319,6 +396,8 @@ async def chat_completions( ), routing={ "model_revision": decision.profile.revision, + "adapter_id": decision.adapter_id, + "adapter_revision": decision.adapter_revision, "task": decision.task.value, "reason": decision.reason, "score": decision.score, diff --git a/src/llm_router/config.py b/src/llm_router/config.py index 163c124..6d49aa2 100644 --- a/src/llm_router/config.py +++ b/src/llm_router/config.py @@ -15,6 +15,7 @@ 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 + registry_path: str = "config/registry.yaml" routing_policy_version: str = "v1" cache_enabled: bool = True cache_ttl_seconds: float = Field(default=300.0, gt=0) diff --git a/src/llm_router/models.py b/src/llm_router/models.py index 05b6f6a..aa981a8 100644 --- a/src/llm_router/models.py +++ b/src/llm_router/models.py @@ -27,6 +27,7 @@ class ChatMessage(BaseModel): class RoutingOptions(BaseModel): task: TaskClass | None = None + domain: str | None = Field(default=None, max_length=64) privacy: PrivacyClass = PrivacyClass.PRIVATE latency_tier: Literal["interactive", "standard", "batch"] = "standard" quality_floor: float = Field(default=0.0, ge=0.0, le=1.0) @@ -70,6 +71,8 @@ class RouteDecision(BaseModel): reason: str score: float candidate_count: int + adapter_id: str | None = None + adapter_revision: str | None = None class ChatCompletionChoice(BaseModel): diff --git a/src/llm_router/registry.py b/src/llm_router/registry.py new file mode 100644 index 0000000..c1672f7 --- /dev/null +++ b/src/llm_router/registry.py @@ -0,0 +1,261 @@ +"""Model, adapter, benchmark, and deployment records from section 11. + +The registry is the governed source of truth for what may be served. It is +loaded from a declarative catalog rather than from request input, so a caller +can never introduce a model path, revision, or adapter. +""" + +from collections.abc import Iterable +from enum import StrEnum +from pathlib import Path +from typing import Any + +import yaml +from pydantic import BaseModel, Field, model_validator + +from llm_router.models import ModelProfile, TaskClass + + +class Quantization(StrEnum): + NONE = "none" + AWQ = "awq" + GPTQ = "gptq" + INT8 = "int8" + + +class LifecycleStage(StrEnum): + DEVELOPMENT = "development" + STAGING = "staging" + PRODUCTION = "production" + DEPRECATED = "deprecated" + + +class ModelTier(StrEnum): + SMALL_SPECIALIST = "small-specialist" + GENERAL_LOCAL = "general-local" + HIGH_CAPABILITY = "high-capability" + EXTERNAL_FALLBACK = "external-fallback" + + +class HardwareRequirement(BaseModel): + accelerator: str + count: int = Field(default=1, ge=1) + minimum_memory_gb: int = Field(ge=1) + tensor_parallel_size: int = Field(default=1, ge=1) + + +class ModelCard(BaseModel): + """Governance record required for every servable model revision.""" + + id: str + revision: str + tier: ModelTier + local: bool = True + license: str + tokenizer: str + context_limit: int = Field(ge=1) + quantization: Quantization = Quantization.NONE + hardware: HardwareRequirement + supported_tasks: frozenset[TaskClass] + quality: float = Field(ge=0.0, le=1.0) + estimated_queue_ms: int = Field(default=0, ge=0) + cost_weight: float = Field(default=0.0, ge=0.0) + stage: LifecycleStage = LifecycleStage.DEVELOPMENT + healthy: bool = True + intended_tasks: str + limitations: str + evaluation_references: tuple[str, ...] = () + + @model_validator(mode="after") + def require_evidence_for_production(self) -> "ModelCard": + if self.stage is LifecycleStage.PRODUCTION and not self.evaluation_references: + raise ValueError(f"model {self.id} cannot reach production without evaluation evidence") + return self + + def to_profile(self) -> ModelProfile: + return ModelProfile( + id=self.id, + revision=self.revision, + local=self.local, + healthy=self.healthy, + context_limit=self.context_limit, + supported_tasks=self.supported_tasks, + quality=self.quality, + estimated_queue_ms=self.estimated_queue_ms, + cost_weight=self.cost_weight, + ) + + +class BenchmarkDelta(BaseModel): + quality_delta: float + regressions: tuple[str, ...] = () + + +class AdapterProfile(BaseModel): + """LoRA or QLoRA adapter bound to one immutable base-model revision.""" + + id: str + base_model_id: str + base_revision: str + adapter_revision: str + domain: str + intended_tasks: frozenset[TaskClass] + dataset_version: str + benchmark: BenchmarkDelta + stage: LifecycleStage = LifecycleStage.DEVELOPMENT + quantized: bool = False + + @model_validator(mode="after") + def block_regressed_adapters_from_production(self) -> "AdapterProfile": + if self.stage is LifecycleStage.PRODUCTION and self.benchmark.regressions: + raise ValueError(f"adapter {self.id} has unresolved regressions: {self.benchmark}") + return self + + +class BenchmarkRun(BaseModel): + """Reproducibility record for one quality or load measurement.""" + + id: str + dataset_version: str + workload_version: str + hardware: str + driver: str + container_digest: str + engine_revision: str + model_revision: str + adapter_revision: str | None = None + concurrency: int = Field(ge=1) + prompt_tokens_p50: int = Field(ge=1) + prompt_tokens_p95: int = Field(ge=1) + engine_settings: dict[str, Any] = Field(default_factory=dict) + quality_score: float = Field(ge=0.0, le=1.0) + latency_p95_ms: float = Field(ge=0.0) + throughput_rps: float = Field(ge=0.0) + gpu_seconds_per_request: float = Field(default=0.0, ge=0.0) + + +class DeploymentRevision(BaseModel): + """Immutable description of what is deployed, and what it rolls back to.""" + + id: str + container_digest: str + model_checksums: dict[str, str] + adapter_checksums: dict[str, str] = Field(default_factory=dict) + ray_config: dict[str, Any] = Field(default_factory=dict) + vllm_config: dict[str, Any] = Field(default_factory=dict) + gpu_pool: str + stage: LifecycleStage = LifecycleStage.STAGING + previous_revision_id: str | None = None + + +class RoutePolicy(BaseModel): + """Operator-owned routing constraints applied before model scoring.""" + + version: str = "v1" + eligible_models: frozenset[str] = frozenset() + eligible_adapters: frozenset[str] = frozenset() + restricted_privacy_is_local_only: bool = True + quality_floor: float = Field(default=0.0, ge=0.0, le=1.0) + resource_ceiling: float = Field(default=10.0, gt=0.0) + fallback_order: tuple[str, ...] = () + + +class RegistryError(RuntimeError): + """Raised when a catalog is inconsistent or references an unknown record.""" + + +class Registry(BaseModel): + """In-memory view of the governed catalog.""" + + models: tuple[ModelCard, ...] + adapters: tuple[AdapterProfile, ...] = () + benchmarks: tuple[BenchmarkRun, ...] = () + deployments: tuple[DeploymentRevision, ...] = () + policy: RoutePolicy = RoutePolicy() + + @model_validator(mode="after") + def validate_references(self) -> "Registry": + model_ids = {card.id for card in self.models} + if len(model_ids) != len(self.models): + raise RegistryError("duplicate model identifiers in catalog") + revisions = {(card.id, card.revision) for card in self.models} + for adapter in self.adapters: + if (adapter.base_model_id, adapter.base_revision) not in revisions: + raise RegistryError( + f"adapter {adapter.id} references unknown base " + f"{adapter.base_model_id}@{adapter.base_revision}" + ) + return self + + def servable_models(self) -> tuple[ModelCard, ...]: + return tuple( + card + for card in self.models + if card.stage in {LifecycleStage.STAGING, LifecycleStage.PRODUCTION} + and (not self.policy.eligible_models or card.id in self.policy.eligible_models) + ) + + def profiles(self) -> tuple[ModelProfile, ...]: + return tuple(card.to_profile() for card in self.servable_models()) + + def model_card(self, model_id: str) -> ModelCard: + for card in self.models: + if card.id == model_id: + return card + raise RegistryError(f"unknown model {model_id}") + + def servable_adapters(self) -> tuple[AdapterProfile, ...]: + servable = {card.id for card in self.servable_models()} + return tuple( + adapter + for adapter in self.adapters + if adapter.stage in {LifecycleStage.STAGING, LifecycleStage.PRODUCTION} + and adapter.base_model_id in servable + and (not self.policy.eligible_adapters or adapter.id in self.policy.eligible_adapters) + ) + + def select_adapter( + self, *, model_id: str, revision: str, domain: str | None, task: TaskClass + ) -> AdapterProfile | None: + """Pick the best-scoring approved adapter for a base revision and domain.""" + + if domain is None: + return None + candidates = [ + adapter + for adapter in self.servable_adapters() + if adapter.base_model_id == model_id + and adapter.base_revision == revision + and adapter.domain == domain + and task in adapter.intended_tasks + ] + if not candidates: + return None + return max(candidates, key=lambda adapter: adapter.benchmark.quality_delta) + + def benchmarks_for(self, model_revision: str) -> tuple[BenchmarkRun, ...]: + return tuple(run for run in self.benchmarks if run.model_revision == model_revision) + + def rollback_target(self, deployment_id: str) -> DeploymentRevision | None: + current = next((item for item in self.deployments if item.id == deployment_id), None) + if current is None: + raise RegistryError(f"unknown deployment {deployment_id}") + if current.previous_revision_id is None: + return None + return next( + (item for item in self.deployments if item.id == current.previous_revision_id), None + ) + + +def load_registry(path: str | Path) -> Registry: + """Load and validate a catalog from a YAML document.""" + + document = yaml.safe_load(Path(path).read_text(encoding="utf-8")) + if not isinstance(document, dict): + raise RegistryError(f"catalog {path} must contain a mapping") + return Registry.model_validate(document) + + +def catalog_revisions(registry: Registry) -> Iterable[str]: + yield from (card.revision for card in registry.servable_models()) + yield from (adapter.adapter_revision for adapter in registry.servable_adapters()) diff --git a/src/llm_router/routing.py b/src/llm_router/routing.py index 9d21dcb..2297d9e 100644 --- a/src/llm_router/routing.py +++ b/src/llm_router/routing.py @@ -1,4 +1,5 @@ from dataclasses import dataclass +from typing import TYPE_CHECKING from llm_router.models import ( ChatCompletionRequest, @@ -8,6 +9,9 @@ TaskClass, ) +if TYPE_CHECKING: # pragma: no cover - import cycle guard for type checking only + from llm_router.registry import Registry + class NoEligibleModelError(RuntimeError): """Raised when policy removes every model candidate.""" @@ -70,6 +74,7 @@ def default_model_profiles() -> tuple[ModelProfile, ...]: class Router: profiles: tuple[ModelProfile, ...] external_fallback_enabled: bool = False + registry: "Registry | None" = None def classify_task(self, request: ChatCompletionRequest) -> TaskClass: if request.routing.task is not None: @@ -130,16 +135,34 @@ def score(profile: ModelProfile) -> float: ) selected = max(candidates, key=score) + adapter = ( + self.registry.select_adapter( + model_id=selected.id, + revision=selected.revision, + domain=request.routing.domain, + task=task, + ) + if self.registry is not None + else None + ) + reason = ( + f"selected highest policy score among {len(candidates)} eligible model(s); " + f"task={task.value}, privacy={request.routing.privacy.value}, " + f"latency_tier={request.routing.latency_tier}" + ) + if adapter is not None: + reason += ( + f"; applied adapter {adapter.id} for domain {adapter.domain} " + f"(measured quality delta {adapter.benchmark.quality_delta:+.3f})" + ) return RouteDecision( profile=selected, task=task, - reason=( - f"selected highest policy score among {len(candidates)} eligible model(s); " - f"task={task.value}, privacy={request.routing.privacy.value}, " - f"latency_tier={request.routing.latency_tier}" - ), + reason=reason, score=round(score(selected), 3), candidate_count=len(candidates), + adapter_id=None if adapter is None else adapter.id, + adapter_revision=None if adapter is None else adapter.adapter_revision, ) @staticmethod diff --git a/tests/e2e/api.spec.ts b/tests/e2e/api.spec.ts index 8cbe04a..75268d1 100644 --- a/tests/e2e/api.spec.ts +++ b/tests/e2e/api.spec.ts @@ -94,3 +94,35 @@ test("never caches restricted-class requests", async ({ request }) => { expect(repeat.status()).toBe(200); expect(repeat.headers()["x-cache"]).toBe("miss"); }); + +test("applies a registered LoRA adapter for a domain request", async ({ request }) => { + const response = await request.post("/v1/chat/completions", { + data: { + model: "auto", + messages: [{ role: "user", content: "Extract the claim fields as JSON" }], + routing: { privacy: "restricted", domain: "claims", task: "extraction" }, + }, + }); + + expect(response.status()).toBe(200); + expect(response.headers()["x-route-adapter"]).toBe("claims-extraction-lora"); + const body = await response.json(); + expect(body.model).toBe("small-specialist"); + expect(body.routing.adapter_revision).toBe("claims-lora@sha256:dev"); +}); + +test("publishes model cards and deployment rollback targets", async ({ request }) => { + const card = await request.get("/v1/registry/models/small-specialist"); + expect(card.status()).toBe(200); + const cardBody = await card.json(); + expect(cardBody.license).toBeTruthy(); + expect(cardBody.limitations).toBeTruthy(); + expect(cardBody.evaluation_references.length).toBeGreaterThan(0); + + const deployments = await request.get("/v1/registry/deployments"); + expect(deployments.status()).toBe(200); + const current = (await deployments.json()).data.find( + (item: { id: string }) => item.id === "deploy-0002", + ); + expect(current.rollback_target).toBe("deploy-0001"); +}); diff --git a/tests/integration/test_registry_api.py b/tests/integration/test_registry_api.py new file mode 100644 index 0000000..7c86364 --- /dev/null +++ b/tests/integration/test_registry_api.py @@ -0,0 +1,132 @@ +from fastapi.testclient import TestClient + +from llm_router.app import create_app +from llm_router.config import Settings + +HEADERS = {"Authorization": "Bearer registry-key"} + + +def build_client(**overrides: object) -> TestClient: + settings = Settings(api_keys="registry-key", **overrides) # type: ignore[arg-type] + return TestClient(create_app(settings)) + + +def test_model_catalog_is_served_from_the_governed_registry() -> None: + with build_client() as client: + payload = client.get("/v1/models", headers=HEADERS).json() + + entries = {model["id"]: model for model in payload["data"]} + assert entries["small-specialist"]["quantization"] == "awq" + assert entries["small-specialist"]["tier"] == "small-specialist" + assert entries["small-specialist"]["stage"] == "production" + assert entries["small-specialist"]["license"] == "apache-2.0" + assert "approved-external-fallback" not in entries + + +def test_model_card_endpoint_returns_governance_and_benchmark_evidence() -> None: + with build_client() as client: + response = client.get("/v1/registry/models/small-specialist", headers=HEADERS) + missing = client.get("/v1/registry/models/absent", headers=HEADERS) + + assert response.status_code == 200 + card = response.json() + assert card["limitations"] + assert card["hardware"]["accelerator"] == "nvidia-l4" + assert card["evaluation_references"] + assert [run["id"] for run in card["benchmarks"]] == ["extraction-v3"] + assert missing.status_code == 404 + + +def test_registry_endpoints_require_authentication() -> None: + with build_client() as client: + assert client.get("/v1/registry/adapters").status_code == 401 + assert client.get("/v1/registry/deployments").status_code == 401 + assert client.get("/v1/registry/models/small-specialist").status_code == 401 + + +def test_adapter_catalog_lists_only_promoted_adapters() -> None: + with build_client() as client: + adapters = client.get("/v1/registry/adapters", headers=HEADERS).json()["data"] + + identifiers = {adapter["id"] for adapter in adapters} + assert {"claims-extraction-lora", "support-classification-lora"} <= identifiers + claims = next(adapter for adapter in adapters if adapter["id"] == "claims-extraction-lora") + assert claims["dataset_version"] == "claims-2026-05" + assert claims["benchmark"]["quality_delta"] > 0 + + +def test_deployment_history_exposes_the_rollback_target() -> None: + with build_client() as client: + deployments = client.get("/v1/registry/deployments", headers=HEADERS).json()["data"] + + current = next(item for item in deployments if item["id"] == "deploy-0002") + assert current["rollback_target"] == "deploy-0001" + assert current["gpu_pool"] == "l4-pool" + assert current["vllm_config"]["enable_prefix_caching"] is True + + +def test_domain_request_is_served_by_the_matching_adapter() -> None: + with build_client() as client: + response = client.post( + "/v1/chat/completions", + headers=HEADERS, + json={ + "model": "auto", + "messages": [{"role": "user", "content": "Extract the claim fields"}], + "routing": {"privacy": "restricted", "domain": "claims", "task": "extraction"}, + }, + ) + + assert response.status_code == 200 + body = response.json() + assert body["model"] == "small-specialist" + assert body["routing"]["adapter_id"] == "claims-extraction-lora" + assert body["routing"]["adapter_revision"] == "claims-lora@sha256:dev" + assert response.headers["x-route-adapter"] == "claims-extraction-lora" + assert "applied adapter claims-extraction-lora" in body["routing"]["reason"] + + +def test_request_without_a_domain_uses_the_base_model_only() -> None: + with build_client() as client: + response = client.post( + "/v1/chat/completions", + headers=HEADERS, + json={ + "model": "auto", + "messages": [{"role": "user", "content": "Extract the claim fields"}], + "routing": {"privacy": "restricted", "task": "extraction"}, + }, + ) + + assert response.json()["routing"]["adapter_id"] is None + assert "x-route-adapter" not in response.headers + + +def test_unknown_domain_does_not_invent_an_adapter() -> None: + with build_client() as client: + response = client.post( + "/v1/chat/completions", + headers=HEADERS, + json={ + "model": "auto", + "messages": [{"role": "user", "content": "Extract the fields"}], + "routing": { + "privacy": "restricted", + "domain": "unregistered", + "task": "extraction", + }, + }, + ) + + assert response.status_code == 200 + assert response.json()["routing"]["adapter_id"] is None + + +def test_gateway_falls_back_to_built_in_profiles_without_a_catalog() -> None: + with build_client(registry_path="config/does-not-exist.yaml") as client: + payload = client.get("/v1/models", headers=HEADERS).json() + + entries = {model["id"]: model for model in payload["data"]} + assert "small-specialist" in entries + assert "tier" not in entries["small-specialist"] + assert client.get("/v1/registry/adapters", headers=HEADERS).json()["data"] == [] diff --git a/tests/unit/test_registry.py b/tests/unit/test_registry.py new file mode 100644 index 0000000..def11bd --- /dev/null +++ b/tests/unit/test_registry.py @@ -0,0 +1,206 @@ +from pathlib import Path + +import pytest +from pydantic import ValidationError + +from llm_router.models import TaskClass +from llm_router.registry import ( + AdapterProfile, + BenchmarkDelta, + HardwareRequirement, + LifecycleStage, + ModelCard, + ModelTier, + Quantization, + Registry, + RegistryError, + catalog_revisions, + load_registry, +) + +CATALOG = Path("config/registry.yaml") + + +def make_card(**overrides: object) -> ModelCard: + values: dict[str, object] = { + "id": "small-specialist", + "revision": "rev-1", + "tier": ModelTier.SMALL_SPECIALIST, + "license": "apache-2.0", + "tokenizer": "tok", + "context_limit": 8192, + "quantization": Quantization.AWQ, + "hardware": HardwareRequirement(accelerator="nvidia-l4", minimum_memory_gb=24), + "supported_tasks": frozenset({TaskClass.EXTRACTION}), + "quality": 0.82, + "stage": LifecycleStage.PRODUCTION, + "intended_tasks": "extraction", + "limitations": "narrow", + "evaluation_references": ("benchmark:extraction-v3",), + } + values.update(overrides) + return ModelCard.model_validate(values) + + +def make_adapter(**overrides: object) -> AdapterProfile: + values: dict[str, object] = { + "id": "claims-lora", + "base_model_id": "small-specialist", + "base_revision": "rev-1", + "adapter_revision": "claims@1", + "domain": "claims", + "intended_tasks": frozenset({TaskClass.EXTRACTION}), + "dataset_version": "claims-2026-05", + "benchmark": BenchmarkDelta(quality_delta=0.06), + "stage": LifecycleStage.PRODUCTION, + } + values.update(overrides) + return AdapterProfile.model_validate(values) + + +def test_catalog_file_loads_and_exposes_production_records() -> None: + registry = load_registry(CATALOG) + + assert {card.id for card in registry.servable_models()} == { + "small-specialist", + "general-local", + "high-capability", + "approved-external-fallback", + } + assert registry.policy.version == "v1" + assert len(list(catalog_revisions(registry))) == 7 + + +def test_model_card_requires_evaluation_evidence_for_production() -> None: + with pytest.raises(ValidationError, match="evaluation evidence"): + make_card(evaluation_references=()) + + +def test_model_card_without_evidence_is_allowed_below_production() -> None: + card = make_card(stage=LifecycleStage.STAGING, evaluation_references=()) + + assert card.stage is LifecycleStage.STAGING + + +def test_adapter_with_regressions_cannot_reach_production() -> None: + with pytest.raises(ValidationError, match="unresolved regressions"): + make_adapter(benchmark=BenchmarkDelta(quality_delta=0.01, regressions=("json-validity",))) + + +def test_registry_rejects_adapter_with_unknown_base_revision() -> None: + with pytest.raises(RegistryError, match="unknown base"): + Registry(models=(make_card(),), adapters=(make_adapter(base_revision="rev-9"),)) + + +def test_registry_rejects_duplicate_model_identifiers() -> None: + with pytest.raises(RegistryError, match="duplicate model identifiers"): + Registry(models=(make_card(), make_card(revision="rev-2"))) + + +def test_profiles_exclude_development_and_deprecated_models() -> None: + registry = Registry( + models=( + make_card(), + make_card(id="draft", revision="rev-2", stage=LifecycleStage.DEVELOPMENT), + make_card(id="old", revision="rev-3", stage=LifecycleStage.DEPRECATED), + ) + ) + + assert {profile.id for profile in registry.profiles()} == {"small-specialist"} + + +def test_policy_eligibility_narrows_servable_models_and_adapters() -> None: + registry = Registry( + models=(make_card(), make_card(id="general-local", revision="rev-2")), + adapters=(make_adapter(),), + policy={"eligible_models": {"general-local"}}, # type: ignore[arg-type] + ) + + assert {card.id for card in registry.servable_models()} == {"general-local"} + assert registry.servable_adapters() == () + + +def test_select_adapter_prefers_the_largest_measured_quality_gain() -> None: + registry = Registry( + models=(make_card(),), + adapters=( + make_adapter(), + make_adapter( + id="claims-lora-next", + adapter_revision="claims@2", + benchmark=BenchmarkDelta(quality_delta=0.09), + ), + ), + ) + + selected = registry.select_adapter( + model_id="small-specialist", revision="rev-1", domain="claims", task=TaskClass.EXTRACTION + ) + + assert selected is not None + assert selected.id == "claims-lora-next" + + +def test_select_adapter_requires_domain_task_and_matching_base_revision() -> None: + registry = Registry(models=(make_card(),), adapters=(make_adapter(),)) + + assert ( + registry.select_adapter( + model_id="small-specialist", revision="rev-1", domain=None, task=TaskClass.EXTRACTION + ) + is None + ) + assert ( + registry.select_adapter( + model_id="small-specialist", + revision="rev-1", + domain="support", + task=TaskClass.EXTRACTION, + ) + is None + ) + assert ( + registry.select_adapter( + model_id="small-specialist", + revision="rev-1", + domain="claims", + task=TaskClass.CLASSIFICATION, + ) + is None + ) + assert ( + registry.select_adapter( + model_id="small-specialist", + revision="rev-9", + domain="claims", + task=TaskClass.EXTRACTION, + ) + is None + ) + + +def test_model_card_lookup_reports_unknown_identifiers() -> None: + registry = load_registry(CATALOG) + + assert registry.model_card("general-local").license == "apache-2.0" + with pytest.raises(RegistryError, match="unknown model"): + registry.model_card("absent") + + +def test_benchmarks_and_rollback_targets_are_resolvable() -> None: + registry = load_registry(CATALOG) + + assert [run.id for run in registry.benchmarks_for("mock-small@sha256:dev")] == ["extraction-v3"] + target = registry.rollback_target("deploy-0002") + assert target is not None and target.id == "deploy-0001" + assert registry.rollback_target("deploy-0001") is None + with pytest.raises(RegistryError, match="unknown deployment"): + registry.rollback_target("deploy-9999") + + +def test_loader_rejects_a_non_mapping_document(tmp_path: Path) -> None: + path = tmp_path / "catalog.yaml" + path.write_text("- not-a-mapping\n", encoding="utf-8") + + with pytest.raises(RegistryError, match="must contain a mapping"): + load_registry(path)