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:
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,
|
||||
}
|
||||
Reference in New Issue
Block a user