"""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, }