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>
696 lines
31 KiB
Python
696 lines
31 KiB
Python
"""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,
|
|
}
|