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>
297 lines
12 KiB
Python
297 lines
12 KiB
Python
"""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
|