bing
80eef4ce6a
Three always-running vLLM services (text TP=2 GPU0+1, ocr + embed on GPU2, sleep mode) behind a FastAPI router that auto-wakes models on request. Tiered idle (sleep 15 min / offload 3 h), depth-aware 503s with Retry-After, persisted wake-intent recovery, admin API on 127.0.0.1:8010. Routine control via vllmctl is pure HTTP — no docker on the request path. Verified: 91 router unit tests + 15-test E2E on real hardware (measurements in CALIBRATION.md; design record in .claude/memory/router-front-door-plan.md). Old nginx stack files removed before git init; design survives in .claude/memory/sleep-mode-implementation-plan.md. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
537 lines
21 KiB
Python
537 lines
21 KiB
Python
"""Public listener: path allowlist, model resolution, streaming proxy.
|
|
|
|
Security shape (plan section 6):
|
|
* the *normalized* path decides what is served, so `/v1/../sleep` and
|
|
`/v1%2f..%2fsleep` both collapse to `/sleep` and get a 404;
|
|
* `/admin/*` and every vLLM dev endpoint are simply not in the allowlist;
|
|
* the admin listener lives on another socket (see admin_api.py).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import posixpath
|
|
import re
|
|
import time
|
|
import urllib.parse
|
|
from typing import Any
|
|
|
|
import httpx
|
|
from fastapi import FastAPI, Request
|
|
from fastapi.responses import JSONResponse, Response, StreamingResponse
|
|
from starlette.background import BackgroundTask
|
|
from starlette.datastructures import MutableHeaders
|
|
from starlette.middleware.body_limit import RequestBodyLimitMiddleware
|
|
|
|
from config import (
|
|
DEPTH_AWAKE,
|
|
DEPTH_OFFLOADED,
|
|
DEPTH_RESTARTING,
|
|
DEPTH_SLEEPING,
|
|
Config,
|
|
)
|
|
from services import ServiceManager, ServiceState
|
|
|
|
log = logging.getLogger("vllm_router.http")
|
|
|
|
# Endpoints whose bodies legitimately omit `model` (plan 6.2 / 6.6).
|
|
DEFAULT_TEXT_PATHS = frozenset({"/v1/chat/completions", "/v1/completions"})
|
|
# A multipart request on the chat path is the image-bearing (OCR) path.
|
|
MULTIPART_DEFAULT_PATH = "/v1/chat/completions"
|
|
|
|
# Headers that must never be forwarded in either direction.
|
|
HOP_BY_HOP = frozenset({
|
|
"connection", "keep-alive", "proxy-authenticate", "proxy-authorization",
|
|
"te", "trailers", "transfer-encoding", "upgrade",
|
|
})
|
|
REQUEST_ONLY_STRIP = HOP_BY_HOP | {"host", "content-length", "expect"}
|
|
RESPONSE_STRIP = HOP_BY_HOP | {"content-length"}
|
|
|
|
_BOUNDARY_RE = re.compile(r'boundary="?([^";,]+)"?', re.IGNORECASE)
|
|
_NAME_RE = re.compile(r'name="((?:[^"\\]|\\.)*)"', re.IGNORECASE)
|
|
# A text field we are willing to read fully (and re-send) while sniffing.
|
|
_MAX_MODEL_FIELD_BYTES = 4096
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# errors (OpenAI-shaped so existing clients can parse them)
|
|
# --------------------------------------------------------------------------
|
|
def _error(status: int, message: str, *, err_type: str = "invalid_request_error",
|
|
code: str | None = None, param: str | None = None,
|
|
extra: dict[str, Any] | None = None, headers: dict[str, str] | None = None
|
|
) -> JSONResponse:
|
|
err: dict[str, Any] = {"message": message, "type": err_type}
|
|
if param is not None:
|
|
err["param"] = param
|
|
if code is not None:
|
|
err["code"] = code
|
|
if extra:
|
|
err.update(extra)
|
|
return JSONResponse({"error": err}, status_code=status, headers=headers)
|
|
|
|
|
|
def not_found_response(message: str = "Not found") -> JSONResponse:
|
|
"""Uniform 404 for everything outside the allowlist (paths *and* models)."""
|
|
return _error(404, message, code="not_found")
|
|
|
|
|
|
def shutting_down_response(detail: str = "router is shutting down") -> JSONResponse:
|
|
"""503 for in-flight work that dies because the router is stopping."""
|
|
return JSONResponse(
|
|
{"error": {
|
|
"type": "router_shutting_down",
|
|
"code": "router_shutting_down",
|
|
"message": f"Request aborted: {detail}; retry shortly",
|
|
"retry_after_seconds": 5,
|
|
}},
|
|
status_code=503,
|
|
headers={"Retry-After": "5", "connection": "close"},
|
|
)
|
|
|
|
|
|
_SHUTDOWN_BODY = json.dumps({
|
|
"error": {
|
|
"type": "router_shutting_down",
|
|
"code": "router_shutting_down",
|
|
"message": "Request aborted: the router is shutting down; retry shortly",
|
|
"retry_after_seconds": 5,
|
|
}
|
|
}).encode()
|
|
|
|
|
|
class ShutdownGuardMiddleware:
|
|
"""Turns shutdown-induced cancellation of an in-flight request into a
|
|
depth-aware 503 instead of uvicorn's bare 21-byte 500 (E2E case 12).
|
|
|
|
uvicorn cancels pending ASGI tasks once `timeout_graceful_shutdown`
|
|
expires; that CancelledError is a BaseException, so Starlette's
|
|
ServerErrorMiddleware (and every FastAPI exception_handler) ignores it and
|
|
uvicorn answers with a plain-text 500. This is the only layer that can
|
|
intercept it. A response whose body has already started streaming cannot
|
|
be un-sent: those connections are simply closed (client sees a truncated
|
|
stream), which is the best available behaviour.
|
|
"""
|
|
|
|
def __init__(self, app, manager: ServiceManager) -> None:
|
|
self.app = app
|
|
self.manager = manager
|
|
|
|
async def __call__(self, scope, receive, send) -> None:
|
|
if scope["type"] != "http":
|
|
await self.app(scope, receive, send)
|
|
return
|
|
started = False
|
|
|
|
async def send_wrapper(message) -> None:
|
|
nonlocal started
|
|
if message["type"] == "http.response.start":
|
|
started = True
|
|
await send(message)
|
|
|
|
try:
|
|
await self.app(scope, receive, send_wrapper)
|
|
except BaseException as exc: # noqa: BLE001 - CancelledError is one
|
|
if started or not self.manager.shutting_down:
|
|
raise
|
|
log.info("abort_shutdown method=%s path=%s error=%s",
|
|
scope.get("method"), scope.get("path"), type(exc).__name__)
|
|
try:
|
|
await send({"type": "http.response.start", "status": 503, "headers": [
|
|
(b"content-type", b"application/json"),
|
|
(b"content-length", str(len(_SHUTDOWN_BODY)).encode("latin-1")),
|
|
(b"retry-after", b"5"),
|
|
(b"connection", b"close"),
|
|
]})
|
|
await send({"type": "http.response.body", "body": _SHUTDOWN_BODY})
|
|
except Exception: # pragma: no cover - client may be gone already
|
|
pass
|
|
|
|
|
|
def model_not_found_response(model: str) -> JSONResponse:
|
|
return _error(
|
|
404,
|
|
f"The model `{model}` does not exist or is not served by this endpoint.",
|
|
param="model",
|
|
code="model_not_found",
|
|
)
|
|
|
|
|
|
def waking_response(cfg: Config, svc: ServiceState, depth: str,
|
|
detail: str = "") -> JSONResponse:
|
|
"""Depth-aware 503 (plan 6.2.1)."""
|
|
policy = cfg.policy(depth)
|
|
body = policy.error_body(svc.cfg.model)
|
|
if detail:
|
|
body["error"]["message"] += f" ({detail})"
|
|
return JSONResponse(
|
|
body,
|
|
status_code=503,
|
|
headers={
|
|
"Retry-After": str(policy.retry_after_s),
|
|
"x-vllm-router": "model-waking",
|
|
},
|
|
)
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# path handling
|
|
# --------------------------------------------------------------------------
|
|
def normalize_raw_path(raw: bytes | str) -> str:
|
|
"""Decode %xx and collapse `.`/`..` segments BEFORE any allowlist check."""
|
|
if isinstance(raw, (bytes, bytearray)):
|
|
raw = bytes(raw).decode("latin-1")
|
|
path = urllib.parse.unquote(raw)
|
|
path = posixpath.normpath(path)
|
|
if not path.startswith("/"):
|
|
path = "/" + path
|
|
return path
|
|
|
|
|
|
def classify_public_path(norm: str, cfg: Config) -> str | None:
|
|
"""Returns the handler kind, or None for 404."""
|
|
if norm == "/health":
|
|
return "health"
|
|
if norm == "/metrics" and cfg.metrics_enabled:
|
|
return "metrics"
|
|
if norm == "/v1/models":
|
|
return "models"
|
|
if norm.startswith("/v1/") and ".." not in norm.split("/") and "\\" not in norm:
|
|
return "api"
|
|
return None
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# model resolution (plan 6.2 step 1)
|
|
# --------------------------------------------------------------------------
|
|
def extract_multipart_model(body: bytes, content_type: str) -> str | None:
|
|
"""Pull a small `model` text field out of a multipart body.
|
|
|
|
Deliberately hand-rolled: Starlette's form parser would consume (and
|
|
buffer/spool) every file part. We walk the boundary delimiters, look at
|
|
part headers only, and never slice a file part's payload.
|
|
"""
|
|
match = _BOUNDARY_RE.search(content_type)
|
|
if match is None:
|
|
return None
|
|
delimiter = b"--" + match.group(1).encode("latin-1")
|
|
pos = 0
|
|
while True:
|
|
start = body.find(delimiter, pos)
|
|
if start < 0:
|
|
return None
|
|
line_end = body.find(b"\r\n", start + len(delimiter))
|
|
if line_end < 0:
|
|
return None
|
|
next_delim = body.find(delimiter, line_end)
|
|
if next_delim < 0:
|
|
return None
|
|
headers = body[line_end + 2:next_delim]
|
|
head, sep, _ = headers.partition(b"\r\n\r\n")
|
|
if not sep:
|
|
pos = next_delim
|
|
continue
|
|
disposition = b""
|
|
for line in head.split(b"\r\n"):
|
|
if line.lower().startswith(b"content-disposition:"):
|
|
disposition = line
|
|
break
|
|
if b"filename=" in disposition.lower():
|
|
pos = next_delim # file part: leave it untouched
|
|
continue
|
|
name_match = _NAME_RE.search(disposition.decode("latin-1", "replace"))
|
|
if name_match is None or name_match.group(1).lower() != "model":
|
|
pos = next_delim
|
|
continue
|
|
value = headers[len(head) + 4:].rstrip(b"\r\n")
|
|
if len(value) > _MAX_MODEL_FIELD_BYTES:
|
|
return None
|
|
return value.decode("utf-8", "replace").strip() or None
|
|
|
|
|
|
def resolve_model_name(cfg: Config, name: str) -> str | None:
|
|
"""Registry lookup (case-insensitive over model names, aliases and keys)."""
|
|
return cfg.index.get(name.strip().lower())
|
|
|
|
|
|
async def resolve_target(manager: ServiceManager, request: Request, norm_path: str
|
|
) -> tuple[ServiceState | None, Response | None]:
|
|
"""Decide which service a request belongs to. Returns (service, error);
|
|
exactly one is None. Never wakes anything when the model is unknown."""
|
|
cfg = manager.cfg
|
|
|
|
# 1. /v1/embeddings always goes to the embed service.
|
|
if norm_path == "/v1/embeddings":
|
|
return manager.services["embed"], None
|
|
|
|
# 2. Model named in the path (GET /v1/models/{id}).
|
|
if norm_path.startswith("/v1/models/") and request.method in ("GET", "HEAD"):
|
|
key = resolve_model_name(cfg, urllib.parse.unquote(norm_path[len("/v1/models/"):]))
|
|
if key is None:
|
|
return None, model_not_found_response(norm_path[len("/v1/models/"):])
|
|
return manager.services[key], None
|
|
|
|
content_type = request.headers.get("content-type", "")
|
|
is_json = content_type.split(";")[0].strip().lower() in ("application/json",) or \
|
|
content_type.strip().lower().endswith("+json")
|
|
is_multipart = content_type.split(";")[0].strip().lower().startswith("multipart/")
|
|
|
|
if is_json:
|
|
raw = await request.body()
|
|
try:
|
|
payload = json.loads(raw or b"{}")
|
|
except (ValueError, UnicodeDecodeError):
|
|
return None, _error(400, "Request body is not valid JSON.", code="invalid_json")
|
|
if not isinstance(payload, dict):
|
|
return None, _error(400, "Request body must be a JSON object.", code="invalid_json")
|
|
model = payload.get("model")
|
|
if isinstance(model, str) and model.strip():
|
|
key = resolve_model_name(cfg, model)
|
|
if key is None:
|
|
return None, model_not_found_response(model)
|
|
return manager.services[key], None
|
|
if norm_path in DEFAULT_TEXT_PATHS:
|
|
return manager.services["text"], None
|
|
return None, _error(
|
|
400,
|
|
"Missing required parameter: 'model'.",
|
|
param="model",
|
|
code="missing_model",
|
|
)
|
|
|
|
if is_multipart:
|
|
model = extract_multipart_model(await request.body(), content_type)
|
|
if model is not None:
|
|
key = resolve_model_name(cfg, model)
|
|
if key is None:
|
|
return None, model_not_found_response(model)
|
|
return manager.services[key], None
|
|
if norm_path == MULTIPART_DEFAULT_PATH:
|
|
return manager.services["ocr"], None # image-bearing path
|
|
if norm_path in DEFAULT_TEXT_PATHS:
|
|
return manager.services["text"], None
|
|
return None, _error(
|
|
400,
|
|
"Missing required parameter: 'model'.",
|
|
param="model",
|
|
code="missing_model",
|
|
)
|
|
|
|
# Any other body shape (empty, text/plain, ...): only the two chat paths
|
|
# have a defensible default.
|
|
if norm_path in DEFAULT_TEXT_PATHS:
|
|
return manager.services["text"], None
|
|
return None, _error(
|
|
400,
|
|
"Missing required parameter: 'model'.",
|
|
param="model",
|
|
code="missing_model",
|
|
)
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# proxy
|
|
# --------------------------------------------------------------------------
|
|
def _forwarded_request_headers(request: Request) -> list[tuple[str, str]]:
|
|
return [
|
|
(k, v) for k, v in request.headers.items()
|
|
if k.lower() not in REQUEST_ONLY_STRIP
|
|
]
|
|
|
|
|
|
async def _open_upstream(manager: ServiceManager, request: Request, svc: ServiceState,
|
|
norm_path: str) -> httpx.Response:
|
|
query = request.scope.get("query_string", b"").decode("latin-1")
|
|
url = svc.cfg.base_url + norm_path + (f"?{query}" if query else "")
|
|
upstream_req = manager.proxy_client.build_request(
|
|
request.method,
|
|
url,
|
|
headers=_forwarded_request_headers(request),
|
|
content=await request.body(),
|
|
)
|
|
# Stream: the response body is handed to the client chunk by chunk and is
|
|
# never buffered here (SSE safe, 900s read timeout upstream).
|
|
return await manager.proxy_client.send(upstream_req, stream=True)
|
|
|
|
|
|
async def _proxy_response(manager: ServiceManager, request: Request, svc: ServiceState,
|
|
norm_path: str) -> Response:
|
|
upstream = await _open_upstream(manager, request, svc, norm_path)
|
|
headers = MutableHeaders()
|
|
for key, value in upstream.headers.items():
|
|
if key.lower() not in RESPONSE_STRIP:
|
|
headers.append(key, value)
|
|
|
|
async def finish() -> None:
|
|
# Runs when the last chunk has been sent (or the client hung up) --
|
|
# this is what keeps the idle manager from sleeping mid-stream.
|
|
try:
|
|
await upstream.aclose()
|
|
finally:
|
|
manager.end_request(svc)
|
|
|
|
return StreamingResponse(
|
|
upstream.aiter_raw(),
|
|
status_code=upstream.status_code,
|
|
headers=headers,
|
|
background=BackgroundTask(finish),
|
|
)
|
|
|
|
|
|
async def forward(manager: ServiceManager, request: Request, svc: ServiceState,
|
|
norm_path: str) -> Response:
|
|
"""Ensure the target is awake, then proxy. The request stays counted as
|
|
active for the whole life of the response body."""
|
|
cfg = manager.cfg
|
|
manager.begin_request(svc)
|
|
handed_off = False
|
|
try:
|
|
for attempt in (1, 2):
|
|
outcome = await manager.ensure_awake(svc.cfg.key)
|
|
if not outcome.ok:
|
|
log.info("reject service=%s status=503 depth=%s reason=%s",
|
|
svc.cfg.key, outcome.depth, outcome.reason)
|
|
detail = "" if outcome.reason == "hold_timeout" else outcome.detail
|
|
return waking_response(cfg, svc, outcome.depth, detail)
|
|
try:
|
|
response = await _proxy_response(manager, request, svc, norm_path)
|
|
handed_off = True
|
|
return response
|
|
except httpx.TransportError as exc:
|
|
# Backend went away between the health check and the proxy call.
|
|
svc.reachable = False
|
|
svc.depth = None
|
|
svc.state_checked_at = float("-inf")
|
|
svc.last_error = f"proxy: {type(exc).__name__}"
|
|
log.warning("proxy_transport_error service=%s attempt=%d error=%s",
|
|
svc.cfg.key, attempt, type(exc).__name__)
|
|
if attempt == 2:
|
|
return waking_response(cfg, svc, DEPTH_RESTARTING,
|
|
"backend connection failed")
|
|
finally:
|
|
if not handed_off:
|
|
manager.end_request(svc)
|
|
return waking_response(cfg, svc, DEPTH_RESTARTING, "unreachable") # pragma: no cover
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# public endpoints
|
|
# --------------------------------------------------------------------------
|
|
def _service_summary(manager: ServiceManager, live: bool) -> dict[str, Any]:
|
|
out: dict[str, Any] = {}
|
|
for svc in manager.services.values():
|
|
out[svc.cfg.key] = {
|
|
"service": svc.cfg.service,
|
|
"model": svc.cfg.model,
|
|
"reachable": svc.reachable,
|
|
"sleeping": None if svc.depth is None else svc.depth != DEPTH_AWAKE,
|
|
"depth": svc.depth,
|
|
"wake_in_progress": svc.wake_in_progress,
|
|
"active_requests": svc.active_requests,
|
|
}
|
|
return out
|
|
|
|
|
|
def build_public_app(manager: ServiceManager) -> FastAPI:
|
|
cfg = manager.cfg
|
|
app = FastAPI(
|
|
title="vllm-router",
|
|
version="1.0.0",
|
|
docs_url=None,
|
|
redoc_url=None,
|
|
openapi_url=None,
|
|
)
|
|
# FastAPI does not forward **extra to Starlette, so the request-body cap
|
|
# (ROUTER_MAX_BODY_BYTES) is wired as middleware explicitly.
|
|
app.add_middleware(RequestBodyLimitMiddleware, max_body_size=cfg.max_body_bytes)
|
|
# Added last => outermost. Must sit outside everything because the
|
|
# exception it converts is a BaseException that Starlette's
|
|
# ServerErrorMiddleware deliberately does not catch.
|
|
app.add_middleware(ShutdownGuardMiddleware, manager=manager)
|
|
|
|
@app.exception_handler(Exception)
|
|
async def internal_error(_request: Request, exc: Exception) -> JSONResponse:
|
|
"""Anything unhandled becomes a parseable JSON error, never a bare
|
|
21-byte "Internal Server Error". During shutdown that is a 503 with
|
|
Retry-After (E2E case 12)."""
|
|
if manager.shutting_down:
|
|
log.info("abort_shutdown method=%s path=%s error=%s",
|
|
_request.method, _request.url.path, type(exc).__name__)
|
|
return shutting_down_response(f"{type(exc).__name__} during shutdown")
|
|
log.exception("internal_error error=%s", type(exc).__name__)
|
|
return _error(500, f"Unhandled router error: {type(exc).__name__}",
|
|
err_type="internal_error", code="internal_error",
|
|
headers={"Retry-After": "1"})
|
|
|
|
@app.get("/health")
|
|
async def health() -> JSONResponse:
|
|
"""Router liveness + cached per-service summary. Always 200: a sleeping
|
|
backend is normal, not an outage (docker healthchecks must not flap)."""
|
|
return JSONResponse({
|
|
"status": "ok",
|
|
"router": {
|
|
"uptime_s": round(time.monotonic() - manager.started_at, 1),
|
|
"public_port": cfg.public_port,
|
|
"admin_port": cfg.admin_port,
|
|
"services": len(manager.services),
|
|
},
|
|
"services": _service_summary(manager, live=False),
|
|
})
|
|
|
|
@app.get("/v1/models")
|
|
async def models() -> JSONResponse:
|
|
return JSONResponse({
|
|
"object": "list",
|
|
"data": [
|
|
{"id": svc.cfg.model, "object": "model", "created": 0,
|
|
"owned_by": "vllm-router"}
|
|
for svc in manager.services.values()
|
|
],
|
|
})
|
|
|
|
# NOTE: /metrics is deliberately NOT registered as an explicit route -- the
|
|
# catch-all below re-checks it against the allowlist so it 404s whenever
|
|
# ROUTER_METRICS is off.
|
|
async def metrics() -> Response:
|
|
lines = [
|
|
"# HELP vllm_router_service_depth 0=awake 1=sleeping 2=offloaded 3=restarting",
|
|
"# TYPE vllm_router_service_depth gauge",
|
|
]
|
|
rank = {DEPTH_AWAKE: 0, DEPTH_SLEEPING: 1, DEPTH_OFFLOADED: 2, DEPTH_RESTARTING: 3}
|
|
for svc in manager.services.values():
|
|
lines.append(f'vllm_router_service_depth{{service="{svc.cfg.key}"}} '
|
|
f'{rank.get(svc.depth, 2)}')
|
|
lines.append(f'vllm_router_active_requests{{service="{svc.cfg.key}"}} '
|
|
f'{svc.active_requests}')
|
|
return Response("\n".join(lines) + "\n", media_type="text/plain; version=0.0.4")
|
|
|
|
async def gate(request: Request) -> Response:
|
|
raw = request.scope.get("raw_path") or request.url.path.encode()
|
|
norm = normalize_raw_path(raw)
|
|
kind = classify_public_path(norm, cfg)
|
|
if kind is None:
|
|
log.info("reject_path method=%s path=%s", request.method, norm)
|
|
return not_found_response()
|
|
if kind == "health":
|
|
return await health()
|
|
if kind == "models":
|
|
return await models()
|
|
if kind == "metrics":
|
|
return await metrics()
|
|
|
|
svc, error = await resolve_target(manager, request, norm)
|
|
if error is not None:
|
|
log.info("reject_model method=%s path=%s status=%s",
|
|
request.method, norm, error.status_code)
|
|
return error
|
|
return await forward(manager, request, svc, norm)
|
|
|
|
# Registered last: explicit routes above win, everything else lands here
|
|
# and is re-checked against the allowlist on the *normalized* path.
|
|
app.router.add_route(
|
|
"/{path:path}",
|
|
gate,
|
|
methods=["GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"],
|
|
)
|
|
return app
|