fix: hold streamed requests to the concurrency bound and define engine failure

README.md

@@ -215,6 +215,29 @@ be served. A request can never introduce a model path, revision, or adapter.
 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.
 
+## Failure behaviour
+
+| Condition | What the gateway does |
+|---|---|
+| Capacity saturated | `503` `overloaded` after the bounded admission wait, with `Retry-After: 1`. A streamed response holds its slot until the stream ends, so streaming cannot escape the concurrency bound. |
+| GPU out of memory | `503` `engine_out_of_memory`, `Retry-After: 10`. The same request at the same size will fail again, so the message says to shorten the prompt or lower `max_tokens`. Counts toward the circuit. |
+| Node lost or engine unreachable | `502` `backend_unavailable` for each of the first `ROUTER_ENGINE_FAILURE_THRESHOLD` consecutive failures; then the circuit opens. |
+| Circuit open | `503` `engine_unavailable` immediately, without contacting the engine, with `Retry-After` set to the remaining cooldown. `/readyz` fails, so the orchestrator stops routing here. |
+| Engine recovers | A healthy probe ends the cooldown early and lets exactly one trial request through. Only that request succeeding closes the circuit; if it fails, the cooldown restarts. |
+| Shutdown | `/readyz` fails first, then admitted requests get `ROUTER_SHUTDOWN_GRACE_SECONDS` to finish before the engine client closes. Keep it under the pod's `terminationGracePeriodSeconds` (60). |
+
+An engine error body is inspected for an out-of-memory report and then discarded, never forwarded:
+it can echo the prompt it rejected. `router_engine_circuit_open` reports 0 closed, 0.5 half-open,
+1 open.
+
+A request that fails is not retried on another model. Every local model shares the one engine a
+gateway faces, so a retry would meet the same failure; the caller is told when to come back.
+
+Cold start is measured, not assumed (see `router_model_load_seconds` below). Tiers that keep a
+warm replica never pay it on the request path. The high-capability tier scales to zero, so its
+first request after an idle period waits for a full model load: budget for it, or raise its
+`min_replicas`.
+
 ## Tenants
 
 What a caller may use is governance and lives in the catalog; the credential that proves which
@@ -380,6 +403,9 @@ All settings use the `ROUTER_` prefix.
 | `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_ENGINE_FAILURE_THRESHOLD` | `5` | Consecutive engine failures before the circuit opens. |
+| `ROUTER_ENGINE_COOLDOWN_SECONDS` | `30` | How long the circuit stays open before a trial request. |
+| `ROUTER_SHUTDOWN_GRACE_SECONDS` | `20` | How long shutdown waits for admitted requests. |
 | `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. |

src/llm_router/admission.py

@@ -17,17 +17,51 @@ class AdmissionController:
     def __init__(self, max_concurrency: int, timeout_seconds: float) -> None:
         self._semaphore = asyncio.Semaphore(max_concurrency)
         self._timeout_seconds = timeout_seconds
+        self._in_flight = 0
+        self._idle = asyncio.Event()
+        self._idle.set()
+
+    @property
+    def in_flight(self) -> int:
+        return self._in_flight
+
+    async def acquire(self) -> None:
+        """Take a slot, or refuse once the bounded wait has run out.
+
+        A streamed response outlives the handler that admitted it, so it
+        holds its slot through `acquire` and `release` rather than a block
+        that would end, and free the slot, before generation began.
+        """
 
-    @asynccontextmanager
-    async def slot(self) -> AsyncIterator[None]:
         try:
             await asyncio.wait_for(self._semaphore.acquire(), timeout=self._timeout_seconds)
         except TimeoutError as error:
             raise AdmissionRejectedError("inference capacity is saturated") from error
+        self._in_flight += 1
+        self._idle.clear()
+
+    def release(self) -> None:
+        self._in_flight -= 1
+        if self._in_flight == 0:
+            self._idle.set()
+        self._semaphore.release()
+
+    @asynccontextmanager
+    async def slot(self) -> AsyncIterator[None]:
+        await self.acquire()
         try:
             yield
         finally:
-            self._semaphore.release()
+            self.release()
+
+    async def drain(self, timeout_seconds: float) -> bool:
+        """Wait for admitted work to finish; False if the grace period ran out."""
+
+        try:
+            await asyncio.wait_for(self._idle.wait(), timeout=timeout_seconds)
+        except TimeoutError:
+            return False
+        return True
 
 
 class SlidingWindowQuota:

src/llm_router/app.py

@@ -7,7 +7,7 @@
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Annotated
+from typing import Annotated, Any
 
 import httpx
 from fastapi import Depends, FastAPI, Header, HTTPException, Request, Response, status
@@ -21,6 +21,7 @@
     SlidingWindowQuota,
 )
 from llm_router.backends import (
+    BackendOutOfMemoryError,
     BackendResult,
     BackendUnavailableError,
     InferenceBackend,
@@ -62,6 +63,7 @@
     load_registry,
     strictest_privacy,
 )
+from llm_router.resilience import CircuitBreaker, EngineCircuitOpenError, ResilientBackend
 from llm_router.routing import NoEligibleModelError, Router, default_model_profiles
 from llm_router.tracing import RequestSpan, Tracing, build_tracer_provider
 
@@ -144,7 +146,7 @@ def create_app(
     engine_client = (
         httpx.AsyncClient() if backend is None and runtime_settings.backend == "vllm" else None
     )
-    inference_backend: InferenceBackend = backend or (
+    raw_backend: InferenceBackend = backend or (
         VLLMBackend(
             base_url=runtime_settings.vllm_base_url,
             client=engine_client,
@@ -153,6 +155,11 @@ def create_app(
         if engine_client is not None
         else MockInferenceBackend()
     )
+    circuit = CircuitBreaker(
+        failure_threshold=runtime_settings.engine_failure_threshold,
+        cooldown_seconds=runtime_settings.engine_cooldown_seconds,
+    )
+    inference_backend: InferenceBackend = ResilientBackend(raw_backend, circuit)
     telemetry = metrics if metrics is not None else Metrics()
     engine_telemetry = engine_stats or (
         EngineStatsCollector(base_url=runtime_settings.vllm_base_url, client=engine_client)
@@ -193,7 +200,11 @@ def create_app(
     async def lifespan(app: FastAPI) -> AsyncIterator[None]:
         app.state.ready = True
         yield
+        # Readiness fails first so no new traffic arrives, then admitted
+        # requests are given the grace period to finish before the engine
+        # client they depend on is closed underneath them.
         app.state.ready = False
+        await admission.drain(runtime_settings.shutdown_grace_seconds)
         if engine_client is not None:
             await engine_client.aclose()
 
@@ -246,6 +257,26 @@ async def admission_handler(_: Request, error: AdmissionRejectedError) -> JSONRe
             content={"error": {"message": str(error), "type": "overloaded"}},
         )
 
+    @app.exception_handler(BackendOutOfMemoryError)
+    async def out_of_memory_handler(_: Request, error: BackendOutOfMemoryError) -> JSONResponse:
+        # Not a retry-as-is condition: the same request at the same size will
+        # fail again, so the message says what to change.
+        telemetry.record_rejection("engine_out_of_memory")
+        return JSONResponse(
+            status_code=503,
+            headers={"Retry-After": "10"},
+            content={"error": {"message": str(error), "type": "engine_out_of_memory"}},
+        )
+
+    @app.exception_handler(EngineCircuitOpenError)
+    async def circuit_handler(_: Request, error: EngineCircuitOpenError) -> JSONResponse:
+        telemetry.record_rejection("engine_unavailable")
+        return JSONResponse(
+            status_code=503,
+            headers={"Retry-After": str(max(1, round(error.retry_after_seconds)))},
+            content={"error": {"message": str(error), "type": "engine_unavailable"}},
+        )
+
     @app.exception_handler(BackendUnavailableError)
     async def backend_handler(_: Request, error: BackendUnavailableError) -> JSONResponse:
         telemetry.record_rejection("backend_unavailable")
@@ -287,6 +318,7 @@ async def prometheus_metrics() -> Response:
         # Engine and accelerator state is pulled through on scrape so the
         # gateway stays the single scrape target for the whole serving path and
         # no background poller runs when nobody is collecting.
+        telemetry.record_circuit_state(circuit.state, engine=engine_label)
         if engine_telemetry is not None:
             stats = await engine_telemetry.sample()
             if stats is not None:
@@ -513,6 +545,39 @@ def _chunk(completion_id: str, created: int, model_id: str, **choice: object) ->
         }
         return f"data: {json.dumps(document)}\n\n"
 
+    class _StreamLease:
+        """What a live stream holds until it ends: a slot, a gauge, a span."""
+
+        def __init__(self, span: RequestSpan) -> None:
+            self._span = span
+            self._closed = False
+
+        def close(self) -> None:
+            if self._closed:
+                return
+            self._closed = True
+            telemetry.inflight_requests.dec()
+            admission.release()
+            self._span.end()
+
+    class _LeasedStreamingResponse(StreamingResponse):
+        """Releases the lease even if the body iterator is never started.
+
+        A generator that is never iterated never runs its own cleanup, which
+        happens when the client disconnects before the first chunk. The
+        response is always invoked, so it closes the lease as well.
+        """
+
+        def __init__(self, *args: Any, lease: _StreamLease, **kwargs: Any) -> None:
+            super().__init__(*args, **kwargs)
+            self._lease = lease
+
+        async def __call__(self, scope: Any, receive: Any, send: Any) -> None:
+            try:
+                await super().__call__(scope, receive, send)
+            finally:
+                self._lease.close()
+
     async def _stream_completion(
         payload: ChatCompletionRequest,
         prompt: str,
@@ -522,6 +587,7 @@ async def _stream_completion(
         started: float,
         queue_seconds: float,
         span: RequestSpan,
+        lease: "_StreamLease",
     ) -> AsyncIterator[str]:
         completion_id = f"chatcmpl-{uuid.uuid4().hex}"
         created = int(time.time())
@@ -565,7 +631,7 @@ async def _stream_completion(
             span.fail(error)
             raise
         finally:
-            span.end()
+            lease.close()
 
     @app.post("/v1/chat/completions", response_model=None)
     async def chat_completions(
@@ -679,36 +745,43 @@ async def _complete(
 
         telemetry.queued_requests.inc()
         try:
-            async with admission.slot():
-                telemetry.queued_requests.dec()
-                queue_seconds = time.perf_counter() - started
-                # What this request actually waited becomes the next
-                # request's estimate for the same model.
-                load.observe_queue(decision.profile.id, queue_seconds * 1000)
-                telemetry.inflight_requests.inc()
-                try:
-                    if payload.stream:
-                        span.handed_off = True
-                        return StreamingResponse(
-                            _stream_completion(
-                                payload,
-                                prompt,
-                                cache_key,
-                                subject,
-                                decision,
-                                started,
-                                queue_seconds,
-                                span,
-                            ),
-                            media_type="text/event-stream",
-                            headers=_route_headers(decision, "miss"),
-                        )
-                    result = await inference_backend.generate(payload, decision)
-                finally:
-                    telemetry.inflight_requests.dec()
-        except AdmissionRejectedError:
+            await admission.acquire()
+        finally:
             telemetry.queued_requests.dec()
-            raise
+        queue_seconds = time.perf_counter() - started
+        # What this request actually waited becomes the next request's
+        # estimate for the same model.
+        load.observe_queue(decision.profile.id, queue_seconds * 1000)
+        telemetry.inflight_requests.inc()
+
+        if payload.stream:
+            # The stream holds its admission slot until it ends. Returning it
+            # from inside a slot block would free the slot before generation
+            # began, and streamed work would escape the concurrency bound.
+            lease = _StreamLease(span)
+            span.handed_off = True
+            return _LeasedStreamingResponse(
+                _stream_completion(
+                    payload,
+                    prompt,
+                    cache_key,
+                    subject,
+                    decision,
+                    started,
+                    queue_seconds,
+                    span,
+                    lease,
+                ),
+                lease=lease,
+                media_type="text/event-stream",
+                headers=_route_headers(decision, "miss"),
+            )
+
+        try:
+            result = await inference_backend.generate(payload, decision)
+        finally:
+            telemetry.inflight_requests.dec()
+            admission.release()
 
         span.set_usage(
             prompt_tokens=result.prompt_tokens, completion_tokens=result.completion_tokens

src/llm_router/backends.py

@@ -19,6 +19,28 @@ class BackendUnavailableError(RuntimeError):
     """Raised when an inference engine is unreachable or returns an error."""
 
 
+class BackendOutOfMemoryError(BackendUnavailableError):
+    """Raised when the engine reports it ran out of GPU memory for a request."""
+
+
+OUT_OF_MEMORY_MARKERS = ("out of memory", "outofmemoryerror", "cuda oom")
+
+
+def engine_failure(status_code: int, body: str, context: str) -> BackendUnavailableError:
+    """Classify an engine error response without repeating what it said.
+
+    The body is inspected for an out-of-memory report and then discarded: an
+    engine error can echo the prompt it rejected, so it is never forwarded.
+    """
+
+    if any(marker in body.lower() for marker in OUT_OF_MEMORY_MARKERS):
+        return BackendOutOfMemoryError(
+            f"inference engine ran out of GPU memory {context}; "
+            "retry with a shorter prompt or a smaller max_tokens"
+        )
+    return BackendUnavailableError(f"inference engine returned {status_code} {context}")
+
+
 @dataclass(frozen=True)
 class BackendResult:
     text: str
@@ -103,9 +125,8 @@ async def generate(
             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)}"
+            raise engine_failure(
+                response.status_code, response.text, f"for {served_model_name(decision)}"
             )
 
         body = response.json()
@@ -133,9 +154,8 @@ async def stream(
                 timeout=self.request_timeout_seconds,
             ) as response:
                 if response.status_code >= 400:
-                    raise BackendUnavailableError(
-                        f"inference engine returned {response.status_code} while streaming"
-                    )
+                    body = (await response.aread()).decode(errors="replace")
+                    raise engine_failure(response.status_code, body, "while streaming")
                 async for line in response.aiter_lines():
                     delta = _parse_stream_line(line)
                     if delta:

src/llm_router/config.py

@@ -24,6 +24,13 @@ class Settings(BaseSettings):
     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)
+    # Consecutive engine failures before requests fail fast, and how long
+    # the gateway waits before letting a trial request through again.
+    engine_failure_threshold: int = Field(default=5, ge=1)
+    engine_cooldown_seconds: float = Field(default=30.0, gt=0)
+    # How long shutdown waits for admitted requests before closing the engine
+    # client; keep it under the pod's termination grace period.
+    shutdown_grace_seconds: float = Field(default=20.0, ge=0)
     registry_path: str = "config/registry.yaml"
     routing_policy_version: str = "v1"
     # Labelled prompts the task and complexity classifier is trained on at