Initial commit: router front-door vLLM stack
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>
This commit is contained in:
28
router/Dockerfile
Normal file
28
router/Dockerfile
Normal file
@@ -0,0 +1,28 @@
|
||||
# vllm-router: single process, one event loop, two sockets (public :8000,
|
||||
# admin :8010). No GPU, no model weights -- pure HTTP front door.
|
||||
FROM python:3.12-slim
|
||||
|
||||
# Non-root runtime user.
|
||||
RUN groupadd --system --gid 10001 router \
|
||||
&& useradd --system --uid 10001 --gid router --home-dir /app router
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Deps first so code changes don't bust the layer.
|
||||
COPY requirements.txt ./
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
COPY app.py config.py services.py routing.py admin_api.py ./
|
||||
|
||||
# Wake-intent state dir: a named volume seeded from this ownership, so the
|
||||
# non-root runtime user can always write it (bind mounts inherit host uids).
|
||||
RUN mkdir -p /state && chown router:router /state
|
||||
|
||||
USER router
|
||||
EXPOSE 8000 8010
|
||||
|
||||
# Liveness only (a sleeping backend is normal, not an outage).
|
||||
HEALTHCHECK --interval=30s --timeout=5s --start-period=15s --retries=3 \
|
||||
CMD ["python", "-c", "import urllib.request;urllib.request.urlopen('http://127.0.0.1:8000/health', timeout=3)"]
|
||||
|
||||
CMD ["python", "app.py"]
|
||||
143
router/admin_api.py
Normal file
143
router/admin_api.py
Normal file
@@ -0,0 +1,143 @@
|
||||
"""Admin listener (the :8010 socket) -- `vllmctl` / debugging surface.
|
||||
|
||||
This app is mounted on a *separate socket*, never on the public one (plan
|
||||
review C5: a public /admin/sleep would be a trivial remote DoS). It shares
|
||||
the exact same ServiceManager, locks and depth tracking as the request path.
|
||||
|
||||
The socket binds 0.0.0.0 *inside the container*; the host-side "localhost
|
||||
only" restriction comes from compose publishing `127.0.0.1:8010:8010`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
from config import DEPTH_AWAKE
|
||||
from services import ServiceManager
|
||||
|
||||
log = logging.getLogger("vllm_router.admin")
|
||||
|
||||
|
||||
def build_admin_app(manager: ServiceManager) -> FastAPI:
|
||||
cfg = manager.cfg
|
||||
app = FastAPI(title="vllm-router admin", version="1.0.0",
|
||||
docs_url=None, redoc_url=None, openapi_url=None)
|
||||
|
||||
def service_or_404(key: str) -> tuple[Any, JSONResponse | None]:
|
||||
# Accepts the service key ("text") or the model name / alias
|
||||
# ("OvisOCR2", "ocr", "qwen3.6-35b-a3b-fp8"), case-insensitively.
|
||||
resolved = key.strip().lower()
|
||||
if resolved not in manager.services:
|
||||
resolved = cfg.index.get(key.strip().lower(), "")
|
||||
svc = manager.services.get(resolved)
|
||||
if svc is None:
|
||||
body = {"error": {"message": f"Unknown service or model '{key}'.",
|
||||
"type": "invalid_request_error", "code": "not_found",
|
||||
"known": sorted(manager.services)}}
|
||||
return None, JSONResponse(body, status_code=404)
|
||||
return svc, None
|
||||
|
||||
@app.get("/admin/status")
|
||||
async def admin_status() -> JSONResponse:
|
||||
status = await manager.status(live_probe=True)
|
||||
return JSONResponse({
|
||||
"router": {
|
||||
"uptime_s": status["uptime_s"],
|
||||
"public_port": cfg.public_port,
|
||||
"admin_port": cfg.admin_port,
|
||||
"idle": {
|
||||
"enabled": cfg.idle_enabled,
|
||||
"sleep_min": cfg.idle_sleep_min,
|
||||
"offload_min": cfg.idle_offload_min,
|
||||
},
|
||||
},
|
||||
"services": status["services"],
|
||||
"api": "http://127.0.0.1:%d/v1" % cfg.public_port,
|
||||
})
|
||||
|
||||
@app.post("/admin/wake/{key}")
|
||||
async def admin_wake(key: str) -> JSONResponse:
|
||||
svc, err = service_or_404(key)
|
||||
if err is not None:
|
||||
return err
|
||||
outcome = await manager.ensure_awake(svc.cfg.key)
|
||||
svc.last_activity = time.monotonic() # idle clock restarts
|
||||
body = {
|
||||
"service": svc.cfg.key,
|
||||
"model": svc.cfg.model,
|
||||
"ok": outcome.ok,
|
||||
"depth": outcome.depth,
|
||||
"reason": outcome.reason,
|
||||
"wake_in_progress": svc.wake_in_progress,
|
||||
}
|
||||
if outcome.latency_s is not None:
|
||||
body["latency_s"] = round(outcome.latency_s, 2)
|
||||
if not outcome.ok:
|
||||
policy = cfg.policy(outcome.depth)
|
||||
body["retry_after_s"] = policy.retry_after_s
|
||||
body["estimated_wake_seconds"] = policy.est_wake_s
|
||||
body["error"] = {"type": "model_waking", "code": "model_waking",
|
||||
"sleep_depth": outcome.depth,
|
||||
"estimated_wake_seconds": policy.est_wake_s}
|
||||
if outcome.detail:
|
||||
body["error"]["message"] = outcome.detail
|
||||
log.info("admin_wake service=%s ok=false depth=%s reason=%s",
|
||||
svc.cfg.key, outcome.depth, outcome.reason)
|
||||
return JSONResponse(body, status_code=503,
|
||||
headers={"Retry-After": str(policy.retry_after_s)})
|
||||
log.info("admin_wake service=%s ok=true depth=%s reason=%s",
|
||||
svc.cfg.key, outcome.depth, outcome.reason)
|
||||
return JSONResponse(body)
|
||||
|
||||
@app.post("/admin/sleep/{key}")
|
||||
async def admin_sleep(key: str, request: Request) -> JSONResponse:
|
||||
svc, err = service_or_404(key)
|
||||
if err is not None:
|
||||
return err
|
||||
level = 1
|
||||
if request.query_params.get("level"):
|
||||
try:
|
||||
level = int(request.query_params["level"])
|
||||
except ValueError:
|
||||
return JSONResponse(
|
||||
{"error": {"message": "level must be 1 or 2",
|
||||
"type": "invalid_request_error"}}, status_code=400)
|
||||
if level not in (1, 2):
|
||||
return JSONResponse(
|
||||
{"error": {"message": "level must be 1 or 2",
|
||||
"type": "invalid_request_error"}}, status_code=400)
|
||||
result = await manager.sleep_service(svc.cfg.key, level, reason="admin")
|
||||
status_code = 200 if result.get("ok") else 409
|
||||
result["level_requested"] = level
|
||||
result["model"] = svc.cfg.model
|
||||
log.info("admin_sleep service=%s level=%d ok=%s reason=%s",
|
||||
svc.cfg.key, level, result.get("ok"), result.get("reason"))
|
||||
return JSONResponse(result, status_code=status_code)
|
||||
|
||||
@app.get("/health")
|
||||
@app.get("/admin/health", include_in_schema=False)
|
||||
async def admin_health() -> JSONResponse:
|
||||
return JSONResponse({"status": "ok",
|
||||
"awake": [k for k, s in manager.services.items()
|
||||
if s.depth == DEPTH_AWAKE],
|
||||
"wake_recovery_pending": [k for k, s in manager.services.items()
|
||||
if s.pending_reload]})
|
||||
|
||||
@app.exception_handler(Exception)
|
||||
async def internal_error(_request: Request, exc: Exception) -> JSONResponse:
|
||||
if manager.shutting_down:
|
||||
return JSONResponse(
|
||||
{"error": {"type": "router_shutting_down", "code": "router_shutting_down",
|
||||
"message": "Router is shutting down; retry shortly"}},
|
||||
status_code=503, headers={"Retry-After": "5"})
|
||||
return JSONResponse(
|
||||
{"error": {"type": "internal_error", "code": "internal_error",
|
||||
"message": f"Unhandled router error: {type(exc).__name__}"}},
|
||||
status_code=500)
|
||||
|
||||
return app
|
||||
187
router/app.py
Normal file
187
router/app.py
Normal file
@@ -0,0 +1,187 @@
|
||||
"""vllm-router -- single-process front door for the vLLM serving stack.
|
||||
|
||||
Plan v3.3 section 6. One process, ONE asyncio event loop, TWO listening
|
||||
sockets:
|
||||
|
||||
* public :8000 -> OpenAI API (/v1/*), /health, (/metrics)
|
||||
* admin :8010 -> /admin/* (for `vllmctl`; published by compose as
|
||||
`127.0.0.1:8010:8010`, so it is host-loopback only. Inside the
|
||||
container it binds 0.0.0.0 because docker cannot publish to a
|
||||
container-loopback bind -- see config.ROUTER_ADMIN_HOST.)
|
||||
|
||||
Two uvicorn.Server instances are run as coroutines on the same loop because
|
||||
uvicorn cannot bind two ports in one worker, and two *processes* would break
|
||||
the in-process wake locks -- the admin listener must share the very same
|
||||
ServiceManager, asyncio locks and depth tracking as the request path.
|
||||
|
||||
What it owns: model-name -> service mapping, wake-on-request (single-flight,
|
||||
depth-aware hold), streaming proxy, tiered idle sleep, dev-endpoint hiding
|
||||
(404 for anything outside the allowlist, matched on the *normalized* path).
|
||||
|
||||
Configuration (environment; authoritative list + defaults in config.py):
|
||||
|
||||
registry VLLM_TEXT_MODEL=Qwen3.6-35B-A3B-FP8 VLLM_OCR_MODEL=OvisOCR2
|
||||
VLLM_EMBED_MODEL=Qwen3-Embedding-8B (+ *_SERVICE, *_ALIASES)
|
||||
listeners ROUTER_PUBLIC_PORT=8000 ROUTER_ADMIN_PORT=8010
|
||||
ROUTER_ADMIN_HOST=127.0.0.1
|
||||
proxy ROUTER_CONNECT_TIMEOUT=5 ROUTER_READ_TIMEOUT=900
|
||||
ROUTER_POOL_SIZE=32 ROUTER_MAX_BODY_BYTES=134217728
|
||||
ROUTER_STATE_CACHE_TTL=2 (keep short: /health returns 200 on a
|
||||
sleeping backend, so readiness is gated on /is_sleeping)
|
||||
wake WAKE_HOLD_SLEEP_S=30 WAKE_HOLD_OFFLOAD_S=180
|
||||
WAKE_HOLD_RESTART_S=600 (text cold boot measured ~10 min),
|
||||
WAKE_RETRY_AFTER_{SLEEP,OFFLOAD,RESTART}=10/60/600,
|
||||
WAKE_ATTEMPTS=2, WAKE_HTTP_TIMEOUT=300
|
||||
idle IDLE_SLEEP_MIN=15 (level 1) IDLE_OFFLOAD_MIN=180 (level 2)
|
||||
IDLE_POLL_SECONDS=30 IDLE_ENABLED=1
|
||||
|
||||
Run: python app.py (see Dockerfile; `pytest` from router/tests)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
import signal
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# Sibling modules (config/services/routing/admin_api) -- makes the app work
|
||||
# both as `python app.py` and when imported by the test suite.
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
||||
|
||||
import uvicorn # noqa: E402
|
||||
|
||||
from admin_api import build_admin_app # noqa: E402
|
||||
from config import Config, load_config # noqa: E402
|
||||
from routing import build_public_app # noqa: E402
|
||||
from services import ServiceManager # noqa: E402
|
||||
|
||||
log = logging.getLogger("vllm_router")
|
||||
|
||||
ROUTES = ("GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS")
|
||||
|
||||
|
||||
def setup_logging(level: str = "info") -> None:
|
||||
logging.basicConfig(
|
||||
level=getattr(logging, level.upper(), logging.INFO),
|
||||
format="%(asctime)s %(levelname)-7s %(name)s %(message)s",
|
||||
datefmt="%Y-%m-%dT%H:%M:%S",
|
||||
)
|
||||
for noisy in ("httpx", "httpcore", "uvicorn.access"):
|
||||
logging.getLogger(noisy).setLevel(logging.WARNING)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _no_signal_capture() -> "contextlib.AbstractContextManager[None]":
|
||||
"""Replaces uvicorn.Server.capture_signals (see serve())."""
|
||||
yield
|
||||
|
||||
|
||||
def build_apps(cfg: Config | None = None, transport=None):
|
||||
"""Build the public + admin apps around one shared ServiceManager.
|
||||
|
||||
`transport` is an httpx transport override used by the test suite to stub
|
||||
the vLLM backends without any network, GPU or docker.
|
||||
"""
|
||||
cfg = cfg or load_config()
|
||||
manager = ServiceManager(cfg, transport=transport)
|
||||
public_app = build_public_app(manager)
|
||||
admin_app = build_admin_app(manager)
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def lifespan(_app):
|
||||
manager.start() # idempotent: called by both listeners
|
||||
yield
|
||||
await manager.stop() # idempotent as well
|
||||
|
||||
# Both apps share one manager; the lifespan is attached post-construction
|
||||
# so there is exactly one definition of it.
|
||||
public_app.router.lifespan_context = lifespan
|
||||
admin_app.router.lifespan_context = lifespan
|
||||
return public_app, admin_app, manager
|
||||
|
||||
|
||||
def _make_server(app, host: str, port: int, cfg: Config) -> uvicorn.Server:
|
||||
return uvicorn.Server(uvicorn.Config(
|
||||
app,
|
||||
host=host,
|
||||
port=port,
|
||||
log_level=cfg.log_level,
|
||||
access_log=cfg.access_log,
|
||||
# Long generations must not be cut off, but shutdown must terminate.
|
||||
timeout_graceful_shutdown=10,
|
||||
# The app does its own (structured) request logging.
|
||||
timeout_keep_alive=75,
|
||||
))
|
||||
|
||||
|
||||
async def serve(cfg: Config | None = None) -> None:
|
||||
cfg = cfg or load_config()
|
||||
setup_logging(cfg.log_level)
|
||||
public_app, admin_app, manager = build_apps(cfg)
|
||||
|
||||
servers = [
|
||||
_make_server(public_app, cfg.public_host, cfg.public_port, cfg),
|
||||
_make_server(admin_app, cfg.admin_host, cfg.admin_port, cfg),
|
||||
]
|
||||
|
||||
# uvicorn's signal capture is per-Server; with two servers the inner
|
||||
# capture would restore its own handler and swallow SIGTERM for the other
|
||||
# listener. Disable it and own the signals here (falling back gracefully
|
||||
# on uvicorn versions without that method).
|
||||
for server in servers:
|
||||
if hasattr(server, "capture_signals"):
|
||||
server.capture_signals = _no_signal_capture
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
stopping = asyncio.Event()
|
||||
|
||||
def _request_shutdown() -> None:
|
||||
if stopping.is_set():
|
||||
return
|
||||
stopping.set()
|
||||
log.info("router_stop signal_received")
|
||||
# Flag first, so requests dying in the shutdown window answer with a
|
||||
# depth-aware 503 instead of a bare 500 (E2E case 12).
|
||||
manager.begin_shutdown()
|
||||
for server in servers:
|
||||
server.should_exit = True
|
||||
|
||||
for sig in (signal.SIGINT, signal.SIGTERM):
|
||||
with contextlib.suppress(NotImplementedError, ValueError, RuntimeError):
|
||||
loop.add_signal_handler(sig, _request_shutdown)
|
||||
|
||||
async def _propagate_exit() -> None:
|
||||
"""Belt and braces: if only one server saw the shutdown, tell the other."""
|
||||
while not stopping.is_set():
|
||||
if any(s.should_exit for s in servers):
|
||||
for s in servers:
|
||||
s.should_exit = True
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
watchdog = asyncio.create_task(_propagate_exit(), name="exit-watchdog")
|
||||
log.info("router_start public=%s:%d admin=%s:%d services=%s",
|
||||
cfg.public_host, cfg.public_port, cfg.admin_host, cfg.admin_port,
|
||||
",".join(f"{s.key}={s.service}/{s.model}" for s in cfg.services.values()))
|
||||
try:
|
||||
await asyncio.gather(*(server.serve() for server in servers))
|
||||
finally:
|
||||
watchdog.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await watchdog
|
||||
await manager.stop()
|
||||
log.info("router_stopped")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
try:
|
||||
asyncio.run(serve())
|
||||
except KeyboardInterrupt: # pragma: no cover - interactive runs
|
||||
return 0
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
293
router/config.py
Normal file
293
router/config.py
Normal file
@@ -0,0 +1,293 @@
|
||||
"""Router configuration: model registry + tunables.
|
||||
|
||||
Everything is environment-driven with the defaults from plan v3.3 (see
|
||||
`.claude/memory/router-front-door-plan.md` sections 4/6). The `Config`
|
||||
object is a plain mutable dataclass so tests (and calibration) can override
|
||||
a single knob without re-reading the environment.
|
||||
|
||||
Environment variables (all optional):
|
||||
|
||||
Registry
|
||||
VLLM_TEXT_SERVICE default vllm-text docker service name (== DNS name)
|
||||
VLLM_TEXT_MODEL default Qwen3.6-35B-A3B-FP8
|
||||
VLLM_TEXT_ALIASES comma list, default "qwen,qwen3,qwen3.6,qwen3.6-35b,
|
||||
qwen3.6-35b-a3b,text,chat,default"
|
||||
VLLM_OCR_SERVICE default vllm-ocr
|
||||
VLLM_OCR_MODEL default OvisOCR2
|
||||
VLLM_OCR_ALIASES default "ocr,ovis,ovisocr,ovis-ocr"
|
||||
VLLM_EMBED_SERVICE default vllm-embed
|
||||
VLLM_EMBED_MODEL default Qwen3-Embedding-8B
|
||||
VLLM_EMBED_ALIASES default "embed,embedding,embeddings,qwen3-embedding,
|
||||
qwen3-embedding-8b"
|
||||
VLLM_BACKEND_PORT default 8000 (container port each vllm service listens on)
|
||||
|
||||
Listeners
|
||||
ROUTER_PUBLIC_HOST default 0.0.0.0
|
||||
ROUTER_PUBLIC_PORT default 8000 (the only public ingress)
|
||||
ROUTER_ADMIN_HOST default 0.0.0.0 (inside the container; the host-side
|
||||
127.0.0.1 restriction comes from compose's
|
||||
"127.0.0.1:8010:8010" port mapping -- docker cannot
|
||||
publish to a container-loopback bind. Set 127.0.0.1
|
||||
when running the router on bare metal.)
|
||||
ROUTER_ADMIN_PORT default 8010 (vllmctl / debugging)
|
||||
|
||||
Proxy (plan 6.2 "Proxy timeouts")
|
||||
ROUTER_CONNECT_TIMEOUT default 5 seconds
|
||||
ROUTER_READ_TIMEOUT default 900 seconds (capped, never disabled)
|
||||
ROUTER_POOL_SIZE default 32 (>= 16 concurrent requests)
|
||||
ROUTER_PROBE_TIMEOUT default 3 seconds for /is_sleeping + /health
|
||||
ROUTER_MAX_BODY_BYTES default 134217728 (128 MiB request body cap)
|
||||
|
||||
Wake-state cache
|
||||
ROUTER_STATE_CACHE_TTL default 2.0 seconds (must stay short: /health
|
||||
returns 200 on a sleeping backend and a proxied
|
||||
request to one HANGS, so a stale "awake" must not
|
||||
bypass the is_sleeping gate for long)
|
||||
|
||||
Depth-aware hold/503 policy (plan 6.2.1; restart deadline from calibration)
|
||||
WAKE_HOLD_SLEEP_S default 30 hold deadline, wake from level 1
|
||||
(measured wake 2.5-3.8s)
|
||||
WAKE_HOLD_OFFLOAD_S default 180 hold deadline, wake from level 2
|
||||
WAKE_HOLD_RESTART_S default 600 hold deadline, container restarting
|
||||
(measured text cold boot up to
|
||||
~10 min from NFS)
|
||||
WAKE_RETRY_AFTER_SLEEP default 10 Retry-After header value
|
||||
WAKE_RETRY_AFTER_OFFLOAD default 60
|
||||
WAKE_RETRY_AFTER_RESTART default 600
|
||||
WAKE_EST_SLEEP_S default 6 body estimated_wake_seconds
|
||||
WAKE_EST_OFFLOAD_S default 60
|
||||
WAKE_EST_RESTART_S default 600
|
||||
WAKE_ATTEMPTS default 2 (sequence retried once, plan 6.2.1)
|
||||
WAKE_HTTP_TIMEOUT default 300 per-call timeout inside the sequence
|
||||
WAKE_HEALTH_POLL_S default 2.0 /health poll interval while waking
|
||||
|
||||
Tiered idle (plan 6.3)
|
||||
IDLE_SLEEP_MIN default 15 minutes -> POST /sleep?level=1
|
||||
IDLE_OFFLOAD_MIN default 180 minutes -> POST /sleep?level=2
|
||||
IDLE_POLL_SECONDS default 30 idle-manager scan interval
|
||||
IDLE_ENABLED default 1 (0 disables the background tiering)
|
||||
|
||||
Misc
|
||||
ROUTER_LOG_LEVEL default info
|
||||
ROUTER_METRICS default 0 (1 = also expose /metrics publicly)
|
||||
ROUTER_ACCESS_LOG default 0 (uvicorn access log; off by default,
|
||||
the app logs one structured line per request)
|
||||
ROUTER_STATE_FILE default /state/router-state.json -- persisted wake
|
||||
intent + depth, so a restarted router can tell
|
||||
"awake and ready" from "awake because the previous
|
||||
router died between wake_up and reload_weights".
|
||||
Compose mounts the named volume `router-state`
|
||||
at /state (seeded with image ownership so the
|
||||
non-root user can write it).
|
||||
Set empty to disable persistence.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field, replace
|
||||
|
||||
|
||||
def _str(name: str, default: str) -> str:
|
||||
v = os.environ.get(name)
|
||||
return default if v is None or v == "" else v
|
||||
|
||||
|
||||
def _int(name: str, default: int) -> int:
|
||||
try:
|
||||
return int(_str(name, str(default)))
|
||||
except ValueError:
|
||||
return default
|
||||
|
||||
|
||||
def _float(name: str, default: float) -> float:
|
||||
try:
|
||||
return float(_str(name, str(default)))
|
||||
except ValueError:
|
||||
return default
|
||||
|
||||
|
||||
def _bool(name: str, default: bool) -> bool:
|
||||
return _str(name, "1" if default else "0").strip().lower() in ("1", "true", "yes", "on")
|
||||
|
||||
|
||||
def _csv(name: str, default: str) -> tuple[str, ...]:
|
||||
return tuple(p.strip() for p in _str(name, default).split(",") if p.strip())
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ServiceConfig:
|
||||
"""One registry entry: key -> (docker service, served model name)."""
|
||||
|
||||
key: str
|
||||
service: str
|
||||
model: str
|
||||
aliases: tuple[str, ...] = ()
|
||||
|
||||
@property
|
||||
def base_url(self) -> str:
|
||||
return f"http://{self.service}:{BACKEND_PORT}"
|
||||
|
||||
|
||||
# Read once at import; the base_url property depends on it.
|
||||
BACKEND_PORT = _int("VLLM_BACKEND_PORT", 8000)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DepthPolicy:
|
||||
"""Hold/retry parameters for one sleep depth (plan 6.2.1 table)."""
|
||||
|
||||
depth: str
|
||||
hold_s: float
|
||||
retry_after_s: int
|
||||
est_wake_s: int
|
||||
phrase: str # "... is waking from <phrase>; retry shortly"
|
||||
|
||||
def error_body(self, model: str) -> dict:
|
||||
return {
|
||||
"error": {
|
||||
"type": "model_waking",
|
||||
"code": "model_waking",
|
||||
"message": f"Model '{model}' is waking from {self.phrase}; retry shortly",
|
||||
"sleep_depth": self.depth,
|
||||
"estimated_wake_seconds": self.est_wake_s,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class Config:
|
||||
services: dict[str, ServiceConfig]
|
||||
# name (lower-cased) -> key, for model resolution
|
||||
index: dict[str, str] = field(default_factory=dict)
|
||||
|
||||
public_host: str = "0.0.0.0"
|
||||
public_port: int = 8000
|
||||
# NOTE: 0.0.0.0 is required *inside the container* -- docker port
|
||||
# publishing cannot reach a container-loopback bind. The "localhost only"
|
||||
# guarantee of plan section 4 comes from compose mapping the port as
|
||||
# `127.0.0.1:8010:8010` (host side). On bare metal set
|
||||
# ROUTER_ADMIN_HOST=127.0.0.1.
|
||||
admin_host: str = "0.0.0.0"
|
||||
admin_port: int = 8010
|
||||
|
||||
connect_timeout: float = 5.0
|
||||
read_timeout: float = 900.0
|
||||
pool_size: int = 32
|
||||
probe_timeout: float = 3.0
|
||||
max_body_bytes: int = 128 * 1024 * 1024
|
||||
|
||||
state_cache_ttl: float = 2.0
|
||||
|
||||
hold_sleep_s: float = 30.0
|
||||
hold_offload_s: float = 180.0
|
||||
hold_restart_s: float = 600.0
|
||||
retry_after_sleep: int = 10
|
||||
retry_after_offload: int = 60
|
||||
retry_after_restart: int = 600
|
||||
est_sleep_s: int = 6
|
||||
est_offload_s: int = 60
|
||||
est_restart_s: int = 600
|
||||
wake_attempts: int = 2
|
||||
wake_http_timeout: float = 300.0
|
||||
wake_health_poll_s: float = 2.0
|
||||
|
||||
idle_sleep_min: float = 15.0
|
||||
idle_offload_min: float = 180.0
|
||||
idle_poll_s: float = 30.0
|
||||
idle_enabled: bool = True
|
||||
|
||||
metrics_enabled: bool = False
|
||||
log_level: str = "info"
|
||||
access_log: bool = False
|
||||
# Wake-intent state file (bind-mounted so it survives a router restart).
|
||||
state_file: str | None = "/state/router-state.json"
|
||||
|
||||
def policy(self, depth: str) -> DepthPolicy:
|
||||
if depth == DEPTH_SLEEPING:
|
||||
return DepthPolicy(depth, self.hold_sleep_s, self.retry_after_sleep,
|
||||
self.est_sleep_s, "sleep")
|
||||
if depth == DEPTH_RESTARTING:
|
||||
return DepthPolicy(depth, self.hold_restart_s, self.retry_after_restart,
|
||||
self.est_restart_s, "a container restart")
|
||||
# DEPTH_OFFLOADED and anything unknown: be conservative (plan 6.2.1).
|
||||
return DepthPolicy(DEPTH_OFFLOADED, self.hold_offload_s, self.retry_after_offload,
|
||||
self.est_offload_s, "offload")
|
||||
|
||||
|
||||
# Sleep depths tracked by the router.
|
||||
DEPTH_AWAKE = "awake"
|
||||
DEPTH_SLEEPING = "sleeping" # /sleep?level=1 - weights moved to host RAM
|
||||
DEPTH_OFFLOADED = "offloaded" # /sleep?level=2 - RAM freed, weights on NFS
|
||||
DEPTH_RESTARTING = "restarting" # container down / cold-starting
|
||||
|
||||
|
||||
def load_config() -> Config:
|
||||
services: dict[str, ServiceConfig] = {
|
||||
"text": ServiceConfig(
|
||||
key="text",
|
||||
service=_str("VLLM_TEXT_SERVICE", "vllm-text"),
|
||||
model=_str("VLLM_TEXT_MODEL", "Qwen3.6-35B-A3B-FP8"),
|
||||
aliases=_csv("VLLM_TEXT_ALIASES",
|
||||
"qwen,qwen3,qwen3.6,qwen3.6-35b,qwen3.6-35b-a3b,text,chat,default"),
|
||||
),
|
||||
"ocr": ServiceConfig(
|
||||
key="ocr",
|
||||
service=_str("VLLM_OCR_SERVICE", "vllm-ocr"),
|
||||
model=_str("VLLM_OCR_MODEL", "OvisOCR2"),
|
||||
aliases=_csv("VLLM_OCR_ALIASES", "ocr,ovis,ovisocr,ovis-ocr"),
|
||||
),
|
||||
"embed": ServiceConfig(
|
||||
key="embed",
|
||||
service=_str("VLLM_EMBED_SERVICE", "vllm-embed"),
|
||||
model=_str("VLLM_EMBED_MODEL", "Qwen3-Embedding-8B"),
|
||||
aliases=_csv("VLLM_EMBED_ALIASES",
|
||||
"embed,embedding,embeddings,qwen3-embedding,qwen3-embedding-8b"),
|
||||
),
|
||||
}
|
||||
|
||||
index: dict[str, str] = {}
|
||||
for cfg in services.values():
|
||||
names = {cfg.model.lower(), cfg.key.lower(), *(a.lower() for a in cfg.aliases)}
|
||||
for name in names:
|
||||
index[name] = cfg.key
|
||||
|
||||
return Config(
|
||||
services=services,
|
||||
index=index,
|
||||
public_host=_str("ROUTER_PUBLIC_HOST", "0.0.0.0"),
|
||||
public_port=_int("ROUTER_PUBLIC_PORT", 8000),
|
||||
admin_host=_str("ROUTER_ADMIN_HOST", "0.0.0.0"),
|
||||
admin_port=_int("ROUTER_ADMIN_PORT", 8010),
|
||||
connect_timeout=_float("ROUTER_CONNECT_TIMEOUT", 5.0),
|
||||
read_timeout=_float("ROUTER_READ_TIMEOUT", 900.0),
|
||||
pool_size=_int("ROUTER_POOL_SIZE", 32),
|
||||
probe_timeout=_float("ROUTER_PROBE_TIMEOUT", 3.0),
|
||||
max_body_bytes=_int("ROUTER_MAX_BODY_BYTES", 128 * 1024 * 1024),
|
||||
state_cache_ttl=_float("ROUTER_STATE_CACHE_TTL", 2.0),
|
||||
hold_sleep_s=_float("WAKE_HOLD_SLEEP_S", 30.0),
|
||||
hold_offload_s=_float("WAKE_HOLD_OFFLOAD_S", 180.0),
|
||||
hold_restart_s=_float("WAKE_HOLD_RESTART_S", 600.0),
|
||||
retry_after_sleep=_int("WAKE_RETRY_AFTER_SLEEP", 10),
|
||||
retry_after_offload=_int("WAKE_RETRY_AFTER_OFFLOAD", 60),
|
||||
retry_after_restart=_int("WAKE_RETRY_AFTER_RESTART", 600),
|
||||
est_sleep_s=_int("WAKE_EST_SLEEP_S", 6),
|
||||
est_offload_s=_int("WAKE_EST_OFFLOAD_S", 60),
|
||||
est_restart_s=_int("WAKE_EST_RESTART_S", 600),
|
||||
wake_attempts=max(1, _int("WAKE_ATTEMPTS", 2)),
|
||||
wake_http_timeout=_float("WAKE_HTTP_TIMEOUT", 300.0),
|
||||
wake_health_poll_s=_float("WAKE_HEALTH_POLL_S", 2.0),
|
||||
idle_sleep_min=_float("IDLE_SLEEP_MIN", 15.0),
|
||||
idle_offload_min=_float("IDLE_OFFLOAD_MIN", 180.0),
|
||||
idle_poll_s=_float("IDLE_POLL_SECONDS", 30.0),
|
||||
idle_enabled=_bool("IDLE_ENABLED", True),
|
||||
metrics_enabled=_bool("ROUTER_METRICS", False),
|
||||
log_level=_str("ROUTER_LOG_LEVEL", "info").lower(),
|
||||
access_log=_bool("ROUTER_ACCESS_LOG", False),
|
||||
state_file=_str("ROUTER_STATE_FILE", "/state/router-state.json") or None,
|
||||
)
|
||||
|
||||
|
||||
def clone(cfg: Config) -> Config:
|
||||
"""Copy for tests."""
|
||||
return replace(cfg, services=dict(cfg.services), index=dict(cfg.index))
|
||||
5
router/pytest.ini
Normal file
5
router/pytest.ini
Normal file
@@ -0,0 +1,5 @@
|
||||
[pytest]
|
||||
testpaths = tests
|
||||
asyncio_mode = auto
|
||||
filterwarnings =
|
||||
error::DeprecationWarning:tests.*
|
||||
10
router/requirements.txt
Normal file
10
router/requirements.txt
Normal file
@@ -0,0 +1,10 @@
|
||||
# Pinned (plan section 6). Direct dependencies only; the transitive set
|
||||
# (starlette, pydantic, anyio, ...) is resolved by pip at build time and
|
||||
# recorded by the image build.
|
||||
#
|
||||
# fastapi == 0.141.1 (current stable, verified with starlette 1.6)
|
||||
# uvicorn == 0.52.3 (two Server instances on one asyncio loop)
|
||||
# httpx == 0.28.1 (AsyncClient streaming proxy)
|
||||
fastapi==0.141.1
|
||||
uvicorn==0.52.3
|
||||
httpx==0.28.1
|
||||
536
router/routing.py
Normal file
536
router/routing.py
Normal file
@@ -0,0 +1,536 @@
|
||||
"""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
|
||||
695
router/services.py
Normal file
695
router/services.py
Normal file
@@ -0,0 +1,695 @@
|
||||
"""Per-service state, wake single-flight, tiered idle manager.
|
||||
|
||||
Everything here is deliberately single-event-loop, in-process: the wake locks
|
||||
must be shared between the public listener (:8000) and the admin listener
|
||||
(:8010), which is why the whole router is one process (plan section 6).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from config import ( # noqa: F401 (config is a sibling module, see app.py)
|
||||
DEPTH_AWAKE,
|
||||
DEPTH_OFFLOADED,
|
||||
DEPTH_RESTARTING,
|
||||
DEPTH_SLEEPING,
|
||||
Config,
|
||||
ServiceConfig,
|
||||
)
|
||||
|
||||
log = logging.getLogger("vllm_router.services")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# httpx client factories (module level so tests can monkeypatch / inject)
|
||||
# --------------------------------------------------------------------------
|
||||
def make_proxy_client(cfg: Config, transport: httpx.AsyncBaseTransport | None = None
|
||||
) -> httpx.AsyncClient:
|
||||
"""Client used to proxy requests and to drive the wake/sleep calls.
|
||||
|
||||
connect 5s / read 900s (plan 6.2 "Proxy timeouts": capped, never disabled)
|
||||
and a pool sized for >= 16 concurrent requests.
|
||||
"""
|
||||
return httpx.AsyncClient(
|
||||
transport=transport,
|
||||
timeout=httpx.Timeout(cfg.read_timeout, connect=cfg.connect_timeout),
|
||||
limits=httpx.Limits(
|
||||
max_connections=cfg.pool_size,
|
||||
max_keepalive_connections=max(4, cfg.pool_size // 2),
|
||||
),
|
||||
# The vLLM backends are plain HTTP on the docker network.
|
||||
trust_env=False,
|
||||
)
|
||||
|
||||
|
||||
def make_ctrl_client(cfg: Config, transport: httpx.AsyncBaseTransport | None = None
|
||||
) -> httpx.AsyncClient:
|
||||
"""Short-timeout client for state probes (/is_sleeping, /health)."""
|
||||
return httpx.AsyncClient(
|
||||
transport=transport,
|
||||
timeout=httpx.Timeout(cfg.probe_timeout, connect=cfg.connect_timeout),
|
||||
limits=httpx.Limits(max_connections=8, max_keepalive_connections=4),
|
||||
trust_env=False,
|
||||
)
|
||||
|
||||
|
||||
class WakeError(RuntimeError):
|
||||
"""A step of the wake sequence failed."""
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# persisted wake intent (plan 6.4, E2E case 12)
|
||||
# --------------------------------------------------------------------------
|
||||
class StateStore:
|
||||
"""Tiny JSON file that survives router restarts.
|
||||
|
||||
`is_sleeping=false` does NOT mean "reload done": if the previous router
|
||||
died between POST /wake_up and collective_rpc reload_weights, the backend
|
||||
happily reports awake+healthy while serving garbage. The file records
|
||||
wake intent and depth so the new router can tell the two apart.
|
||||
|
||||
Writes are atomic (tmp + rename) and synchronous -- the file is a few
|
||||
hundred bytes on a local bind mount and is only written on state
|
||||
transitions, never per request.
|
||||
"""
|
||||
|
||||
VERSION = 1
|
||||
|
||||
def __init__(self, path: str | None) -> None:
|
||||
self.path = path
|
||||
self.enabled = bool(path)
|
||||
self._warned = False
|
||||
|
||||
def load(self) -> dict[str, dict]:
|
||||
if not self.path:
|
||||
return {}
|
||||
try:
|
||||
with open(self.path, encoding="utf-8") as fh:
|
||||
data = json.load(fh)
|
||||
entries = data.get("services")
|
||||
return dict(entries) if isinstance(entries, dict) else {}
|
||||
except FileNotFoundError:
|
||||
return {}
|
||||
except Exception as exc: # corrupt/unreadable -> start from scratch
|
||||
log.warning("state_file_read_failed path=%s error=%s", self.path, exc)
|
||||
return {}
|
||||
|
||||
def write(self, services: dict[str, "ServiceState"]) -> None:
|
||||
if not self.path or not self.enabled:
|
||||
return
|
||||
payload = {
|
||||
"version": self.VERSION,
|
||||
"updated": time.time(),
|
||||
"services": {
|
||||
key: {
|
||||
"depth": svc.depth,
|
||||
"wake_in_progress": svc.wake_intent,
|
||||
"level": _LEVEL_BY_DEPTH.get(svc.depth, 0),
|
||||
"pending_reload": svc.pending_reload,
|
||||
}
|
||||
for key, svc in services.items()
|
||||
},
|
||||
}
|
||||
tmp = f"{self.path}.tmp"
|
||||
try:
|
||||
with open(tmp, "w", encoding="utf-8") as fh:
|
||||
json.dump(payload, fh)
|
||||
os.replace(tmp, self.path)
|
||||
except Exception as exc:
|
||||
self.enabled = False # never let state IO break serving
|
||||
if not self._warned:
|
||||
self._warned = True
|
||||
log.warning("state_file_write_failed path=%s disabled=true error=%s",
|
||||
self.path, exc)
|
||||
|
||||
|
||||
_LEVEL_BY_DEPTH = {DEPTH_SLEEPING: 1, DEPTH_OFFLOADED: 2}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# state
|
||||
# --------------------------------------------------------------------------
|
||||
@dataclass
|
||||
class WakeOutcome:
|
||||
"""Result of ensure_awake(); `ok=False` maps to a depth-aware 503."""
|
||||
|
||||
ok: bool
|
||||
key: str
|
||||
depth: str
|
||||
reason: str # cached | probe | woke | hold_timeout | wake_failed
|
||||
latency_s: float | None = None
|
||||
detail: str = ""
|
||||
|
||||
@property
|
||||
def held(self) -> bool:
|
||||
return self.reason in ("hold_timeout", "wake_failed")
|
||||
|
||||
|
||||
class ServiceState:
|
||||
"""Mutable per-service bookkeeping. Single event loop -> no extra locks
|
||||
beyond `lock`, which serializes wake vs. sleep vs. the idle manager."""
|
||||
|
||||
__slots__ = ("cfg", "lock", "active_requests", "last_activity", "depth",
|
||||
"state_checked_at", "reachable", "wake_task", "wake_depth",
|
||||
"last_wake_latency", "wake_count", "sleep_count", "last_error",
|
||||
"pending_reload", "wake_intent", "on_change")
|
||||
|
||||
def __init__(self, cfg: ServiceConfig) -> None:
|
||||
self.cfg = cfg
|
||||
self.lock = asyncio.Lock()
|
||||
self.active_requests = 0
|
||||
self.last_activity = time.monotonic()
|
||||
self.depth: str | None = None # None == unknown (fresh router)
|
||||
self.state_checked_at = float("-inf")
|
||||
self.reachable: bool | None = None
|
||||
self.wake_task: asyncio.Task | None = None
|
||||
self.wake_depth: str | None = None
|
||||
self.last_wake_latency: float | None = None
|
||||
self.wake_count = 0
|
||||
self.sleep_count = 0
|
||||
self.last_error: str | None = None
|
||||
# True when a previous router died between wake_up and reload_weights:
|
||||
# the backend *looks* awake but its weights were never reloaded.
|
||||
self.pending_reload = False
|
||||
# Persisted wake intent: True from the first POST /wake_up of a wake
|
||||
# until that wake completes or definitively fails.
|
||||
self.wake_intent = False
|
||||
self.on_change: Any = None # callback(services) -> None
|
||||
|
||||
# -- cached-state helpers ------------------------------------------------
|
||||
def set_depth(self, depth: str, reachable: bool = True) -> None:
|
||||
changed = self.depth != depth
|
||||
self.depth = depth
|
||||
self.reachable = reachable
|
||||
self.state_checked_at = time.monotonic()
|
||||
if changed and self.on_change is not None:
|
||||
self.on_change()
|
||||
|
||||
def awake_cached(self, ttl: float) -> bool:
|
||||
return (
|
||||
self.depth == DEPTH_AWAKE
|
||||
and self.state_checked_at + ttl >= time.monotonic()
|
||||
)
|
||||
|
||||
@property
|
||||
def wake_in_progress(self) -> bool:
|
||||
return self.wake_task is not None and not self.wake_task.done()
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# manager
|
||||
# --------------------------------------------------------------------------
|
||||
class ServiceManager:
|
||||
def __init__(self, cfg: Config, transport: httpx.AsyncBaseTransport | None = None) -> None:
|
||||
self.cfg = cfg
|
||||
self.services: dict[str, ServiceState] = {
|
||||
key: ServiceState(sc) for key, sc in cfg.services.items()
|
||||
}
|
||||
self.transport = transport
|
||||
self.proxy_client = make_proxy_client(cfg, transport)
|
||||
self.ctrl_client = make_ctrl_client(cfg, transport)
|
||||
self.state_store = StateStore(cfg.state_file)
|
||||
self.shutting_down = False
|
||||
self._idle_task: asyncio.Task | None = None
|
||||
self._recover_task: asyncio.Task | None = None
|
||||
self._started = False
|
||||
self._stopping = False
|
||||
self.started_at = time.monotonic()
|
||||
for svc in self.services.values():
|
||||
svc.on_change = self._persist_state
|
||||
# Read the persisted wake intent immediately, so no request can be
|
||||
# served on the strength of a stale "awake" before the recovery task
|
||||
# has had a chance to look at the backends.
|
||||
self._load_startup_state()
|
||||
|
||||
# ---- lifecycle --------------------------------------------------------
|
||||
def start(self) -> None:
|
||||
"""Idempotent: both listeners' lifespans call this."""
|
||||
if self._started:
|
||||
return
|
||||
self._started = True
|
||||
self._stopping = False
|
||||
self.started_at = time.monotonic()
|
||||
if self.cfg.idle_enabled:
|
||||
self._idle_task = asyncio.create_task(self._idle_loop(), name="idle-manager")
|
||||
log.info("idle_manager_start sleep_min=%s offload_min=%s poll_s=%s",
|
||||
self.cfg.idle_sleep_min, self.cfg.idle_offload_min, self.cfg.idle_poll_s)
|
||||
# Heals an interrupted wake even with no traffic (E2E case 12).
|
||||
self._recover_task = asyncio.create_task(
|
||||
self._recover_interrupted_wakes(), name="wake-recovery")
|
||||
|
||||
def begin_shutdown(self) -> None:
|
||||
"""Set before the listeners stop, so in-flight failures can say why."""
|
||||
self.shutting_down = True
|
||||
|
||||
def _persist_state(self) -> None:
|
||||
self.state_store.write(self.services)
|
||||
|
||||
def _load_startup_state(self) -> None:
|
||||
"""Restore depth, and flag services whose wake was interrupted."""
|
||||
entries = self.state_store.load()
|
||||
for key, svc in self.services.items():
|
||||
entry = entries.get(key)
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
depth = entry.get("depth")
|
||||
if depth in (DEPTH_SLEEPING, DEPTH_OFFLOADED, DEPTH_RESTARTING):
|
||||
svc.depth = depth
|
||||
svc.state_checked_at = float("-inf") # force a re-probe
|
||||
if entry.get("pending_reload"):
|
||||
svc.pending_reload = True
|
||||
if entry.get("wake_in_progress"):
|
||||
# The previous router died mid-wake. Until proven otherwise
|
||||
# this backend is "awake but not reloaded" -- the most
|
||||
# dangerous state, because /is_sleeping and /health both lie.
|
||||
svc.pending_reload = True
|
||||
log.warning("startup_wake_interrupted service=%s depth=%s", key, depth)
|
||||
if any(svc.pending_reload for svc in self.services.values()):
|
||||
log.warning("startup_recovery_pending services=%s",
|
||||
",".join(k for k, s in self.services.items() if s.pending_reload))
|
||||
self._persist_state()
|
||||
|
||||
async def _recover_interrupted_wakes(self) -> None:
|
||||
"""Complete (or discard) wakes a previous router left half-done."""
|
||||
for key, svc in list(self.services.items()):
|
||||
if not svc.pending_reload:
|
||||
continue
|
||||
try:
|
||||
reachable, sleeping = await self.probe_is_sleeping(svc)
|
||||
except Exception: # pragma: no cover - probe is defensive
|
||||
reachable, sleeping = False, None
|
||||
if not reachable:
|
||||
# Container (re)booting: a fresh vLLM boots with fresh
|
||||
# weights, so there is nothing to complete.
|
||||
svc.pending_reload = False
|
||||
svc.depth = None
|
||||
svc.state_checked_at = float("-inf")
|
||||
log.info("recovery_cleared service=%s reason=unreachable", key)
|
||||
self._persist_state()
|
||||
continue
|
||||
if sleeping:
|
||||
# The wake never took effect; the normal request path will
|
||||
# run a full wake when traffic arrives.
|
||||
svc.pending_reload = False
|
||||
log.info("recovery_cleared service=%s reason=still_sleeping", key)
|
||||
self._persist_state()
|
||||
continue
|
||||
log.info("recovery_start service=%s action=complete_reload_sequence", key)
|
||||
outcome = await self.ensure_awake(key)
|
||||
log.info("recovery_finish service=%s ok=%s reason=%s",
|
||||
key, outcome.ok, outcome.reason)
|
||||
self._persist_state()
|
||||
|
||||
async def stop(self) -> None:
|
||||
if not self._started:
|
||||
return
|
||||
self._stopping = True
|
||||
for name in ("_idle_task", "_recover_task"):
|
||||
task = getattr(self, name, None)
|
||||
if task is not None:
|
||||
task.cancel()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
setattr(self, name, None)
|
||||
for svc in self.services.values():
|
||||
if svc.wake_task is not None:
|
||||
svc.wake_task.cancel()
|
||||
await asyncio.gather(svc.wake_task, return_exceptions=True)
|
||||
svc.wake_task = None
|
||||
for client in (self.proxy_client, self.ctrl_client):
|
||||
try:
|
||||
await client.aclose()
|
||||
except Exception: # pragma: no cover - shutdown best effort
|
||||
pass
|
||||
self._started = False
|
||||
log.info("manager_stop")
|
||||
|
||||
# ---- activity tracking (plan 6.3 race safety) -------------------------
|
||||
def begin_request(self, svc: ServiceState) -> None:
|
||||
svc.active_requests += 1
|
||||
svc.last_activity = time.monotonic()
|
||||
|
||||
def end_request(self, svc: ServiceState) -> None:
|
||||
svc.active_requests = max(0, svc.active_requests - 1)
|
||||
# Refreshed at *completion* too: a long stream must not be seen as idle.
|
||||
svc.last_activity = time.monotonic()
|
||||
|
||||
# ---- probes -----------------------------------------------------------
|
||||
async def probe_is_sleeping(self, svc: ServiceState) -> tuple[bool, bool | None]:
|
||||
"""Returns (reachable, is_sleeping). is_sleeping is None when the
|
||||
endpoint answered but the payload was unusable."""
|
||||
try:
|
||||
resp = await self.ctrl_client.get(f"{svc.cfg.base_url}/is_sleeping")
|
||||
except httpx.HTTPError as exc:
|
||||
svc.reachable = False
|
||||
svc.last_error = f"is_sleeping: {type(exc).__name__}"
|
||||
return False, None
|
||||
svc.reachable = True
|
||||
if resp.status_code != 200:
|
||||
return True, None
|
||||
try:
|
||||
value = resp.json().get("is_sleeping")
|
||||
except Exception:
|
||||
return True, None
|
||||
if not isinstance(value, bool):
|
||||
return True, None
|
||||
return True, value
|
||||
|
||||
async def probe_health(self, svc: ServiceState) -> bool:
|
||||
try:
|
||||
resp = await self.ctrl_client.get(f"{svc.cfg.base_url}/health")
|
||||
except httpx.HTTPError:
|
||||
return False
|
||||
return resp.status_code == 200
|
||||
|
||||
async def classify(self, svc: ServiceState) -> str:
|
||||
"""Best-effort current depth. Callers hold svc.lock (or are the wake
|
||||
task, which is mutually exclusive with the lock by construction)."""
|
||||
if svc.pending_reload:
|
||||
# is_sleeping=false + /health 200 does NOT mean "ready" here: the
|
||||
# previous router may have died between wake_up and reload_weights
|
||||
# (E2E case 12). Conservative -> the full sequence runs again.
|
||||
return DEPTH_OFFLOADED
|
||||
reachable, sleeping = await self.probe_is_sleeping(svc)
|
||||
if not reachable:
|
||||
return DEPTH_RESTARTING
|
||||
if sleeping is None:
|
||||
return DEPTH_AWAKE if await self.probe_health(svc) else DEPTH_RESTARTING
|
||||
if not sleeping:
|
||||
if await self.probe_health(svc):
|
||||
svc.set_depth(DEPTH_AWAKE)
|
||||
return DEPTH_AWAKE
|
||||
return DEPTH_RESTARTING
|
||||
# Sleeping. If we do not know the level (fresh router, or a sleep done
|
||||
# behind our back) assume the worst case (plan 6.2.1).
|
||||
if svc.depth in (DEPTH_SLEEPING, DEPTH_OFFLOADED):
|
||||
return svc.depth
|
||||
return DEPTH_OFFLOADED
|
||||
|
||||
# ---- wake (plan 6.2 step 2/3) ----------------------------------------
|
||||
async def ensure_awake(self, key: str, hold_s: float | None = None) -> WakeOutcome:
|
||||
"""Single-flight wake for one service.
|
||||
|
||||
The first caller runs the wake sequence as a *shielded* task; everyone
|
||||
else (and late arrivals) join the same task. A caller that exceeds its
|
||||
depth-dependent hold deadline walks away with a 503 while the wake keeps
|
||||
running -- nobody ever proxies to a half-awake backend.
|
||||
"""
|
||||
svc = self.services[key]
|
||||
ttl = self.cfg.state_cache_ttl
|
||||
|
||||
if svc.awake_cached(ttl):
|
||||
return WakeOutcome(True, key, DEPTH_AWAKE, "cached")
|
||||
|
||||
async with svc.lock: # single-flight; also excludes the idle sleeper
|
||||
if svc.awake_cached(ttl):
|
||||
return WakeOutcome(True, key, DEPTH_AWAKE, "cached")
|
||||
task = svc.wake_task
|
||||
if task is None or task.done():
|
||||
depth = await self.classify(svc)
|
||||
if depth == DEPTH_AWAKE:
|
||||
return WakeOutcome(True, key, DEPTH_AWAKE, "probe")
|
||||
log.info("wake_start service=%s depth=%s", key, depth)
|
||||
svc.wake_depth = depth
|
||||
svc.wake_intent = True # persisted before /wake_up fires
|
||||
self._persist_state()
|
||||
svc.wake_task = asyncio.create_task(
|
||||
self._wake_sequence(svc, depth), name=f"wake-{key}"
|
||||
)
|
||||
task = svc.wake_task
|
||||
depth = svc.wake_depth or DEPTH_OFFLOADED
|
||||
|
||||
policy = self.cfg.policy(depth)
|
||||
deadline = policy.hold_s if hold_s is None else hold_s
|
||||
try:
|
||||
# shield(): our timeout must not cancel a wake other callers share.
|
||||
outcome = await asyncio.wait_for(asyncio.shield(task), timeout=deadline)
|
||||
return outcome
|
||||
except (asyncio.TimeoutError, TimeoutError):
|
||||
log.warning("wake_hold_timeout service=%s depth=%s hold_s=%s",
|
||||
key, depth, deadline)
|
||||
return WakeOutcome(False, key, depth, "hold_timeout")
|
||||
|
||||
async def _wake_sequence(self, svc: ServiceState, depth: str) -> WakeOutcome:
|
||||
cfg = self.cfg
|
||||
key = svc.cfg.key
|
||||
for attempt in range(1, cfg.wake_attempts + 1):
|
||||
t0 = time.monotonic()
|
||||
try:
|
||||
if depth == DEPTH_RESTARTING:
|
||||
# Container down / cold-starting (text cold boot measured
|
||||
# up to ~10 min from NFS -> 600s hold). `restart:
|
||||
# unless-stopped` brings it back and vLLM boots awake.
|
||||
if not await self._wait_ready(svc, cfg.hold_restart_s + 60.0):
|
||||
raise WakeError("backend did not become ready")
|
||||
elif depth == DEPTH_SLEEPING:
|
||||
# Level-1: /wake_up ALONE restores bit-identical output
|
||||
# (calibration 2026-08-17, temp-0 verified); the reload
|
||||
# sequence would only add ~20s on the text model
|
||||
# (23.4s -> 2.5-3.8s).
|
||||
await self._wake_level1(svc)
|
||||
else:
|
||||
# Level-2, or depth unknown -> conservative offloaded.
|
||||
await self._wake_level2(svc)
|
||||
latency = time.monotonic() - t0
|
||||
svc.set_depth(DEPTH_AWAKE)
|
||||
svc.pending_reload = False
|
||||
svc.wake_intent = False # restart after this point is cheap
|
||||
svc.last_wake_latency = latency
|
||||
svc.wake_count += 1
|
||||
svc.last_error = None
|
||||
self._persist_state()
|
||||
log.info("wake_finish service=%s from_depth=%s attempt=%d latency_s=%.2f",
|
||||
key, depth, attempt, latency)
|
||||
return WakeOutcome(True, key, DEPTH_AWAKE, "woke", latency)
|
||||
except Exception as exc: # noqa: BLE001 - logged, retried, surfaced
|
||||
svc.last_error = f"{type(exc).__name__}: {exc}"
|
||||
log.warning("wake_failed service=%s depth=%s attempt=%d/%d error=%s",
|
||||
key, depth, attempt, cfg.wake_attempts, svc.last_error)
|
||||
try:
|
||||
# Re-classify: the backend may have restarted under us.
|
||||
depth = await self.classify(svc)
|
||||
if depth == DEPTH_AWAKE:
|
||||
latency = time.monotonic() - t0
|
||||
svc.set_depth(DEPTH_AWAKE)
|
||||
svc.pending_reload = False
|
||||
svc.wake_intent = False
|
||||
svc.wake_count += 1
|
||||
svc.last_wake_latency = latency
|
||||
self._persist_state()
|
||||
log.info("wake_finish service=%s recovered attempt=%d latency_s=%.2f",
|
||||
key, attempt, latency)
|
||||
return WakeOutcome(True, key, DEPTH_AWAKE, "woke", latency)
|
||||
except Exception: # pragma: no cover - classify is defensive
|
||||
pass
|
||||
if attempt < cfg.wake_attempts:
|
||||
await asyncio.sleep(1.0)
|
||||
log.error("wake_exhausted service=%s depth=%s error=%s", key, depth, svc.last_error)
|
||||
# Definitive failure: stop claiming a wake is in flight (the recovery
|
||||
# path would otherwise re-run the sequence after every restart).
|
||||
svc.wake_intent = False
|
||||
svc.pending_reload = False
|
||||
self._persist_state()
|
||||
return WakeOutcome(False, key, depth, "wake_failed", detail=svc.last_error or "")
|
||||
|
||||
async def _wake_level1(self, svc: ServiceState) -> None:
|
||||
"""Level-1 (weights still in host RAM): POST /wake_up alone."""
|
||||
await self._post(f"{svc.cfg.base_url}/wake_up")
|
||||
if not await self._wait_ready(svc, self.cfg.hold_sleep_s + 60.0):
|
||||
raise WakeError("backend still reports is_sleeping after /wake_up")
|
||||
|
||||
async def _wake_level2(self, svc: ServiceState) -> None:
|
||||
"""Level-2 (weights on NFS) or unknown depth: the full sequence is
|
||||
REQUIRED. Requests admitted between wake_up and reload_weights return
|
||||
200 + garbage, so nothing is proxied until this completes."""
|
||||
base = svc.cfg.base_url
|
||||
await self._post(f"{base}/wake_up")
|
||||
await self._post(f"{base}/collective_rpc", json={"method": "reload_weights"})
|
||||
try:
|
||||
await self._post(f"{base}/reset_prefix_cache")
|
||||
except WakeError as exc:
|
||||
log.warning("reset_prefix_cache_failed service=%s error=%s", svc.cfg.key, exc)
|
||||
if not await self._wait_ready(svc, self.cfg.hold_offload_s + 60.0):
|
||||
raise WakeError("backend did not become ready after wake sequence")
|
||||
|
||||
async def _wait_ready(self, svc: ServiceState, budget_s: float) -> bool:
|
||||
"""Readiness gate.
|
||||
|
||||
/health LIES on a sleeping backend (200 while asleep, calibration
|
||||
2026-08-17), and a request proxied to a sleeping backend hangs rather
|
||||
than erroring. So readiness is `GET /is_sleeping -> false`; /health
|
||||
200 is only an extra sanity check once that has happened.
|
||||
"""
|
||||
deadline = time.monotonic() + budget_s
|
||||
while True:
|
||||
reachable, sleeping = await self.probe_is_sleeping(svc)
|
||||
if reachable and sleeping is False and await self.probe_health(svc):
|
||||
return True
|
||||
if time.monotonic() >= deadline:
|
||||
return False
|
||||
await asyncio.sleep(self.cfg.wake_health_poll_s)
|
||||
|
||||
async def _post(self, url: str, json: dict | None = None) -> httpx.Response:
|
||||
try:
|
||||
resp = await self.proxy_client.post(url, json=json)
|
||||
except httpx.HTTPError as exc:
|
||||
raise WakeError(f"POST {url}: {type(exc).__name__}") from exc
|
||||
if resp.status_code >= 300:
|
||||
raise WakeError(f"POST {url} -> HTTP {resp.status_code}")
|
||||
return resp
|
||||
|
||||
# ---- sleep ------------------------------------------------------------
|
||||
async def sleep_service(self, key: str, level: int = 1, *, reason: str = "manual",
|
||||
min_idle_s: float = 0.0) -> dict[str, Any]:
|
||||
"""Sleep one service. All refusals are re-checked *under* the lock so
|
||||
the idle manager cannot lose a race with an arriving request."""
|
||||
svc = self.services[key]
|
||||
level = 2 if int(level) == 2 else 1
|
||||
target = DEPTH_OFFLOADED if level == 2 else DEPTH_SLEEPING
|
||||
|
||||
async with svc.lock:
|
||||
if svc.wake_in_progress:
|
||||
return {"ok": False, "service": key, "reason": "wake_in_progress"}
|
||||
if svc.active_requests > 0:
|
||||
return {"ok": False, "service": key, "reason": "active_requests",
|
||||
"active_requests": svc.active_requests}
|
||||
idle_s = time.monotonic() - svc.last_activity
|
||||
if min_idle_s and idle_s < min_idle_s:
|
||||
return {"ok": False, "service": key, "reason": "activity_resumed",
|
||||
"idle_s": round(idle_s, 1)}
|
||||
|
||||
reachable, sleeping = await self.probe_is_sleeping(svc)
|
||||
if not reachable:
|
||||
svc.depth = DEPTH_RESTARTING
|
||||
return {"ok": False, "service": key, "reason": "unreachable"}
|
||||
|
||||
if sleeping:
|
||||
deeper = (svc.depth == DEPTH_OFFLOADED) or (
|
||||
level == 1 and svc.depth == DEPTH_SLEEPING)
|
||||
if deeper:
|
||||
# Already at (or below) the requested tier.
|
||||
svc.state_checked_at = time.monotonic()
|
||||
return {"ok": True, "service": key, "level": level,
|
||||
"already": True, "depth": svc.depth}
|
||||
if level == 2:
|
||||
# Sleeping at level 1 and asked for level 2. A direct
|
||||
# POST /sleep?level=2 is a well-behaved NO-OP that RETAINS
|
||||
# the host-RAM weights copy (calibration 2026-08-17; the
|
||||
# vLLM allocator never frees existing cpu_backup tensors),
|
||||
# so we must wake into RAM first and then offload.
|
||||
log.info("sleep_escalate service=%s from=sleeping action=wake_then_offload",
|
||||
key)
|
||||
try:
|
||||
await self._post(f"{svc.cfg.base_url}/wake_up")
|
||||
if not await self._wait_ready(svc, self.cfg.hold_sleep_s + 60.0):
|
||||
raise WakeError("escalation wake did not become ready")
|
||||
await self._post(f"{svc.cfg.base_url}/sleep?level=2")
|
||||
except WakeError as exc:
|
||||
svc.state_checked_at = time.monotonic()
|
||||
return {"ok": False, "service": key,
|
||||
"reason": "escalation_failed", "error": str(exc)}
|
||||
else:
|
||||
svc.state_checked_at = time.monotonic()
|
||||
return {"ok": True, "service": key, "level": level,
|
||||
"already": True, "depth": target}
|
||||
else:
|
||||
# Awake: level 1 or level 2 can be entered directly.
|
||||
try:
|
||||
await self._post(f"{svc.cfg.base_url}/sleep?level={level}")
|
||||
except WakeError as exc:
|
||||
svc.state_checked_at = time.monotonic()
|
||||
return {"ok": False, "service": key, "reason": "sleep_failed",
|
||||
"error": str(exc)}
|
||||
|
||||
_, still = await self.probe_is_sleeping(svc)
|
||||
if not still:
|
||||
svc.state_checked_at = time.monotonic()
|
||||
return {"ok": False, "service": key, "reason": "sleep_not_confirmed"}
|
||||
svc.set_depth(target)
|
||||
svc.sleep_count += 1
|
||||
log.info("sleep service=%s level=%d depth=%s reason=%s idle_s=%.0f",
|
||||
key, level, target, reason, idle_s)
|
||||
return {"ok": True, "service": key, "level": level, "depth": target}
|
||||
|
||||
# ---- tiered idle (plan 6.3) ------------------------------------------
|
||||
async def _idle_loop(self) -> None:
|
||||
while not self._stopping:
|
||||
try:
|
||||
await asyncio.sleep(self.cfg.idle_poll_s)
|
||||
await self.idle_tick()
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception: # pragma: no cover - never let the loop die
|
||||
log.exception("idle_tick_error")
|
||||
|
||||
async def idle_tick(self) -> None:
|
||||
cfg = self.cfg
|
||||
sleep_s = cfg.idle_sleep_min * 60.0
|
||||
offload_s = cfg.idle_offload_min * 60.0
|
||||
now = time.monotonic()
|
||||
for svc in self.services.values():
|
||||
if svc.active_requests > 0 or svc.wake_in_progress:
|
||||
continue
|
||||
if svc.depth == DEPTH_OFFLOADED or svc.depth == DEPTH_RESTARTING:
|
||||
continue # nothing left to shed
|
||||
idle_s = now - svc.last_activity
|
||||
if idle_s >= offload_s:
|
||||
log.info("idle_trigger service=%s idle_s=%.0f target=offload", svc.cfg.key, idle_s)
|
||||
await self.sleep_service(svc.cfg.key, 2, reason="idle", min_idle_s=offload_s)
|
||||
elif svc.depth == DEPTH_AWAKE and idle_s >= sleep_s:
|
||||
log.info("idle_trigger service=%s idle_s=%.0f target=sleep", svc.cfg.key, idle_s)
|
||||
await self.sleep_service(svc.cfg.key, 1, reason="idle", min_idle_s=sleep_s)
|
||||
|
||||
# ---- status -----------------------------------------------------------
|
||||
async def status(self, live_probe: bool = True) -> dict[str, Any]:
|
||||
"""Snapshot for /admin/status and /health."""
|
||||
async def probe(svc: ServiceState) -> tuple[bool, bool | None]:
|
||||
if not live_probe:
|
||||
return (svc.reachable, None)
|
||||
return await self.probe_is_sleeping(svc)
|
||||
|
||||
results = await asyncio.gather(*(probe(s) for s in self.services.values()))
|
||||
services: dict[str, Any] = {}
|
||||
now = time.monotonic()
|
||||
for svc, (reachable, sleeping) in zip(self.services.values(), results):
|
||||
if live_probe and reachable and sleeping is not None and not svc.pending_reload:
|
||||
# Keep the router's own depth tracking aligned with reality
|
||||
# (but never while a reload recovery is pending: the backend
|
||||
# *looks* awake then, and that is exactly the lie).
|
||||
if sleeping and svc.depth not in (DEPTH_SLEEPING, DEPTH_OFFLOADED):
|
||||
svc.depth = DEPTH_OFFLOADED # unknown level -> conservative
|
||||
elif not sleeping and svc.depth in (DEPTH_SLEEPING, DEPTH_OFFLOADED):
|
||||
svc.depth = DEPTH_AWAKE
|
||||
svc.state_checked_at = now
|
||||
services[svc.cfg.key] = {
|
||||
"key": svc.cfg.key,
|
||||
"service": svc.cfg.service,
|
||||
"model": svc.cfg.model,
|
||||
"base_url": svc.cfg.base_url,
|
||||
"reachable": reachable,
|
||||
"sleeping": sleeping,
|
||||
"depth": svc.depth,
|
||||
"depth_known": svc.depth is not None,
|
||||
"wake_in_progress": svc.wake_in_progress,
|
||||
"wake_recovery_pending": svc.pending_reload,
|
||||
"active_requests": svc.active_requests,
|
||||
"last_activity_ago_s": round(now - svc.last_activity, 1),
|
||||
"last_wake_latency_s": (
|
||||
round(svc.last_wake_latency, 2) if svc.last_wake_latency is not None else None
|
||||
),
|
||||
"wake_count": svc.wake_count,
|
||||
"sleep_count": svc.sleep_count,
|
||||
"last_error": svc.last_error,
|
||||
}
|
||||
return {
|
||||
"uptime_s": round(now - self.started_at, 1),
|
||||
"services": services,
|
||||
}
|
||||
263
router/tests/conftest.py
Normal file
263
router/tests/conftest.py
Normal file
@@ -0,0 +1,263 @@
|
||||
"""Shared fixtures + a fake vLLM backend transport (no GPU / docker / network).
|
||||
|
||||
The router's outbound httpx clients are built with an injectable transport
|
||||
(`ServiceManager(cfg, transport=...)`), so a single fake serves all three
|
||||
backends and records every call the router makes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
import urllib.parse
|
||||
from types import SimpleNamespace
|
||||
from pathlib import Path
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
|
||||
from app import build_apps # noqa: E402
|
||||
from config import clone, load_config # noqa: E402
|
||||
from services import ServiceManager # noqa: E402
|
||||
|
||||
|
||||
class FakeVLLM(httpx.AsyncBaseTransport):
|
||||
"""Stands in for vllm-text / vllm-ocr / vllm-embed.
|
||||
|
||||
Endpoints: /health, /is_sleeping, /wake_up, /collective_rpc,
|
||||
/reset_prefix_cache, /sleep?level=, plus /v1/* which echoes what it saw.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[str] = []
|
||||
self.services: dict[str, dict] = {
|
||||
"vllm-text": self._svc(),
|
||||
"vllm-ocr": self._svc(),
|
||||
"vllm-embed": self._svc(),
|
||||
}
|
||||
self.stream_chunks: list[bytes] = []
|
||||
self.stream_delay: float = 0.0 # per-chunk delay for slow/SSE streams
|
||||
|
||||
@staticmethod
|
||||
def _svc() -> dict:
|
||||
return {
|
||||
"reachable": True,
|
||||
"healthy": True,
|
||||
"sleeping": False,
|
||||
"wake_delay": 0.0,
|
||||
"wake_fails": 0, # number of /wake_up calls to fail (500)
|
||||
"sleep_fails": 0, # number of /sleep calls to fail (500)
|
||||
"wake_up_calls": 0,
|
||||
"reload_calls": 0,
|
||||
"reset_calls": 0,
|
||||
"sleep_calls": 0,
|
||||
"api_calls": 0,
|
||||
"api_delay": 0.0,
|
||||
}
|
||||
|
||||
# -- helpers -----------------------------------------------------------
|
||||
@staticmethod
|
||||
def _response(status: int = 200, *, payload: bytes = b"", content_type: str | None = None):
|
||||
"""Build a *streamable* response.
|
||||
|
||||
`httpx.Response(200, json=...)` would mark the body as already
|
||||
consumed, which the router's streaming proxy would rightly reject.
|
||||
"""
|
||||
async def gen():
|
||||
if payload:
|
||||
yield payload
|
||||
|
||||
headers = {"content-type": content_type} if content_type else {}
|
||||
return httpx.Response(status, content=gen(), headers=headers)
|
||||
|
||||
def count(self, needle: str) -> int:
|
||||
return sum(1 for call in self.calls if needle in call)
|
||||
|
||||
def last(self, needle: str) -> int:
|
||||
"""Index of the last matching recorded call (-1 if none)."""
|
||||
for i in range(len(self.calls) - 1, -1, -1):
|
||||
if needle in self.calls[i]:
|
||||
return i
|
||||
return -1
|
||||
|
||||
def set_sleeping(self, key: str, sleeping: bool, level: int = 0) -> None:
|
||||
svc = self.services[f"vllm-{key}"]
|
||||
svc["sleeping"] = sleeping
|
||||
svc["level"] = level
|
||||
|
||||
# -- transport ---------------------------------------------------------
|
||||
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
|
||||
host = request.url.host or ""
|
||||
svc = self.services.get(host)
|
||||
if svc is None or not svc["reachable"]:
|
||||
raise httpx.ConnectError(f"connection refused: {host}", request=request)
|
||||
|
||||
path = request.url.path
|
||||
params = request.url.params
|
||||
# record method + full target (path + query) so tests can assert on
|
||||
# things like "level=1"
|
||||
target = path + (f"?{params}" if str(params) else "")
|
||||
self.calls.append(f"{request.method} {host}{target}")
|
||||
|
||||
if path == "/is_sleeping":
|
||||
# Optional hook: report `true` for the next N polls even after a
|
||||
# wake-up, to model a backend whose wake is still in flight.
|
||||
if svc.get("hold_sleeping_polls", 0) > 0:
|
||||
svc["hold_sleeping_polls"] -= 1
|
||||
return self._response(200, payload=b'{"is_sleeping": true}',
|
||||
content_type="application/json")
|
||||
return self._response(200, payload=json.dumps(
|
||||
{"is_sleeping": bool(svc["sleeping"])}).encode(),
|
||||
content_type="application/json")
|
||||
|
||||
if path == "/health":
|
||||
# NOTE: like the real vLLM, /health answers 200 even while the
|
||||
# model is asleep -- readiness must be gated on /is_sleeping.
|
||||
if not svc["healthy"]:
|
||||
raise httpx.ConnectError(f"unhealthy: {host}", request=request)
|
||||
return self._response(200)
|
||||
|
||||
if path == "/wake_up":
|
||||
svc["wake_up_calls"] += 1
|
||||
if svc["wake_fails"] > 0:
|
||||
svc["wake_fails"] -= 1
|
||||
return self._response(500, payload=b"wake failed")
|
||||
if svc.get("wake_flip_midway") and svc["wake_delay"]:
|
||||
# Model the real hazard: is_sleeping flips to false *during*
|
||||
# /wake_up, while reload_weights has not happened yet. A
|
||||
# router that dies in this window leaves a backend that looks
|
||||
# perfectly awake and serves garbage (E2E case 12).
|
||||
await asyncio.sleep(svc["wake_delay"] / 2)
|
||||
svc["sleeping"] = False
|
||||
await asyncio.sleep(svc["wake_delay"] / 2)
|
||||
return self._response(200)
|
||||
await asyncio.sleep(svc["wake_delay"])
|
||||
svc["sleeping"] = False
|
||||
return self._response(200)
|
||||
|
||||
if path == "/collective_rpc":
|
||||
svc["reload_calls"] += 1
|
||||
return self._response(200)
|
||||
|
||||
if path == "/reset_prefix_cache":
|
||||
svc["reset_calls"] += 1
|
||||
return self._response(200)
|
||||
|
||||
if path == "/sleep":
|
||||
svc["sleep_calls"] += 1
|
||||
if svc["sleep_fails"] > 0:
|
||||
svc["sleep_fails"] -= 1
|
||||
return self._response(500, payload=b"sleep failed")
|
||||
svc["sleeping"] = True
|
||||
svc["level"] = int(params.get("level", "1"))
|
||||
return self._response(200)
|
||||
|
||||
if path.startswith("/v1/"):
|
||||
svc["api_calls"] += 1
|
||||
if svc.get("api_delay"):
|
||||
await asyncio.sleep(svc["api_delay"])
|
||||
ctype = request.headers.get("content-type", "")
|
||||
if "text/event-stream" in ctype or self.stream_chunks:
|
||||
async def gen():
|
||||
for chunk in self.stream_chunks:
|
||||
await asyncio.sleep(self.stream_delay)
|
||||
yield chunk
|
||||
return httpx.Response(200, content=gen(),
|
||||
headers={"content-type": "text/event-stream"})
|
||||
body = (request.content or b"").decode("utf-8", "replace")
|
||||
return self._response(
|
||||
200,
|
||||
payload=json.dumps({
|
||||
"echo_path": path,
|
||||
"echo_method": request.method,
|
||||
"echo_body": body,
|
||||
"echo_content_type": ctype,
|
||||
"echo_query": str(request.url.params),
|
||||
"service": host,
|
||||
}).encode(),
|
||||
content_type="application/json",
|
||||
)
|
||||
|
||||
return self._response(404, payload=json.dumps(
|
||||
{"detail": f"no route {path}"}).encode(), content_type="application/json")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def backend() -> FakeVLLM:
|
||||
return FakeVLLM()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def cfg():
|
||||
config = clone(load_config())
|
||||
config.state_cache_ttl = 0.0 # never trust cache -> deterministic
|
||||
config.wake_health_poll_s = 0.005
|
||||
config.idle_enabled = False # idle is exercised via idle_tick()
|
||||
config.hold_sleep_s = 5.0
|
||||
config.hold_offload_s = 5.0
|
||||
config.hold_restart_s = 5.0
|
||||
return config
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def stack(cfg, backend):
|
||||
public_app, admin_app, manager = build_apps(cfg, transport=backend)
|
||||
try:
|
||||
yield SimpleNamespace(public=public_app, admin=admin_app, manager=manager,
|
||||
cfg=cfg, backend=backend)
|
||||
finally:
|
||||
await manager.proxy_client.aclose()
|
||||
await manager.ctrl_client.aclose()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def pub(stack):
|
||||
transport = httpx.ASGITransport(app=stack.public)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://router.test") as client:
|
||||
yield client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def adm(stack):
|
||||
transport = httpx.ASGITransport(app=stack.admin)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://router.test") as client:
|
||||
yield client
|
||||
|
||||
|
||||
async def raw_asgi(app, method: str, raw_path: str, *, headers=None, body: bytes = b""):
|
||||
"""Call the ASGI app with a hand-built scope so the *raw* (percent-encoded)
|
||||
path reaches the router exactly as a hostile client would send it."""
|
||||
scope = {
|
||||
"type": "http", "asgi": {"version": "3.0", "spec_version": "2.3"},
|
||||
"method": method,
|
||||
"path": urllib.parse.unquote(raw_path),
|
||||
"raw_path": raw_path.encode("latin-1"),
|
||||
"query_string": b"",
|
||||
"headers": [(k.lower().encode(), v.encode()) for k, v in (headers or [])],
|
||||
"http_version": "1.1", "scheme": "http",
|
||||
"server": ("router.test", 80), "client": ("127.0.0.1", 1234),
|
||||
"root_path": "",
|
||||
}
|
||||
sent: list[dict] = []
|
||||
incoming: list[dict] = [{"type": "http.request", "body": body, "more_body": False}]
|
||||
closed = asyncio.Event()
|
||||
|
||||
async def receive():
|
||||
# Behaves like a real server: hand over the body once, then wait for
|
||||
# the client to hang up (Starlette's disconnect listener relies on it).
|
||||
if incoming:
|
||||
return incoming.pop(0)
|
||||
await closed.wait()
|
||||
return {"type": "http.disconnect"}
|
||||
|
||||
async def send(message):
|
||||
sent.append(message)
|
||||
|
||||
await app(scope, receive, send)
|
||||
closed.set()
|
||||
status = sent[0]["status"]
|
||||
payload = b"".join(m.get("body", b"") for m in sent if m["type"] == "http.response.body")
|
||||
return status, payload
|
||||
120
router/tests/test_admin_api.py
Normal file
120
router/tests/test_admin_api.py
Normal file
@@ -0,0 +1,120 @@
|
||||
"""Admin API surface (this is what `vllmctl` talks to)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
|
||||
from config import DEPTH_AWAKE, DEPTH_OFFLOADED
|
||||
|
||||
|
||||
async def test_status_shape(adm, backend):
|
||||
backend.set_sleeping("ocr", True, level=1)
|
||||
r = await adm.get("/admin/status")
|
||||
assert r.status_code == 200
|
||||
body = r.json()
|
||||
assert set(body["services"]) == {"text", "ocr", "embed"}
|
||||
text = body["services"]["text"]
|
||||
for field in ("service", "model", "base_url", "reachable", "sleeping",
|
||||
"depth", "wake_in_progress", "active_requests",
|
||||
"last_activity_ago_s", "last_wake_latency_s", "last_error"):
|
||||
assert field in text, field
|
||||
assert text["model"] == "Qwen3.6-35B-A3B-FP8"
|
||||
assert text["base_url"] == "http://vllm-text:8000"
|
||||
assert body["services"]["ocr"]["sleeping"] is True
|
||||
assert body["services"]["ocr"]["depth"] == "offloaded" # unknown level
|
||||
assert "http://127.0.0.1:8000/v1" in body["api"]
|
||||
|
||||
|
||||
async def test_wake_endpoint(adm, backend):
|
||||
backend.set_sleeping("embed", True)
|
||||
r = await adm.post("/admin/wake/embed")
|
||||
assert r.status_code == 200
|
||||
body = r.json()
|
||||
assert body["ok"] is True
|
||||
assert body["depth"] == "awake"
|
||||
assert backend.services["vllm-embed"]["wake_up_calls"] == 1
|
||||
|
||||
|
||||
async def test_wake_endpoint_reports_503_when_it_cannot_wake(adm, backend, stack):
|
||||
backend.set_sleeping("text", True)
|
||||
backend.services["vllm-text"]["wake_fails"] = 5
|
||||
r = await adm.post("/admin/wake/text")
|
||||
assert r.status_code == 503
|
||||
assert r.headers["retry-after"] == "60"
|
||||
body = r.json()
|
||||
assert body["ok"] is False
|
||||
assert body["error"]["sleep_depth"] == "offloaded"
|
||||
assert body["error"]["estimated_wake_seconds"] == 60
|
||||
|
||||
|
||||
async def test_wake_accepts_model_name_and_case(adm, backend):
|
||||
backend.set_sleeping("ocr", True)
|
||||
r = await adm.post("/admin/wake/OvisOCR2")
|
||||
assert r.status_code == 200
|
||||
assert backend.services["vllm-ocr"]["wake_up_calls"] == 1
|
||||
|
||||
|
||||
async def test_wake_unknown_key_is_404(adm):
|
||||
assert (await adm.post("/admin/wake/nope")).status_code == 404
|
||||
|
||||
|
||||
async def test_sleep_level1_and_level2(adm, backend, stack):
|
||||
r = await adm.post("/admin/sleep/text?level=1")
|
||||
assert r.status_code == 200
|
||||
assert r.json()["ok"] is True
|
||||
assert "level=1" in backend.calls[backend.last("POST vllm-text/sleep")]
|
||||
assert stack.manager.services["text"].depth == "sleeping"
|
||||
|
||||
r = await adm.post("/admin/sleep/text?level=2")
|
||||
assert r.status_code == 200
|
||||
assert "level=2" in backend.calls[backend.last("POST vllm-text/sleep")]
|
||||
assert stack.manager.services["text"].depth == DEPTH_OFFLOADED
|
||||
|
||||
|
||||
async def test_sleep_defaults_to_level1(adm, backend):
|
||||
r = await adm.post("/admin/sleep/embed")
|
||||
assert r.status_code == 200
|
||||
assert "level=1" in backend.calls[backend.last("POST vllm-embed/sleep")]
|
||||
|
||||
|
||||
async def test_sleep_rejects_bad_level(adm):
|
||||
assert (await adm.post("/admin/sleep/text?level=3")).status_code == 400
|
||||
assert (await adm.post("/admin/sleep/text?level=abc")).status_code == 400
|
||||
|
||||
|
||||
async def test_sleep_unknown_key_is_404(adm, backend):
|
||||
assert (await adm.post("/admin/sleep/nope")).status_code == 404
|
||||
assert backend.calls == []
|
||||
|
||||
|
||||
async def test_sleep_idempotent_when_already_at_depth(adm, backend, stack):
|
||||
await adm.post("/admin/sleep/ocr?level=1")
|
||||
backend.calls.clear()
|
||||
r = await adm.post("/admin/sleep/ocr?level=1")
|
||||
assert r.status_code == 200
|
||||
assert r.json()["already"] is True
|
||||
assert backend.count("POST vllm-ocr/sleep") == 0 # no second POST /sleep
|
||||
|
||||
|
||||
async def test_admin_wake_resets_the_idle_clock(adm, backend, stack):
|
||||
backend.set_sleeping("text", True)
|
||||
text = stack.manager.services["text"]
|
||||
text.last_activity -= 10_000
|
||||
before = time.monotonic() - text.last_activity
|
||||
assert (await adm.post("/admin/wake/text")).status_code == 200
|
||||
assert text.depth == DEPTH_AWAKE
|
||||
assert (time.monotonic() - text.last_activity) < before
|
||||
|
||||
|
||||
async def test_admin_health(adm):
|
||||
r = await adm.get("/health")
|
||||
assert r.status_code == 200
|
||||
assert r.json()["status"] == "ok"
|
||||
|
||||
|
||||
async def test_admin_health_alias(adm):
|
||||
"""The documented admin surface says /admin/health."""
|
||||
assert (await adm.get("/admin/health")).status_code == 200
|
||||
body = (await adm.get("/admin/health")).json()
|
||||
assert body["status"] == "ok"
|
||||
assert "awake" in body
|
||||
207
router/tests/test_depth_errors.py
Normal file
207
router/tests/test_depth_errors.py
Normal file
@@ -0,0 +1,207 @@
|
||||
"""Depth-aware 503 semantics (plan 6.2.1) and never-proxy-half-awake."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
import httpx
|
||||
|
||||
from conftest import raw_asgi
|
||||
|
||||
|
||||
async def _sleep_at_level(stack, level: int) -> None:
|
||||
result = await stack.manager.sleep_service("text", level, reason="test")
|
||||
assert result["ok"], result
|
||||
|
||||
|
||||
async def test_503_from_level1_sleep(pub, backend, stack):
|
||||
await _sleep_at_level(stack, 1)
|
||||
stack.cfg.hold_sleep_s = 0.05
|
||||
backend.services["vllm-text"]["wake_delay"] = 0.4
|
||||
r = await pub.post("/v1/chat/completions", json={"model": "text"})
|
||||
assert r.status_code == 503
|
||||
assert r.headers["retry-after"] == "10"
|
||||
err = r.json()["error"]
|
||||
assert err["code"] == "model_waking"
|
||||
assert err["type"] == "model_waking"
|
||||
assert err["sleep_depth"] == "sleeping"
|
||||
assert isinstance(err["estimated_wake_seconds"], int)
|
||||
assert err["estimated_wake_seconds"] == 6
|
||||
assert "Qwen3.6-35B-A3B-FP8" in err["message"]
|
||||
|
||||
|
||||
async def test_503_from_level2_offload(pub, backend, stack):
|
||||
await _sleep_at_level(stack, 2)
|
||||
stack.cfg.hold_offload_s = 0.05
|
||||
backend.services["vllm-text"]["wake_delay"] = 0.4
|
||||
r = await pub.post("/v1/chat/completions", json={"model": "text"})
|
||||
assert r.status_code == 503
|
||||
assert r.headers["retry-after"] == "60"
|
||||
err = r.json()["error"]
|
||||
assert err["sleep_depth"] == "offloaded"
|
||||
assert err["estimated_wake_seconds"] == 60
|
||||
|
||||
|
||||
async def test_503_when_container_restarting(pub, backend, stack):
|
||||
backend.services["vllm-text"]["reachable"] = False
|
||||
stack.cfg.hold_restart_s = 0.05
|
||||
r = await pub.post("/v1/chat/completions", json={"model": "text"})
|
||||
assert r.status_code == 503
|
||||
assert r.headers["retry-after"] == "600"
|
||||
err = r.json()["error"]
|
||||
assert err["sleep_depth"] == "restarting"
|
||||
assert err["estimated_wake_seconds"] == 600
|
||||
|
||||
|
||||
async def test_unknown_depth_is_conservative_offloaded(pub, backend, stack):
|
||||
"""Router restarted / slept behind our back: depth unknown -> offloaded."""
|
||||
backend.set_sleeping("ocr", True, level=1) # actually only level 1 asleep
|
||||
stack.cfg.hold_offload_s = 0.05
|
||||
backend.services["vllm-ocr"]["wake_delay"] = 0.3
|
||||
r = await pub.post("/v1/chat/completions", json={"model": "ocr"})
|
||||
assert r.status_code == 503
|
||||
err = r.json()["error"]
|
||||
# conservative: worst-case depth, longest client wait
|
||||
assert err["sleep_depth"] == "offloaded"
|
||||
assert r.headers["retry-after"] == "60"
|
||||
|
||||
|
||||
async def test_wake_sequence_retried_once_then_503(pub, backend, stack):
|
||||
backend.set_sleeping("text", True)
|
||||
backend.services["vllm-text"]["wake_fails"] = 2
|
||||
r = await pub.post("/v1/chat/completions", json={"model": "text"})
|
||||
assert r.status_code == 503
|
||||
assert r.headers["retry-after"] == "60"
|
||||
err = r.json()["error"]
|
||||
assert err["code"] == "model_waking"
|
||||
assert err["sleep_depth"] == "offloaded"
|
||||
# exactly one retry (plan 6.2.1)
|
||||
assert backend.count("POST vllm-text/wake_up") == 2
|
||||
# never proxied to a half-awake backend
|
||||
assert backend.count("POST vllm-text/v1/") == 0
|
||||
|
||||
|
||||
async def test_never_proxy_before_health_ok(pub, backend):
|
||||
"""The /v1 call must happen after the wake sequence, never before it.
|
||||
|
||||
Requests admitted between wake_up and reload_weights return 200 + garbage,
|
||||
so the whole sequence has to finish first.
|
||||
"""
|
||||
backend.set_sleeping("text", True)
|
||||
backend.services["vllm-text"]["wake_delay"] = 0.1
|
||||
r = await pub.post("/v1/chat/completions", json={"model": "text"})
|
||||
assert r.status_code == 200
|
||||
svc = backend.services["vllm-text"]
|
||||
assert svc["wake_up_calls"] == 1
|
||||
assert svc["reload_calls"] == 1
|
||||
assert svc["reset_calls"] == 1
|
||||
calls = backend.calls
|
||||
idx = {needle: backend.last(needle) for needle in (
|
||||
"POST vllm-text/wake_up",
|
||||
"POST vllm-text/collective_rpc",
|
||||
"POST vllm-text/reset_prefix_cache",
|
||||
"POST vllm-text/v1/chat/completions",
|
||||
)}
|
||||
assert idx["POST vllm-text/v1/chat/completions"] > idx["POST vllm-text/reset_prefix_cache"]
|
||||
assert idx["POST vllm-text/reset_prefix_cache"] > idx["POST vllm-text/collective_rpc"]
|
||||
assert idx["POST vllm-text/collective_rpc"] > idx["POST vllm-text/wake_up"]
|
||||
assert calls[-1] == "POST vllm-text/v1/chat/completions"
|
||||
|
||||
|
||||
async def test_level1_wake_uses_the_fast_path(pub, backend, stack):
|
||||
"""From level-1 sleep, /wake_up ALONE is enough (calibration 2026-08-17):
|
||||
bit-identical output at temp 0, and ~20s cheaper than the reload sequence
|
||||
(23.4s -> 2.5-3.8s on the text model)."""
|
||||
assert (await stack.manager.sleep_service("text", 1))["ok"] is True
|
||||
assert stack.manager.services["text"].depth == "sleeping"
|
||||
backend.calls.clear()
|
||||
svc = backend.services["vllm-text"]
|
||||
for field in ("wake_up_calls", "reload_calls", "reset_calls"):
|
||||
svc[field] = 0
|
||||
|
||||
r = await pub.post("/v1/chat/completions", json={"model": "text"})
|
||||
assert r.status_code == 200
|
||||
assert svc["wake_up_calls"] == 1 # wake_up only ...
|
||||
assert svc["reload_calls"] == 0 # ... no reload_weights ...
|
||||
assert svc["reset_calls"] == 0 # ... and no prefix-cache reset
|
||||
assert backend.count("POST vllm-text/v1/chat/completions") == 1
|
||||
|
||||
|
||||
async def test_unknown_depth_uses_the_full_sequence(pub, backend, stack):
|
||||
"""Router restarted / slept out of band: unknown depth -> conservative
|
||||
level-2 treatment, full sequence."""
|
||||
backend.set_sleeping("text", True) # router depth stays unknown
|
||||
r = await pub.post("/v1/chat/completions", json={"model": "text"})
|
||||
assert r.status_code == 200
|
||||
svc = backend.services["vllm-text"]
|
||||
assert (svc["wake_up_calls"], svc["reload_calls"], svc["reset_calls"]) == (1, 1, 1)
|
||||
|
||||
|
||||
async def test_readiness_is_gated_on_is_sleeping_not_health(pub, backend, stack):
|
||||
"""/health answers 200 on a sleeping backend, so /is_sleeping is the gate;
|
||||
nothing is proxied while the backend still reports is_sleeping=true (a
|
||||
request sent to a sleeping backend hangs instead of erroring)."""
|
||||
assert (await stack.manager.sleep_service("text", 2))["ok"] is True
|
||||
backend.calls.clear()
|
||||
backend.services["vllm-text"]["hold_sleeping_polls"] = 3 # wake "in flight"
|
||||
|
||||
r = await pub.post("/v1/chat/completions", json={"model": "text"})
|
||||
assert r.status_code == 200
|
||||
calls = backend.calls
|
||||
probes = [i for i, c in enumerate(calls) if c == "GET vllm-text/is_sleeping"]
|
||||
assert len(probes) >= 4 # kept polling until it flipped
|
||||
assert calls[-1] == "POST vllm-text/v1/chat/completions" # proxied last
|
||||
assert backend.services["vllm-text"]["api_calls"] == 1
|
||||
|
||||
|
||||
async def test_backend_dies_between_health_and_proxy(pub, backend, stack):
|
||||
"""Transport error mid-proxy -> service marked restarting -> depth 503."""
|
||||
text = backend.services["vllm-text"]
|
||||
|
||||
original = backend.handle_async_request
|
||||
|
||||
async def flaky(request: httpx.Request) -> httpx.Response:
|
||||
if request.url.path.startswith("/v1/"):
|
||||
raise httpx.ConnectError("backend gone", request=request)
|
||||
return await original(request)
|
||||
|
||||
backend.handle_async_request = flaky
|
||||
stack.cfg.hold_restart_s = 0.05
|
||||
r = await pub.post("/v1/chat/completions", json={"model": "text"})
|
||||
assert r.status_code == 503
|
||||
assert r.headers["retry-after"] == "600"
|
||||
assert r.json()["error"]["sleep_depth"] == "restarting"
|
||||
assert text["wake_up_calls"] == 0
|
||||
|
||||
|
||||
async def test_wake_after_level2_runs_full_sequence(pub, backend):
|
||||
backend.set_sleeping("embed", True, level=2)
|
||||
r = await pub.post("/v1/embeddings", json={"input": "hi"})
|
||||
assert r.status_code == 200
|
||||
svc = backend.services["vllm-embed"]
|
||||
assert svc["wake_up_calls"] == 1
|
||||
assert svc["reload_calls"] == 1 # reload_weights is mandatory after L2
|
||||
assert svc["reset_calls"] == 1 # prefix cache reset too
|
||||
assert svc["api_calls"] == 1
|
||||
|
||||
|
||||
async def test_error_body_is_openai_shaped(pub, backend, stack):
|
||||
await _sleep_at_level(stack, 1)
|
||||
stack.cfg.hold_sleep_s = 0.01
|
||||
backend.services["vllm-text"]["wake_delay"] = 0.2
|
||||
r = await pub.post("/v1/chat/completions", json={"model": "text"})
|
||||
payload = r.json()
|
||||
assert set(payload) == {"error"}
|
||||
assert set(payload["error"]) == {
|
||||
"type", "code", "message", "sleep_depth", "estimated_wake_seconds"
|
||||
}
|
||||
assert r.headers["content-type"].startswith("application/json")
|
||||
|
||||
|
||||
async def test_depth_survives_raw_traversal_requests(stack, backend):
|
||||
"""Traversal requests are rejected before any backend contact."""
|
||||
backend.set_sleeping("text", True)
|
||||
for raw in ("/v1/../sleep", "/v1%2f..%2fsleep", "//sleep", "/v1/../../wake_up"):
|
||||
status, _ = await raw_asgi(stack.public, "POST", raw)
|
||||
assert status == 404, raw
|
||||
assert backend.calls == []
|
||||
195
router/tests/test_idle.py
Normal file
195
router/tests/test_idle.py
Normal file
@@ -0,0 +1,195 @@
|
||||
"""Tiered idle management: thresholds, lock races, active-request guard."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
from config import DEPTH_AWAKE, DEPTH_OFFLOADED, DEPTH_SLEEPING
|
||||
|
||||
|
||||
def _age(svc, seconds: float) -> None:
|
||||
svc.last_activity = time.monotonic() - seconds
|
||||
|
||||
|
||||
async def test_idle_level1_after_threshold(stack, backend):
|
||||
stack.cfg.idle_sleep_min = 15 / 60.0 # 15 s in "minutes"
|
||||
stack.cfg.idle_offload_min = 180 / 60.0
|
||||
text = stack.manager.services["text"]
|
||||
text.depth = DEPTH_AWAKE
|
||||
_age(text, 20)
|
||||
|
||||
await stack.manager.idle_tick()
|
||||
|
||||
assert backend.count("POST vllm-text/sleep") == 1
|
||||
assert "level=1" in backend.calls[backend.last("POST vllm-text/sleep")]
|
||||
assert text.depth == DEPTH_SLEEPING
|
||||
|
||||
|
||||
async def test_idle_escalates_to_level2(stack, backend):
|
||||
stack.cfg.idle_sleep_min = 15 / 60.0
|
||||
stack.cfg.idle_offload_min = 180 / 60.0
|
||||
text = stack.manager.services["text"]
|
||||
text.depth = DEPTH_AWAKE
|
||||
_age(text, 400)
|
||||
|
||||
await stack.manager.idle_tick()
|
||||
assert "level=2" in backend.calls[backend.last("POST vllm-text/sleep")]
|
||||
assert text.depth == DEPTH_OFFLOADED
|
||||
|
||||
|
||||
async def test_level1_service_escalates_to_level2(stack, backend):
|
||||
"""Already napping (level 1) and still idle -> escalate to offload.
|
||||
|
||||
A direct POST /sleep?level=2 on a level-1-sleeping backend is a
|
||||
well-behaved NO-OP that retains the host-RAM copy (calibration
|
||||
2026-08-17), so the escalation must wake into RAM first and then offload.
|
||||
"""
|
||||
stack.cfg.idle_sleep_min = 1 / 60.0
|
||||
stack.cfg.idle_offload_min = 5 / 60.0
|
||||
text = stack.manager.services["text"]
|
||||
assert (await stack.manager.sleep_service("text", 1))["ok"] is True
|
||||
backend.calls.clear()
|
||||
svc = backend.services["vllm-text"]
|
||||
svc["wake_up_calls"] = svc["sleep_calls"] = 0
|
||||
_age(text, 400)
|
||||
|
||||
await stack.manager.idle_tick()
|
||||
assert text.depth == DEPTH_OFFLOADED
|
||||
assert svc["wake_up_calls"] == 1 # wake into RAM ...
|
||||
assert svc["sleep_calls"] == 1 # ... then offload
|
||||
assert backend.last("POST vllm-text/sleep?level=2") > backend.last("POST vllm-text/wake_up")
|
||||
|
||||
|
||||
async def test_offload_from_awake_is_direct(stack, backend):
|
||||
"""Only from depth 'awake' can level 2 be entered directly."""
|
||||
backend.services["vllm-embed"]["wake_up_calls"] = 0
|
||||
result = await stack.manager.sleep_service("embed", 2)
|
||||
assert result["ok"] is True
|
||||
assert backend.services["vllm-embed"]["wake_up_calls"] == 0
|
||||
assert backend.count("POST vllm-embed/sleep?level=2") == 1
|
||||
assert stack.manager.services["embed"].depth == DEPTH_OFFLOADED
|
||||
|
||||
|
||||
async def test_escalation_failure_is_reported(stack, backend):
|
||||
assert (await stack.manager.sleep_service("ocr", 1))["ok"] is True
|
||||
backend.services["vllm-ocr"]["wake_fails"] = 1
|
||||
backend.calls.clear()
|
||||
result = await stack.manager.sleep_service("ocr", 2)
|
||||
assert result["ok"] is False
|
||||
assert result["reason"] == "escalation_failed"
|
||||
assert backend.count("POST vllm-ocr/sleep") == 0
|
||||
assert stack.manager.services["ocr"].depth == DEPTH_SLEEPING
|
||||
|
||||
|
||||
async def test_offloaded_service_is_left_alone(stack, backend):
|
||||
stack.cfg.idle_offload_min = 1 / 60.0
|
||||
text = stack.manager.services["text"]
|
||||
text.depth = DEPTH_OFFLOADED
|
||||
_age(text, 10_000)
|
||||
await stack.manager.idle_tick()
|
||||
assert backend.calls == []
|
||||
|
||||
|
||||
async def test_active_requests_block_idle_sleep(stack, backend):
|
||||
stack.cfg.idle_sleep_min = 1 / 60.0
|
||||
text = stack.manager.services["text"]
|
||||
text.depth = DEPTH_AWAKE
|
||||
_age(text, 600)
|
||||
|
||||
manager = stack.manager
|
||||
manager.begin_request(text) # long generation in flight
|
||||
try:
|
||||
await manager.idle_tick()
|
||||
assert backend.calls == []
|
||||
finally:
|
||||
manager.end_request(text)
|
||||
|
||||
|
||||
async def test_stream_completes_then_idle_can_sleep(stack, backend, pub):
|
||||
"""end_request (background task of the streamed response) re-opens sleep."""
|
||||
stack.cfg.idle_sleep_min = 1 / 60.0
|
||||
text = stack.manager.services["text"]
|
||||
r = await pub.post("/v1/chat/completions", json={"model": "text"})
|
||||
assert r.status_code == 200
|
||||
assert text.active_requests == 0 # released after the body drained
|
||||
|
||||
_age(text, 600)
|
||||
await stack.manager.idle_tick()
|
||||
assert backend.count("POST vllm-text/sleep") == 1
|
||||
|
||||
|
||||
async def test_race_last_activity_refreshed_under_lock(stack, backend):
|
||||
"""A request landed between the threshold check and the locked re-check."""
|
||||
stack.cfg.idle_sleep_min = 1 / 60.0 # threshold = 60 s
|
||||
text = stack.manager.services["text"]
|
||||
text.depth = DEPTH_AWAKE
|
||||
_age(text, 600) # what the idle scan saw
|
||||
|
||||
# ... and then a request arrived, refreshing last_activity *now*
|
||||
text.last_activity = time.monotonic()
|
||||
result = await stack.manager.sleep_service("text", 1, reason="idle", min_idle_s=60.0)
|
||||
assert result["ok"] is False
|
||||
assert result["reason"] == "activity_resumed"
|
||||
assert backend.count("POST vllm-text/sleep") == 0
|
||||
|
||||
|
||||
async def test_race_active_request_seen_under_lock(stack, backend):
|
||||
text = stack.manager.services["text"]
|
||||
text.depth = DEPTH_AWAKE
|
||||
_age(text, 600)
|
||||
text.active_requests = 2 # arrived while we were scanning
|
||||
result = await stack.manager.sleep_service("text", 1, reason="idle", min_idle_s=60.0)
|
||||
assert result["ok"] is False
|
||||
assert result["reason"] == "active_requests"
|
||||
assert backend.calls == []
|
||||
text.active_requests = 0
|
||||
|
||||
|
||||
async def test_request_at_the_moment_the_timer_expires(pub, backend, stack):
|
||||
"""E2E case 11 in miniature: the request wins, no sleep mid-flight."""
|
||||
stack.cfg.idle_sleep_min = 1 / 60.0
|
||||
stack.cfg.idle_offload_min = 5 / 60.0
|
||||
text = stack.manager.services["text"]
|
||||
text.depth = DEPTH_AWAKE
|
||||
_age(text, 600)
|
||||
|
||||
async def tick_soon():
|
||||
await asyncio.sleep(0) # run after the request started
|
||||
return await stack.manager.idle_tick()
|
||||
|
||||
tick, response = await asyncio.gather(tick_soon(), pub.post(
|
||||
"/v1/chat/completions", json={"model": "text"}))
|
||||
assert response.status_code == 200
|
||||
assert backend.count("POST vllm-text/sleep") == 0
|
||||
|
||||
|
||||
async def test_admin_sleep_refuses_when_active(adm, stack, backend):
|
||||
text = stack.manager.services["text"]
|
||||
text.depth = DEPTH_AWAKE
|
||||
stack.manager.begin_request(text)
|
||||
r = await adm.post("/admin/sleep/text?level=1")
|
||||
assert r.status_code == 409
|
||||
assert r.json()["reason"] == "active_requests"
|
||||
assert backend.count("/sleep") == 0
|
||||
stack.manager.end_request(text)
|
||||
|
||||
|
||||
async def test_idle_loop_runs_in_background_when_started(stack, backend):
|
||||
stack.cfg.idle_enabled = True
|
||||
stack.cfg.idle_poll_s = 0.01
|
||||
stack.cfg.idle_sleep_min = 1 / 60.0 # 60 s
|
||||
stack.cfg.idle_offload_min = 1000.0 # far away: expect the level-1 tier
|
||||
text = stack.manager.services["text"]
|
||||
text.depth = DEPTH_AWAKE
|
||||
_age(text, 600)
|
||||
stack.manager.start()
|
||||
try:
|
||||
for _ in range(100):
|
||||
if backend.count("POST vllm-text/sleep"):
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
finally:
|
||||
await stack.manager.stop()
|
||||
assert backend.count("POST vllm-text/sleep") == 1
|
||||
assert "level=1" in backend.calls[backend.last("POST vllm-text/sleep")]
|
||||
160
router/tests/test_model_resolution.py
Normal file
160
router/tests/test_model_resolution.py
Normal file
@@ -0,0 +1,160 @@
|
||||
"""Model resolution: JSON / multipart / defaults / embeddings / unknown."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
TEXT = "Qwen3.6-35B-A3B-FP8"
|
||||
OCR = "OvisOCR2"
|
||||
EMBED = "Qwen3-Embedding-8B"
|
||||
|
||||
|
||||
async def test_json_model_exact(pub, backend):
|
||||
r = await pub.post("/v1/chat/completions", json={"model": OCR, "messages": []})
|
||||
assert r.status_code == 200
|
||||
assert r.json()["service"] == "vllm-ocr"
|
||||
|
||||
|
||||
async def test_json_model_case_insensitive(pub):
|
||||
r = await pub.post("/v1/chat/completions", json={"model": "qwen3.6-35b-a3b-fp8"})
|
||||
assert r.status_code == 200
|
||||
assert r.json()["service"] == "vllm-text"
|
||||
|
||||
|
||||
async def test_alias_matches(pub):
|
||||
r = await pub.post("/v1/chat/completions", json={"model": "ocr"})
|
||||
assert r.status_code == 200
|
||||
assert r.json()["service"] == "vllm-ocr"
|
||||
r = await pub.post("/v1/chat/completions", json={"model": "Embedding"})
|
||||
assert r.json()["service"] == "vllm-embed"
|
||||
|
||||
|
||||
async def test_missing_model_defaults_to_text_on_chat(pub):
|
||||
r = await pub.post("/v1/chat/completions", json={"messages": [{"role": "user", "content": "hi"}]})
|
||||
assert r.status_code == 200
|
||||
assert r.json()["service"] == "vllm-text"
|
||||
|
||||
|
||||
async def test_missing_model_defaults_to_text_on_completions(pub):
|
||||
r = await pub.post("/v1/completions", json={"prompt": "hi"})
|
||||
assert r.status_code == 200
|
||||
assert r.json()["service"] == "vllm-text"
|
||||
|
||||
|
||||
async def test_missing_model_is_400_elsewhere(pub, backend):
|
||||
r = await pub.post("/v1/rerank", json={"query": "hi"})
|
||||
assert r.status_code == 400
|
||||
assert r.json()["error"]["code"] == "missing_model"
|
||||
assert backend.count("/v1/rerank") == 0
|
||||
|
||||
|
||||
async def test_invalid_json_is_400(pub, backend):
|
||||
r = await pub.post("/v1/chat/completions",
|
||||
content=b"{not json",
|
||||
headers={"content-type": "application/json"})
|
||||
assert r.status_code == 400
|
||||
assert backend.count("/v1/") == 0
|
||||
|
||||
|
||||
async def test_embeddings_always_routes_to_embed(pub):
|
||||
r = await pub.post("/v1/embeddings", json={"input": "hello"})
|
||||
assert r.status_code == 200
|
||||
assert r.json()["service"] == "vllm-embed"
|
||||
# ... even when the body names a different model
|
||||
r = await pub.post("/v1/embeddings", json={"input": "hello", "model": TEXT})
|
||||
assert r.status_code == 200
|
||||
assert r.json()["service"] == "vllm-embed"
|
||||
|
||||
|
||||
async def test_unknown_model_404_and_no_wake(pub, backend):
|
||||
backend.set_sleeping("ocr", True)
|
||||
r = await pub.post("/v1/chat/completions", json={"model": "gpt-4o"})
|
||||
assert r.status_code == 404
|
||||
body = r.json()["error"]
|
||||
assert body["code"] == "model_not_found"
|
||||
assert body["type"] == "invalid_request_error"
|
||||
assert body["param"] == "model"
|
||||
# no probe, no wake, no proxy
|
||||
assert backend.calls == []
|
||||
|
||||
|
||||
async def test_multipart_with_model_field(pub):
|
||||
r = await pub.post(
|
||||
"/v1/chat/completions",
|
||||
data={"model": OCR},
|
||||
files={"image": ("page.png", b"PNGDATA" * 64, "image/png")},
|
||||
)
|
||||
assert r.status_code == 200
|
||||
echo = r.json()
|
||||
assert echo["service"] == "vllm-ocr"
|
||||
# raw body forwarded unchanged: the file bytes and the boundary survive
|
||||
assert "PNGDATA" * 64 in echo["echo_body"]
|
||||
assert echo["echo_content_type"].startswith("multipart/form-data; boundary=")
|
||||
|
||||
|
||||
async def test_multipart_without_model_defaults_to_ocr(pub):
|
||||
r = await pub.post(
|
||||
"/v1/chat/completions",
|
||||
files={"image": ("page.png", b"PNGDATA", "image/png")},
|
||||
)
|
||||
assert r.status_code == 200
|
||||
assert r.json()["service"] == "vllm-ocr"
|
||||
|
||||
|
||||
async def test_multipart_unknown_model_404(pub, backend):
|
||||
r = await pub.post(
|
||||
"/v1/chat/completions",
|
||||
data={"model": "nope"},
|
||||
files={"image": ("page.png", b"x", "image/png")},
|
||||
)
|
||||
assert r.status_code == 404
|
||||
assert r.json()["error"]["code"] == "model_not_found"
|
||||
assert backend.calls == []
|
||||
|
||||
|
||||
async def test_multipart_model_extraction_ignores_file_parts(pub):
|
||||
# a file part literally named "model" must not be read as the model field
|
||||
r = await pub.post(
|
||||
"/v1/chat/completions",
|
||||
data={"prompt": "hi"},
|
||||
files={"model": ("fake.json", b"NOT-A-MODEL-NAME", "application/octet-stream")},
|
||||
)
|
||||
# no usable model field -> multipart default on the chat path is OCR
|
||||
assert r.status_code == 200
|
||||
assert r.json()["service"] == "vllm-ocr"
|
||||
assert "NOT-A-MODEL-NAME" in r.json()["echo_body"]
|
||||
|
||||
|
||||
async def test_model_in_path_for_models_detail(pub):
|
||||
r = await pub.get(f"/v1/models/{OCR}")
|
||||
assert r.status_code == 200
|
||||
assert r.json()["service"] == "vllm-ocr"
|
||||
r = await pub.get("/v1/models/does-not-exist")
|
||||
assert r.status_code == 404
|
||||
assert r.json()["error"]["code"] == "model_not_found"
|
||||
|
||||
|
||||
async def test_body_and_headers_forwarded(pub):
|
||||
payload = {"model": TEXT, "messages": [{"role": "user", "content": "hi"}], "stream": False}
|
||||
r = await pub.post("/v1/chat/completions", json=payload,
|
||||
headers={"Authorization": "Bearer x", "x-custom": "1"})
|
||||
assert r.status_code == 200
|
||||
echo = r.json()
|
||||
assert echo["echo_path"] == "/v1/chat/completions"
|
||||
assert echo["echo_method"] == "POST"
|
||||
assert json.loads(echo["echo_body"]) == payload
|
||||
|
||||
|
||||
async def test_query_string_forwarded(pub):
|
||||
r = await pub.post("/v1/chat/completions?foo=bar&baz=1", json={"model": TEXT})
|
||||
assert r.status_code == 200
|
||||
assert r.json()["echo_query"] == "foo=bar&baz=1"
|
||||
|
||||
|
||||
async def test_streaming_passthrough(pub, backend):
|
||||
backend.stream_chunks = [b"data: {\"a\":1}\n\n", b"data: {\"a\":2}\n\n", b"data: [DONE]\n\n"]
|
||||
r = await pub.post("/v1/chat/completions", json={"model": TEXT, "stream": True})
|
||||
assert r.status_code == 200
|
||||
assert r.headers["content-type"].startswith("text/event-stream")
|
||||
assert r.text == "".join(chunk.decode() for chunk in backend.stream_chunks)
|
||||
backend.stream_chunks = []
|
||||
146
router/tests/test_paths.py
Normal file
146
router/tests/test_paths.py
Normal file
@@ -0,0 +1,146 @@
|
||||
"""Allowlist: only /v1/*, /health (and /metrics when enabled) are public."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
|
||||
from app import build_apps
|
||||
from conftest import raw_asgi
|
||||
|
||||
DEV_ENDPOINTS = [
|
||||
"/sleep",
|
||||
"/wake_up",
|
||||
"/collective_rpc",
|
||||
"/reset_prefix_cache",
|
||||
"/is_sleeping",
|
||||
]
|
||||
ADMIN_PATHS = ["/admin/status", "/admin/wake/text", "/admin/sleep/text?level=1"]
|
||||
|
||||
|
||||
async def test_dev_endpoints_404_on_public(pub, backend):
|
||||
for path in DEV_ENDPOINTS:
|
||||
r = await pub.post(path)
|
||||
assert r.status_code == 404, path
|
||||
assert r.json()["error"]["code"] == "not_found"
|
||||
for path in DEV_ENDPOINTS:
|
||||
r = await pub.get(path)
|
||||
assert r.status_code == 404, path
|
||||
assert backend.calls == []
|
||||
|
||||
|
||||
async def test_admin_not_on_public(pub, backend):
|
||||
for path in ADMIN_PATHS:
|
||||
r = await pub.request("POST", path)
|
||||
assert r.status_code == 404, path
|
||||
r = await pub.get(path)
|
||||
assert r.status_code == 404, path
|
||||
assert backend.calls == []
|
||||
|
||||
|
||||
async def test_docs_and_openapi_404(pub):
|
||||
for path in ("/docs", "/redoc", "/openapi.json", "/"):
|
||||
r = await pub.get(path)
|
||||
assert r.status_code == 404, path
|
||||
|
||||
|
||||
async def test_dotted_paths_404(pub):
|
||||
for path in ("/v1/../sleep", "/v1/..%2fsleep", "/health/../sleep", "/./sleep"):
|
||||
r = await pub.post(path)
|
||||
assert r.status_code == 404, path
|
||||
|
||||
|
||||
async def test_raw_traversal_404(stack, backend):
|
||||
"""Percent-encoded traversal that a normal client would normalise."""
|
||||
cases = [
|
||||
("POST", "/v1%2f..%2fsleep"),
|
||||
("POST", "/v1%2F..%2Fsleep"),
|
||||
("GET", "/v1/%2e%2e/wake_up"),
|
||||
("POST", "/v1/%2e%2e%2fcollective_rpc"),
|
||||
("POST", "/v1//../sleep"),
|
||||
("GET", "/v1/models/../../is_sleeping"),
|
||||
]
|
||||
for method, raw in cases:
|
||||
status, body = await raw_asgi(stack.public, method, raw,
|
||||
headers=[("content-type", "application/json")],
|
||||
body=b"{}")
|
||||
assert status == 404, (method, raw, body)
|
||||
assert backend.calls == []
|
||||
|
||||
|
||||
async def test_normalised_traversal_still_reaches_v1(stack, backend):
|
||||
"""A traversal that lands inside /v1 must still work (not over-block)."""
|
||||
status, body = await raw_asgi(
|
||||
stack.public, "POST", "/v1/../v1/chat/completions",
|
||||
headers=[("content-type", "application/json")],
|
||||
body=b'{"model": "text"}',
|
||||
)
|
||||
assert status == 200, body
|
||||
assert backend.count("POST vllm-text/v1/chat/completions") == 1
|
||||
|
||||
|
||||
async def test_public_health(pub, backend):
|
||||
r = await pub.get("/health")
|
||||
assert r.status_code == 200
|
||||
body = r.json()
|
||||
assert body["status"] == "ok"
|
||||
assert set(body["services"]) == {"text", "ocr", "embed"}
|
||||
# liveness must not depend on backend state
|
||||
backend.services["vllm-text"]["reachable"] = False
|
||||
r = await pub.get("/health")
|
||||
assert r.status_code == 200
|
||||
|
||||
|
||||
async def test_public_models_lists_all_three(pub, backend):
|
||||
r = await pub.get("/v1/models")
|
||||
assert r.status_code == 200
|
||||
body = r.json()
|
||||
assert body["object"] == "list"
|
||||
assert [m["id"] for m in body["data"]] == [
|
||||
"Qwen3.6-35B-A3B-FP8", "OvisOCR2", "Qwen3-Embedding-8B",
|
||||
]
|
||||
# answered by the router itself, no backend contact
|
||||
assert backend.calls == []
|
||||
|
||||
|
||||
async def test_metrics_disabled_by_default(pub):
|
||||
assert (await pub.get("/metrics")).status_code == 404
|
||||
|
||||
|
||||
async def test_metrics_enabled_when_configured(stack, pub):
|
||||
stack.cfg.metrics_enabled = True
|
||||
r = await pub.get("/metrics")
|
||||
assert r.status_code == 200
|
||||
assert "vllm_router_service_depth" in r.text
|
||||
|
||||
|
||||
async def test_unknown_http_method_on_v1(stack):
|
||||
status, _ = await raw_asgi(stack.public, "PROPFIND", "/v1/chat/completions")
|
||||
assert status in (404, 405)
|
||||
|
||||
|
||||
async def test_admin_listener_is_a_separate_app(stack):
|
||||
"""The admin app answers /admin/*; the public app must not."""
|
||||
status, body = await raw_asgi(stack.admin, "GET", "/admin/status")
|
||||
assert status == 200
|
||||
status, _ = await raw_asgi(stack.public, "GET", "/admin/status")
|
||||
assert status == 404
|
||||
|
||||
|
||||
async def test_request_body_cap(cfg, backend):
|
||||
"""ROUTER_MAX_BODY_BYTES must be wired as middleware (413, not OOM)."""
|
||||
cfg.max_body_bytes = 64
|
||||
public_app, _admin, manager = build_apps(cfg, transport=backend)
|
||||
ok_body = b'{"model": "text", "messages": []}'
|
||||
assert len(ok_body) < 64
|
||||
try:
|
||||
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=public_app),
|
||||
base_url="http://router.test") as client:
|
||||
r = await client.post("/v1/chat/completions", content=ok_body,
|
||||
headers={"content-type": "application/json"})
|
||||
assert r.status_code == 200
|
||||
r = await client.post("/v1/chat/completions", content=ok_body + b" " * 512,
|
||||
headers={"content-type": "application/json"})
|
||||
assert r.status_code == 413
|
||||
finally:
|
||||
await manager.proxy_client.aclose()
|
||||
await manager.ctrl_client.aclose()
|
||||
101
router/tests/test_shutdown_and_cli.py
Normal file
101
router/tests/test_shutdown_and_cli.py
Normal file
@@ -0,0 +1,101 @@
|
||||
"""Shutdown behaviour (E2E case 12) and the vllmctl CLI glue."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
|
||||
import pytest
|
||||
|
||||
from conftest import raw_asgi
|
||||
from routing import ShutdownGuardMiddleware
|
||||
|
||||
|
||||
async def test_shutdown_guard_returns_503_when_not_started(stack):
|
||||
"""uvicorn cancels in-flight tasks after timeout_graceful_shutdown; that
|
||||
CancelledError must become a depth-aware 503, not uvicorn's bare 500."""
|
||||
stack.manager.begin_shutdown()
|
||||
|
||||
async def app(scope, receive, send):
|
||||
raise asyncio.CancelledError()
|
||||
|
||||
guard = ShutdownGuardMiddleware(app, stack.manager)
|
||||
status, body = await raw_asgi(guard, "POST", "/v1/chat/completions")
|
||||
assert status == 503
|
||||
payload = json.loads(body)
|
||||
assert payload["error"]["code"] == "router_shutting_down"
|
||||
assert payload["error"]["retry_after_seconds"] == 5
|
||||
|
||||
|
||||
async def test_shutdown_guard_passes_through_when_not_shutting_down(stack):
|
||||
class Boom(RuntimeError):
|
||||
pass
|
||||
|
||||
async def app(scope, receive, send):
|
||||
raise Boom("nope")
|
||||
|
||||
guard = ShutdownGuardMiddleware(app, stack.manager)
|
||||
with pytest.raises(Boom):
|
||||
await raw_asgi(guard, "POST", "/v1/chat/completions")
|
||||
|
||||
|
||||
async def test_shutdown_guard_cannot_unsend_a_started_response(stack):
|
||||
"""A response that already started streaming is closed, not replaced."""
|
||||
stack.manager.begin_shutdown()
|
||||
|
||||
async def app(scope, receive, send):
|
||||
await send({"type": "http.response.start", "status": 200, "headers": []})
|
||||
await send({"type": "http.response.body", "body": b"partial"})
|
||||
raise asyncio.CancelledError()
|
||||
|
||||
guard = ShutdownGuardMiddleware(app, stack.manager)
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await raw_asgi(guard, "POST", "/v1/chat/completions")
|
||||
|
||||
|
||||
async def test_shutdown_guard_forwards_normal_requests(stack):
|
||||
async def app(scope, receive, send):
|
||||
await send({"type": "http.response.start", "status": 204, "headers": []})
|
||||
await send({"type": "http.response.body", "body": b""})
|
||||
|
||||
guard = ShutdownGuardMiddleware(app, stack.manager)
|
||||
status, _ = await raw_asgi(guard, "GET", "/health")
|
||||
assert status == 204
|
||||
|
||||
|
||||
def test_vllmctl_service_name_mapping():
|
||||
"""`logs router` / `restart router` must map to the compose service
|
||||
`router`; the model services map to vllm-*."""
|
||||
script = (
|
||||
"source /data/home/renbaibing/vllm/vllmctl >/dev/null 2>&1; "
|
||||
'for n in router vllm-router Router text vllm-text TEXT ocr OvisOCR2 '
|
||||
"embed Qwen3-Embedding-8B bogus; do "
|
||||
'printf "%s:%s\\n" "$n" "$(docker_service_for "$n" || echo ERR)"; done'
|
||||
)
|
||||
out = subprocess.run(["bash", "-c", script], capture_output=True, text=True,
|
||||
timeout=30)
|
||||
assert out.returncode == 0, out.stderr
|
||||
mapping = dict(line.split(":", 1) for line in out.stdout.strip().splitlines())
|
||||
assert mapping["router"] == "router"
|
||||
assert mapping["vllm-router"] == "router"
|
||||
assert mapping["Router"] == "router"
|
||||
assert mapping["text"] == "vllm-text"
|
||||
assert mapping["vllm-text"] == "vllm-text"
|
||||
assert mapping["TEXT"] == "vllm-text"
|
||||
assert mapping["ocr"] == "vllm-ocr"
|
||||
assert mapping["OvisOCR2"] == "vllm-ocr"
|
||||
assert mapping["embed"] == "vllm-embed"
|
||||
assert mapping["Qwen3-Embedding-8B"] == "vllm-embed"
|
||||
assert mapping["bogus"] == "ERR"
|
||||
|
||||
|
||||
def test_vllmctl_has_no_removed_machinery():
|
||||
source = open("/data/home/renbaibing/vllm/vllmctl", encoding="utf-8").read()
|
||||
assert ".runas.py" not in source
|
||||
assert ".user.env" not in source
|
||||
assert "idle-watch on" not in source
|
||||
assert "cmd_idle_watch" not in source
|
||||
assert "admin/wake/" in source and "admin/sleep/" in source
|
||||
assert os.access("/data/home/renbaibing/vllm/vllmctl", os.X_OK)
|
||||
95
router/tests/test_singleflight.py
Normal file
95
router/tests/test_singleflight.py
Normal file
@@ -0,0 +1,95 @@
|
||||
"""Single-flight wake: N concurrent requests -> exactly one wake sequence."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
from config import DEPTH_AWAKE
|
||||
|
||||
|
||||
async def test_ten_concurrent_requests_trigger_one_wake(pub, backend):
|
||||
backend.set_sleeping("text", True, level=2)
|
||||
backend.services["vllm-text"]["wake_delay"] = 0.15
|
||||
|
||||
async def one(i: int):
|
||||
r = await pub.post("/v1/chat/completions", json={"model": "text", "seed": i})
|
||||
assert r.status_code == 200, r.text
|
||||
return r.json()
|
||||
|
||||
results = await asyncio.gather(*(one(i) for i in range(10)))
|
||||
assert all(r["service"] == "vllm-text" for r in results)
|
||||
|
||||
svc = backend.services["vllm-text"]
|
||||
assert svc["wake_up_calls"] == 1, backend.calls
|
||||
assert svc["reload_calls"] == 1
|
||||
assert svc["api_calls"] == 10
|
||||
|
||||
|
||||
async def test_hold_timeout_does_not_cancel_the_wake(pub, backend, stack):
|
||||
"""The first caller times out; the wake keeps going and finishes for the
|
||||
next caller (no second wake sequence, no half-awake proxy)."""
|
||||
backend.set_sleeping("text", True)
|
||||
backend.services["vllm-text"]["wake_delay"] = 0.3
|
||||
stack.cfg.hold_offload_s = 0.05
|
||||
|
||||
r = await pub.post("/v1/chat/completions", json={"model": "text"})
|
||||
assert r.status_code == 503
|
||||
assert r.json()["error"]["sleep_depth"] == "offloaded"
|
||||
|
||||
await asyncio.sleep(0.5) # let the shielded wake task finish
|
||||
svc_state = stack.manager.services["text"]
|
||||
assert svc_state.depth == DEPTH_AWAKE
|
||||
assert backend.services["vllm-text"]["wake_up_calls"] == 1
|
||||
|
||||
r = await pub.post("/v1/chat/completions", json={"model": "text"})
|
||||
assert r.status_code == 200
|
||||
assert backend.services["vllm-text"]["wake_up_calls"] == 1 # cached awake
|
||||
|
||||
|
||||
async def test_two_services_wake_concurrently_without_cross_talk(pub, backend):
|
||||
backend.set_sleeping("text", True)
|
||||
backend.set_sleeping("embed", True)
|
||||
backend.services["vllm-text"]["wake_delay"] = 0.2
|
||||
backend.services["vllm-embed"]["wake_delay"] = 0.05
|
||||
|
||||
chat = pub.post("/v1/chat/completions", json={"model": "text"})
|
||||
embed = pub.post("/v1/embeddings", json={"input": "hi"})
|
||||
r1, r2 = await asyncio.gather(chat, embed)
|
||||
assert r1.status_code == 200 and r1.json()["service"] == "vllm-text"
|
||||
assert r2.status_code == 200 and r2.json()["service"] == "vllm-embed"
|
||||
assert backend.services["vllm-text"]["wake_up_calls"] == 1
|
||||
assert backend.services["vllm-embed"]["wake_up_calls"] == 1
|
||||
|
||||
|
||||
async def test_wake_state_is_cached_between_requests(pub, backend, stack):
|
||||
stack.cfg.state_cache_ttl = 30.0
|
||||
backend.set_sleeping("ocr", True)
|
||||
assert (await pub.post("/v1/chat/completions", json={"model": "ocr"})).status_code == 200
|
||||
probes_after_first = backend.count("GET vllm-ocr/is_sleeping")
|
||||
assert probes_after_first >= 1
|
||||
for _ in range(5):
|
||||
assert (await pub.post("/v1/chat/completions", json={"model": "ocr"})).status_code == 200
|
||||
# no extra probe per request while the awake state is cached
|
||||
assert backend.count("GET vllm-ocr/is_sleeping") == probes_after_first
|
||||
|
||||
|
||||
async def test_expired_cache_reprobes(pub, backend, stack):
|
||||
stack.cfg.state_cache_ttl = 0.0
|
||||
for _ in range(3):
|
||||
assert (await pub.post("/v1/chat/completions", json={"model": "ocr"})).status_code == 200
|
||||
assert backend.count("GET vllm-ocr/is_sleeping") >= 3
|
||||
|
||||
|
||||
async def test_admin_and_request_path_share_one_lock(adm, pub, backend):
|
||||
"""A wake driven from the admin port is joined by the public request path."""
|
||||
backend.set_sleeping("text", True)
|
||||
backend.services["vllm-text"]["wake_delay"] = 0.2
|
||||
|
||||
admin_wake = asyncio.create_task(adm.post("/admin/wake/text"))
|
||||
await asyncio.sleep(0.05)
|
||||
r = await pub.post("/v1/chat/completions", json={"model": "text"})
|
||||
admin_result = await admin_wake
|
||||
|
||||
assert r.status_code == 200
|
||||
assert admin_result.status_code == 200
|
||||
assert backend.services["vllm-text"]["wake_up_calls"] == 1
|
||||
296
router/tests/test_wake_recovery.py
Normal file
296
router/tests/test_wake_recovery.py
Normal file
@@ -0,0 +1,296 @@
|
||||
"""Router-restart-mid-wake recovery (plan 6.4, E2E case 12).
|
||||
|
||||
`is_sleeping=false` does NOT mean "reload done": if the previous router died
|
||||
between POST /wake_up and collective_rpc reload_weights, the backend happily
|
||||
reports awake+healthy while serving garbage. The persisted wake-intent file
|
||||
is what lets the new router tell the two apart.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
|
||||
import httpx
|
||||
|
||||
from app import build_apps
|
||||
from config import DEPTH_OFFLOADED, clone, load_config
|
||||
from conftest import raw_asgi
|
||||
from services import ServiceManager
|
||||
|
||||
|
||||
def write_state(path, **services) -> None:
|
||||
payload = {"version": 1, "updated": 0.0, "services": {
|
||||
key: {"depth": value.get("depth"),
|
||||
"wake_in_progress": value.get("wake_in_progress", False),
|
||||
"level": value.get("level", 0),
|
||||
"pending_reload": value.get("pending_reload", False)}
|
||||
for key, value in services.items()
|
||||
}}
|
||||
with open(path, "w", encoding="utf-8") as fh:
|
||||
json.dump(payload, fh)
|
||||
|
||||
|
||||
def read_state(path) -> dict:
|
||||
with open(path, encoding="utf-8") as fh:
|
||||
return json.load(fh)["services"]
|
||||
|
||||
|
||||
async def _make(cfg, backend, state_file):
|
||||
public_app, admin_app, manager = build_apps(cfg, transport=backend)
|
||||
try:
|
||||
yield public_app, admin_app, manager
|
||||
finally:
|
||||
await manager.proxy_client.aclose()
|
||||
await manager.ctrl_client.aclose()
|
||||
|
||||
|
||||
async def test_startup_completes_interrupted_wake(tmp_path, backend):
|
||||
"""The E2E failure: old router died after wake_up, before reload_weights.
|
||||
The new router must not proxy until the sequence has been completed."""
|
||||
state_file = str(tmp_path / "state.json")
|
||||
write_state(state_file, text={"depth": "offloaded", "wake_in_progress": True,
|
||||
"level": 2})
|
||||
# backend *looks* perfectly healthy and awake -- that is the lie
|
||||
backend.services["vllm-text"]["sleeping"] = False
|
||||
|
||||
cfg = clone(load_config())
|
||||
cfg.state_file = state_file
|
||||
cfg.idle_enabled = False
|
||||
cfg.state_cache_ttl = 0.0
|
||||
async for _public, _admin, manager in _make(cfg, backend, state_file):
|
||||
manager.start()
|
||||
await asyncio.sleep(0.3) # let the recovery task run
|
||||
|
||||
text = manager.services["text"]
|
||||
assert text.pending_reload is False
|
||||
assert text.depth == "awake"
|
||||
svc = backend.services["vllm-text"]
|
||||
# the interrupted sequence was completed, not skipped
|
||||
assert svc["reload_calls"] == 1
|
||||
assert svc["reset_calls"] == 1
|
||||
# and the file no longer claims a wake in flight
|
||||
assert read_state(state_file)["text"]["wake_in_progress"] is False
|
||||
|
||||
|
||||
async def test_interrupted_wake_never_proxies_before_recovery(tmp_path, backend, pub):
|
||||
"""A request arriving in the recovery window must not be proxied on the
|
||||
strength of is_sleeping=false alone."""
|
||||
state_file = str(tmp_path / "state.json")
|
||||
write_state(state_file, ocr={"depth": "offloaded", "wake_in_progress": True, "level": 2})
|
||||
backend.services["vllm-ocr"]["sleeping"] = False # awake-looking
|
||||
backend.services["vllm-ocr"]["api_delay"] = 0.0
|
||||
backend.calls.clear()
|
||||
|
||||
stack_cfg = pub # unused; keeps the fixture ordering simple
|
||||
del stack_cfg
|
||||
cfg = clone(load_config())
|
||||
cfg.state_file = state_file
|
||||
cfg.idle_enabled = False
|
||||
async for public_app, _admin, manager in _make(cfg, backend, state_file):
|
||||
manager.start()
|
||||
# fire the request immediately, before recovery has finished probing
|
||||
task = asyncio.create_task(_post(public_app, "/v1/chat/completions",
|
||||
{"model": "ocr"}))
|
||||
await asyncio.sleep(0.05)
|
||||
ocr = backend.services["vllm-ocr"]
|
||||
assert ocr["reload_calls"] >= 0
|
||||
response = await task
|
||||
assert response.status_code == 200
|
||||
# proxied only after a completed reload sequence
|
||||
assert backend.last("POST vllm-ocr/v1/chat/completions") > \
|
||||
backend.last("POST vllm-ocr/reset_prefix_cache")
|
||||
assert manager.services["ocr"].pending_reload is False
|
||||
|
||||
|
||||
async def _post(app, path, payload) -> httpx.Response:
|
||||
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app),
|
||||
base_url="http://router.test") as client:
|
||||
return await client.post(path, json=payload)
|
||||
|
||||
|
||||
async def test_startup_clears_state_when_unreachable(tmp_path, backend):
|
||||
"""wake_in_progress + unreachable backend = fresh boot (weights are fresh);
|
||||
nothing to complete."""
|
||||
state_file = str(tmp_path / "state.json")
|
||||
write_state(state_file, text={"depth": "offloaded", "wake_in_progress": True, "level": 2})
|
||||
backend.services["vllm-text"]["reachable"] = False
|
||||
|
||||
cfg = clone(load_config())
|
||||
cfg.state_file = state_file
|
||||
cfg.idle_enabled = False
|
||||
async for _public, _admin, manager in _make(cfg, backend, state_file):
|
||||
manager.start()
|
||||
await asyncio.sleep(0.2)
|
||||
text = manager.services["text"]
|
||||
assert text.pending_reload is False
|
||||
assert backend.services["vllm-text"]["reload_calls"] == 0
|
||||
assert read_state(state_file)["text"]["wake_in_progress"] is False
|
||||
|
||||
|
||||
async def test_startup_clears_state_when_still_sleeping(tmp_path, backend):
|
||||
"""wake_in_progress but the backend never woke: the normal request path
|
||||
will run a full wake when traffic arrives."""
|
||||
state_file = str(tmp_path / "state.json")
|
||||
write_state(state_file, embed={"depth": "sleeping", "wake_in_progress": True, "level": 1})
|
||||
backend.set_sleeping("embed", True, level=1)
|
||||
|
||||
cfg = clone(load_config())
|
||||
cfg.state_file = state_file
|
||||
cfg.idle_enabled = False
|
||||
async for _public, _admin, manager in _make(cfg, backend, state_file):
|
||||
manager.start()
|
||||
await asyncio.sleep(0.2)
|
||||
embed = manager.services["embed"]
|
||||
assert embed.pending_reload is False
|
||||
assert embed.depth == "sleeping" # restored from the file
|
||||
assert backend.services["vllm-embed"]["reload_calls"] == 0
|
||||
|
||||
|
||||
async def test_no_reload_when_no_wake_in_progress(tmp_path, backend, pub):
|
||||
"""Plain restart with a settled backend: no reload, no state churn."""
|
||||
state_file = str(tmp_path / "state.json")
|
||||
write_state(state_file,
|
||||
text={"depth": "awake", "wake_in_progress": False},
|
||||
ocr={"depth": "offloaded", "wake_in_progress": False, "level": 2})
|
||||
|
||||
cfg = clone(load_config())
|
||||
cfg.state_file = state_file
|
||||
cfg.idle_enabled = False
|
||||
async for _public, _admin, manager in _make(cfg, backend, state_file):
|
||||
manager.start()
|
||||
await asyncio.sleep(0.1)
|
||||
assert manager.services["ocr"].depth == "offloaded"
|
||||
assert backend.count("POST vllm-ocr/wake_up") == 0
|
||||
assert backend.count("POST vllm-text/collective_rpc") == 0
|
||||
|
||||
# the awake service serves immediately, straight through the proxy
|
||||
status, body = await raw_asgi(_public_app(manager), "POST",
|
||||
"/v1/chat/completions",
|
||||
headers=[("content-type", "application/json")],
|
||||
body=b'{"model": "text"}')
|
||||
assert status == 200, body
|
||||
svc = backend.services["vllm-text"]
|
||||
assert svc["wake_up_calls"] == 0 # no pointless wake
|
||||
assert svc["reload_calls"] == 0 # no pointless reload
|
||||
|
||||
|
||||
def _public_app(manager):
|
||||
from routing import build_public_app
|
||||
return build_public_app(manager)
|
||||
|
||||
|
||||
async def test_interrupted_wake_end_to_end(tmp_path, backend):
|
||||
"""In-process replay of E2E case 12: the wake is killed after wake_up has
|
||||
flipped is_sleeping but before reload_weights; a new manager on the same
|
||||
state file must complete the sequence before proxying."""
|
||||
state_file = str(tmp_path / "state.json")
|
||||
cfg = clone(load_config())
|
||||
cfg.state_file = state_file
|
||||
cfg.idle_enabled = False
|
||||
cfg.wake_health_poll_s = 0.005
|
||||
|
||||
backend.set_sleeping("text", True, level=2)
|
||||
backend.services["vllm-text"]["wake_delay"] = 0.4
|
||||
backend.services["vllm-text"]["wake_flip_midway"] = True
|
||||
|
||||
first = ServiceManager(cfg, transport=backend)
|
||||
try:
|
||||
task = asyncio.create_task(first.ensure_awake("text"))
|
||||
await asyncio.sleep(0.3) # is_sleeping flipped at t+0.2s
|
||||
task.cancel() # ... and the router "dies"
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
finally:
|
||||
await first.proxy_client.aclose()
|
||||
await first.ctrl_client.aclose()
|
||||
|
||||
assert backend.services["vllm-text"]["sleeping"] is False # looks awake
|
||||
assert backend.services["vllm-text"]["reload_calls"] == 0 # ... but isn't
|
||||
assert read_state(state_file)["text"]["wake_in_progress"] is True
|
||||
|
||||
backend.calls.clear()
|
||||
async for public_app, _admin, manager in _make(cfg, backend, state_file):
|
||||
manager.start()
|
||||
response = await _post(public_app, "/v1/chat/completions", {"model": "text"})
|
||||
assert response.status_code == 200
|
||||
svc = backend.services["vllm-text"]
|
||||
assert svc["reload_calls"] == 1 # sequence was completed ...
|
||||
assert backend.last("POST vllm-text/v1/chat/completions") > \
|
||||
backend.last("POST vllm-text/reset_prefix_cache") # ... first
|
||||
assert manager.services["text"].pending_reload is False
|
||||
|
||||
|
||||
async def test_wake_intent_is_written_and_cleared(tmp_path, backend):
|
||||
"""The file says wake_in_progress=true while a wake is running and false
|
||||
once it completes -- which is exactly what a restart in the middle leaves
|
||||
behind."""
|
||||
state_file = str(tmp_path / "state.json")
|
||||
cfg = clone(load_config())
|
||||
cfg.state_file = state_file
|
||||
cfg.idle_enabled = False
|
||||
cfg.wake_health_poll_s = 0.005
|
||||
backend.set_sleeping("text", True, level=2)
|
||||
backend.services["vllm-text"]["wake_delay"] = 0.2
|
||||
|
||||
manager = ServiceManager(cfg, transport=backend)
|
||||
try:
|
||||
task = asyncio.create_task(manager.ensure_awake("text"))
|
||||
await asyncio.sleep(0.05) # mid-wake
|
||||
assert read_state(state_file)["text"]["wake_in_progress"] is True
|
||||
outcome = await task
|
||||
assert outcome.ok
|
||||
entry = read_state(state_file)["text"]
|
||||
assert entry["wake_in_progress"] is False
|
||||
assert entry["depth"] == "awake"
|
||||
finally:
|
||||
await manager.proxy_client.aclose()
|
||||
await manager.ctrl_client.aclose()
|
||||
|
||||
|
||||
async def test_failed_wake_clears_the_intent(tmp_path, backend):
|
||||
state_file = str(tmp_path / "state.json")
|
||||
cfg = clone(load_config())
|
||||
cfg.state_file = state_file
|
||||
cfg.idle_enabled = False
|
||||
cfg.wake_health_poll_s = 0.005
|
||||
backend.set_sleeping("text", True)
|
||||
backend.services["vllm-text"]["wake_fails"] = 9
|
||||
|
||||
manager = ServiceManager(cfg, transport=backend)
|
||||
try:
|
||||
outcome = await manager.ensure_awake("text")
|
||||
assert not outcome.ok
|
||||
entry = read_state(state_file)["text"]
|
||||
assert entry["wake_in_progress"] is False
|
||||
finally:
|
||||
await manager.proxy_client.aclose()
|
||||
await manager.ctrl_client.aclose()
|
||||
|
||||
|
||||
async def test_state_file_failure_is_not_fatal(tmp_path, backend):
|
||||
"""An unwritable state path disables persistence, never serving."""
|
||||
cfg = clone(load_config())
|
||||
cfg.state_file = str(tmp_path / "no-such-dir" / "state.json")
|
||||
cfg.idle_enabled = False
|
||||
async for public_app, _admin, manager in _make(cfg, backend, cfg.state_file):
|
||||
manager.start()
|
||||
response = await _post(public_app, "/v1/chat/completions", {"model": "text"})
|
||||
assert response.status_code == 200
|
||||
assert not os.path.exists(cfg.state_file)
|
||||
|
||||
|
||||
async def test_admin_reports_recovery(tmp_path, backend):
|
||||
state_file = str(tmp_path / "state.json")
|
||||
write_state(state_file, ocr={"depth": "offloaded", "wake_in_progress": True, "level": 2})
|
||||
backend.services["vllm-ocr"]["sleeping"] = False
|
||||
|
||||
cfg = clone(load_config())
|
||||
cfg.state_file = state_file
|
||||
cfg.idle_enabled = False
|
||||
async for _public, admin_app, manager in _make(cfg, backend, state_file):
|
||||
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=admin_app),
|
||||
base_url="http://router.test") as client:
|
||||
r = await client.get("/admin/status")
|
||||
assert r.status_code == 200
|
||||
assert r.json()["services"]["ocr"]["wake_recovery_pending"] is True
|
||||
Reference in New Issue
Block a user