bing
80eef4ce6a
Three always-running vLLM services (text TP=2 GPU0+1, ocr + embed on GPU2, sleep mode) behind a FastAPI router that auto-wakes models on request. Tiered idle (sleep 15 min / offload 3 h), depth-aware 503s with Retry-After, persisted wake-intent recovery, admin API on 127.0.0.1:8010. Routine control via vllmctl is pure HTTP — no docker on the request path. Verified: 91 router unit tests + 15-test E2E on real hardware (measurements in CALIBRATION.md; design record in .claude/memory/router-front-door-plan.md). Old nginx stack files removed before git init; design survives in .claude/memory/sleep-mode-implementation-plan.md. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
264 lines
9.8 KiB
Python
264 lines
9.8 KiB
Python
"""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
|