From 80eef4ce6a7d601e6d93ac0088d26ec6e392494e Mon Sep 17 00:00:00 2001 From: bing Date: Mon, 17 Aug 2026 10:17:42 +0000 Subject: [PATCH] Initial commit: router front-door vLLM stack MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .claude/memory/MEMORY.md | 57 ++ .claude/memory/front-door-proxy-design.md | 164 ++++ .claude/memory/router-front-door-plan.md | 339 ++++++++ .../memory/sleep-mode-implementation-plan.md | 783 ++++++++++++++++++ .env.example | 9 + .gitignore | 19 + CALIBRATION.md | 266 ++++++ GPU-report.txt | 69 ++ NOTES-2026-08-13.md | 210 +++++ NOTES-2026-08-17-gpu2-recheck.md | 146 ++++ README.md | 156 ++++ TODO.md | 53 ++ claude_code_env.sh | 205 +++++ compose.yml | 119 +++ diag/gpu_integrity.py | 33 + diag/p2p_check.py | 58 ++ diag/weight_check.py | 43 + router/Dockerfile | 28 + router/admin_api.py | 143 ++++ router/app.py | 187 +++++ router/config.py | 293 +++++++ router/pytest.ini | 5 + router/requirements.txt | 10 + router/routing.py | 536 ++++++++++++ router/services.py | 695 ++++++++++++++++ router/tests/conftest.py | 263 ++++++ router/tests/test_admin_api.py | 120 +++ router/tests/test_depth_errors.py | 207 +++++ router/tests/test_idle.py | 195 +++++ router/tests/test_model_resolution.py | 160 ++++ router/tests/test_paths.py | 146 ++++ router/tests/test_shutdown_and_cli.py | 101 +++ router/tests/test_singleflight.py | 95 +++ router/tests/test_wake_recovery.py | 296 +++++++ vllmctl | 297 +++++++ 35 files changed, 6506 insertions(+) create mode 100644 .claude/memory/MEMORY.md create mode 100644 .claude/memory/front-door-proxy-design.md create mode 100644 .claude/memory/router-front-door-plan.md create mode 100644 .claude/memory/sleep-mode-implementation-plan.md create mode 100644 .env.example create mode 100644 .gitignore create mode 100644 CALIBRATION.md create mode 100644 GPU-report.txt create mode 100644 NOTES-2026-08-13.md create mode 100644 NOTES-2026-08-17-gpu2-recheck.md create mode 100644 README.md create mode 100644 TODO.md create mode 100644 claude_code_env.sh create mode 100644 compose.yml create mode 100644 diag/gpu_integrity.py create mode 100644 diag/p2p_check.py create mode 100644 diag/weight_check.py create mode 100644 router/Dockerfile create mode 100644 router/admin_api.py create mode 100644 router/app.py create mode 100644 router/config.py create mode 100644 router/pytest.ini create mode 100644 router/requirements.txt create mode 100644 router/routing.py create mode 100644 router/services.py create mode 100644 router/tests/conftest.py create mode 100644 router/tests/test_admin_api.py create mode 100644 router/tests/test_depth_errors.py create mode 100644 router/tests/test_idle.py create mode 100644 router/tests/test_model_resolution.py create mode 100644 router/tests/test_paths.py create mode 100644 router/tests/test_shutdown_and_cli.py create mode 100644 router/tests/test_singleflight.py create mode 100644 router/tests/test_wake_recovery.py create mode 100755 vllmctl diff --git a/.claude/memory/MEMORY.md b/.claude/memory/MEMORY.md new file mode 100644 index 0000000..c75613a --- /dev/null +++ b/.claude/memory/MEMORY.md @@ -0,0 +1,57 @@ +# vLLM Project Memory + +Index of memory files for this project. + +## Architecture & Design + +- [Router Front Door Plan v3.3](router-front-door-plan.md) — **IMPLEMENTED & VERIFIED 2026-08-17**: per-model vLLM services (all sleeping) + FastAPI router that auto-wakes on request. Pure-HTTP load/unload, no docker on the request path, replaced nginx. Two review rounds + full E2E. + +- [Front Door Proxy Design](front-door-proxy-design.md) — original research (vLLM Sleep Mode, Triton, vllm-proxy) that led here. + +- [Sleep Mode Implementation Plan v2](sleep-mode-implementation-plan.md) — the earlier Sleep Mode + nginx stack (implemented 2026-08-14, then superseded by the router plan). + +## Project Documentation + +- [`../README.md`](../README.md) — architecture, request behavior (503/Retry-After semantics), vllmctl, router tunables +- [`../TODO.md`](../TODO.md) — implementation record, verification trail, known limitations, housekeeping +- [`../CALIBRATION.md`](../CALIBRATION.md) — measured GPU footprints, wake latencies, backend quirks +- [`../NOTES-2026-08-13.md`](../NOTES-2026-08-13.md) — 2026-08-13 rework + GPU P2P diagnosis +- [`../NOTES-2026-08-17-gpu2-recheck.md`](../NOTES-2026-08-17-gpu2-recheck.md) — P2P fault moved (not fixed); root-cause analysis (VT-d peer-DMA translation); fix ladder for admin + +## Model Configuration (live) + +| Key | Served model | Service / GPUs | util | +|-----|--------------|----------------|------| +| text | `Qwen3.6-35B-A3B-FP8` | vllm-text, TP=2, GPU0+1 | 0.85, 262K ctx | +| ocr | `OvisOCR2` | vllm-ocr, GPU2 | 0.10 | +| embed | `Qwen3-Embedding-8B` | vllm-embed, GPU2 | 0.25 | + +Public API `:8000/v1` (all three always available, auto-wake); admin API +`127.0.0.1:8010`; debug ports `127.0.0.1:8001-8003`. + +## Hard-won facts (do not regress) + +- **P2P corrupts data host-wide** (moves between GPUs across reboots; VT-d + peer-DMA platform fault). `NCCL_P2P_DISABLE=1` + `--disable-custom-all-reduce` + mandatory everywhere. GPU2 is fine for single-GPU services. +- **vLLM backend quirks**: `/health` returns 200 while asleep (gate on + `/is_sleeping`); requests to a sleeping backend hang; requests between + `wake_up` and `reload_weights` return garbage; L1 wake needs `wake_up` + only; L1→L2 re-sleep is a no-op (wake-then-sleep to actually offload). +- **`gpu_memory_utilization` counts TOTAL GPU memory per process**; keep + per-GPU sums ≤ 0.88 including CUDA contexts. +- **Router state volume must be a named volume** (`router-state`), not a + bind mount — the non-root container user (10001) can't write a host-uid + dir, and the state store fails OPEN (fix silently inert). +- Docker: `renbaibing` in docker group since 2026-08-14; shells older than + that need `sg docker -c "…"`. +- OCR needs `--max-num-seqs 256` (GDN block limit at util 0.10). +- FP8 model dir was corrupted by racing `hf download`s on 2026-08-14; + repaired + sha256-verified 2026-08-17 (corrupt copy kept, see TODO). + +## Housekeeping reminders + +- Delete `Qwen3.6-35B-A3B-FP8.corrupt-20260817` (43 GB, in MODEL_ROOT) once trusted. +- Repo under git since 2026-08-17 (initial commit = router stack). The v2 + nginx stack's runnable files were removed before git init; its design + survives in `sleep-mode-implementation-plan.md`. diff --git a/.claude/memory/front-door-proxy-design.md b/.claude/memory/front-door-proxy-design.md new file mode 100644 index 0000000..cd84f97 --- /dev/null +++ b/.claude/memory/front-door-proxy-design.md @@ -0,0 +1,164 @@ +# Front Door Proxy Design — Handling Requests to Unloaded Models + +**Date:** 2026-08-13 +**Status:** Design exploration — vLLM Sleep Mode available as better alternative + +> **TL;DR:** vLLM 0.11.0+ includes **Sleep Mode** — a built-in feature that enables fast model switching (18-200x faster than full reload) by hibernating weights while preserving process state. **This is likely the solution we need instead of building a custom front door.** + +## Current Behavior: Request to Unloaded Model + +When a client makes a request to `http://:8000/v1/...` and no model is loaded: + +1. **Container is down** (auto-unload via `idle-watch` or manual `down`) +2. **Immediate connection failure** — client gets `Connection refused` or similar +3. **No retry guidance** — client doesn't know if the service is down, moving, or loading +4. **Manual recovery required** — someone must run `./vllmctl up ` + +### What the client sees + +```bash +$ curl http://host:8000/v1/models +curl: (7) Failed to connect to host port 8000: Connection refused +``` + +No indication of: +- Whether the service is permanently down +- Which model should be loaded +- How long loading would take (~2-10 min for large models from NFS) + +## The Shared Front Door Concept + +A reverse proxy that sits in front of vLLM and handles the "unloaded model" case gracefully. + +### Architecture + +``` +Client → Front Door (port 8000) → vLLM Container (port 8001 or similar) + ↓ + [Model unloaded?] + ↓ + Trigger load / Queue / Inform +``` + +### Design Options + +#### Option 1: Wake-on-Request (Auto-Load) + +On receiving a request when no model is loaded: + +1. **Spawn the container** with the requested model (from URL path or header) +2. **Queue the request** (or return 503 with `Retry-After`) +3. **Client retries** after model is ready (or proxy forwards queued request) + +**Pros:** Transparent to clients, no manual intervention +**Cons:** Long wait times (2-10 min), cold start, need to decide which model to load +**Complexity:** Need request queue, loading state tracking, model selection logic + +#### Option 2: Friendly Error + Status Endpoint + +On receiving a request when no model is loaded: + +1. **Return 503 Service Unavailable** with clear JSON body +2. **Include status info:** "No model loaded. Load with `vllmctl up `" +3. **Provide a status endpoint** clients can poll: `GET /front-door/status` + - Shows: models available, currently loaded, loading state (if any) + +**Pros:** Simple to implement, clear client guidance +**Cons:** Still requires manual load or client-side logic +**Complexity:** Low — just a small nginx/envoy/python proxy + +#### Option 3: Request Queue + Load Notification + +On receiving a request when no model is loaded: + +1. **Return 202 Accepted** with a `Location` header pointing to a job status URL +2. **Trigger model load** in background +3. **Client polls status** or gets notified (webhook) when ready +4. **When ready, client retries** original request + +**Pros:** Clear async flow, client knows what's happening +**Cons:** More complex, needs job tracking, notification or polling +**Complexity:** Medium — need job store, status API + +## Implementation Sketch (Option 2: Friendly Error) + +### Simple Python Front Door + +```python +# front-door.py +from http.server import HTTPServer, BaseHTTPRequestHandler +import subprocess +import requests +import json + +VLLM_PORT = 8001 +FRONT_PORT = 8000 +MODELS_ROOT = "/data/home/renbaibing/huggingface" + +class FrontDoorHandler(BaseHTTPRequestHandler): + def do_GET(self): + # Check if vLLM is healthy + try: + resp = requests.get(f"http://127.0.0.1:{VLLM_PORT}/health", timeout=1) + # Proxy the request + proxy_resp = requests.get(f"http://127.0.0.1:{VLLM_PORT}{self.path}") + self.wfile.write(proxy_resp.content) + except requests.exceptions.RequestException: + # vLLM is not ready + self.send_response(503) + self.send_header("Content-Type", "application/json") + self.end_headers() + response = { + "error": "No model currently loaded", + "message": "Load a model with: ./vllmctl up ", + "available_models": self.list_models(), + "status_url": f"http://{self.server.server_port}/status" + } + self.wfile.write(json.dumps(response).encode()) + + def list_models(self): + # Scan MODEL_ROOT for model directories + # Return list of available model names + pass + +HTTPServer(("", FRONT_PORT), FrontDoorHandler).serve_forever() +``` + +### Integration with existing setup + +1. Move vLLM to port 8001 (`VLLM_PORT=8001` in .env) +2. Run front door on port 8000 (or via compose as a separate service) +3. Front door checks vLLM health before proxying + +## Open Questions + +1. **Model selection:** If auto-loading, how does the proxy know which model to load? + - From URL path? (`/v1/models/NAME/completions`?) + - From header? (`X-Model-Name: Qwen3.6-35B-A3B`) + - Configured default? + +2. **Cold start cost:** Large models take 2-10 minutes to load from NFS + - Is this acceptable for clients? + - Should we cache recently-used models? + +3. **Multi-model concurrent serving:** Does vLLM support serving multiple models simultaneously? + - If yes, the proxy could load models on-demand and keep them warm + - If no, we need model switching logic (potentially disruptive to existing requests) + +4. **Authentication:** Should the front door handle API keys? + - Current setup has no auth + - Adding auth at the front door would be natural + +## Next Steps + +1. **Confirm requirements:** Decide between Option 1 (auto-load), 2 (friendly error), or 3 (async queue) +2. **Prototype:** Build a minimal front door for the chosen option +3. **Integration:** Wire into existing `compose.yml` and `vllmctl` workflow +4. **Testing:** Simulate various failure modes and client behaviors + +## References + +- Current setup: `README.md`, `NOTES-2026-08-13.md` +- vLLM health endpoint: `http://127.0.0.1:8000/health` +- vLLM metrics: `http://127.0.0.1:8000/metrics` (used by idle-watch) +- OpenAI API compatibility: `http://127.0.0.1:8000/v1/...` diff --git a/.claude/memory/router-front-door-plan.md b/.claude/memory/router-front-door-plan.md new file mode 100644 index 0000000..9c08691 --- /dev/null +++ b/.claude/memory/router-front-door-plan.md @@ -0,0 +1,339 @@ +# Plan v3.3 — Router Front Door (auto wake-on-request) + +**Date:** 2026-08-14, revised 2026-08-17 (v3.3) +**Status:** **IMPLEMENTED & VERIFIED 2026-08-17** — v3.2 passed two review rounds (5 critical + 6 major fixed and confirmed; round-2 minors applied). v3.3 = hardware-driven placement change (GPU2 cleared for single-GPU services → OCR+embed on GPU2, text exclusive on GPU0/1). Implementation record and verification trail: `../../TODO.md`; measured numbers: `../../CALIBRATION.md`. One post-plan correctness fix: persisted wake-intent + startup recovery (E2E case 12), see TODO "Verification trail". +**Decisions made with user:** multiple models may be awake concurrently; router = Python FastAPI; tiered idle (level-1 sleep after 15 min, level-2 offload after 3 hr); depth-aware 503 errors. + +## 1. Problem & constraints + +Calling services must be able to hit `http://:8000/v1/...` with any of our +three models at any time, with no knowledge of load state. Hard constraints: + +1. **Native to the Docker stack** — the entire request path runs inside + containers; no host-side cron, watchers, or helpers. +2. **No docker permissions on the request path** — model load/unload/switch is + pure HTTP (vLLM Sleep Mode). Docker is only needed for rare manual ops + (image upgrade, pulling new models). +3. **Highly reliable** — failures degrade to depth-aware 503 + `Retry-After`, + never a hang, never a proxy to a half-awake backend. + +The Plan v2 nginx front door hid dev endpoints but did **not** wake models on +request; that inverted priority is what this plan fixes. + +## 2. Hardware reality (re-verified 2026-08-17 — see NOTES-2026-08-17-gpu2-recheck.md) + +- **3× A800-SXM4-80GB**. P2P is **still broken host-wide, 4/6 paths corrupt**, + but the fault MOVED: now every transfer **sourced from GPU0 or GPU1** + corrupts; GPU2-sourced transfers are the only clean ones (was the inverse + on 2026-08-13). Per-GPU integrity (H2D/D2H, matmul, stress) passes on all + three, twice. +- Consequences: `NCCL_P2P_DISABLE=1` + `--disable-custom-all-reduce` stay + mandatory for any TP>1 workload (host-staged copies verify clean — which is + why TP=2 on GPU0+1 with P2P off works). **GPU2 is usable for single-GPU + services** (no P2P involved; H2D/D2H verified clean). +- On-disk models (verified): `Qwen3.6-35B-A3B-FP8` (43GB dir), `OvisOCR2` + (1.8GB), `Qwen3-Embedding-8B` (12GB). **User confirmed 2026-08-14: the 8B + embedding model is the intended one** (earlier docs mentioning a 0.6B were + wrong; slices are sized for the 8B). + +## 3. Architecture + +``` + ┌──────────────────── Docker network ────────────────────┐ +callers ─► :8000 ─► ┌─────────┐ ┌──────────────────────┐ ┌───────────────────┐ + │ router │ ─►│ vllm-text (TP=2) │ │ vllm-ocr GPU2 │ + │ FastAPI │ │ GPU0 + GPU1 only │ │ vllm-embed GPU2 │ + └─────────┘ └──────────────────────┘ └───────────────────┘ + 127.0.0.1:8010 (admin) +``` + +- **One always-running vLLM service per model**, each with + `--enable-sleep-mode` + `VLLM_SERVER_DEV_MODE=1`, each pinned to a GPU + memory slice so all three can be awake simultaneously. +- **Router** is the only public ingress (port 8000). It owns: model-name → + service mapping, wake-on-request, streaming proxy, tiered idle management, + and dev-endpoint hiding (404 for anything not an allowed path, matched on + the **normalized** path to defeat `/v1/../sleep` traversal). +- **Admin API is a separate listener** published only on `127.0.0.1:8010` — + NOT paths on the public port (review C5: a public `/admin/sleep` would be a + trivial remote DoS). +- **nginx is removed.** The router replaces it. +- **Placement (v3.3, user-approved 2026-08-17):** text TP=2 gets GPU0+GPU1 + **exclusively** (0.85 slice, full 262K context); OCR and embed run TP=1 + **on GPU2** — single-GPU services never issue P2P transfers, and GPU2's + H2D/D2H paths verify clean (§2). `CUDA_VISIBLE_DEVICES` is set explicitly + per service. Rollback if GPU2 ever misbehaves: change one env var to move + a service back to GPU0/1 (slices would need re-tightening). +- All services keep `restart: unless-stopped` and `ipc: host` (NCCL shm). + Each service's TP TCP store is isolated by its own network namespace — no + distributed-init port collision (would only become a risk if someone later + sets `network_mode: host`; do not). + +## 4. Port layout + +| Port | Binding | Purpose | +|---|---|---| +| 8000 | 0.0.0.0 (host) → router | Public API (only public ingress) | +| 8010 | 127.0.0.1 (host) → router admin listener | `vllmctl` / debugging | +| 8001 | 127.0.0.1 (host) → vllm-text :8000 | Direct debug access | +| 8002 | 127.0.0.1 (host) → vllm-ocr :8000 | Direct debug access | +| 8003 | 127.0.0.1 (host) → vllm-embed :8000 | Direct debug access | +| — | Docker net only | router → vllm-* | + +**Security boundary, stated explicitly:** the loopback debug ports 8001–8003 +expose each service's dev endpoints (`/sleep`, `/wake_up`, …) to **any local +user** on this multi-user host, bypassing the router's 404s. Same exposure as +the current stack's port 8001; accepted, but "dev endpoints unreachable" +means "unreachable from the network," not "unreachable locally." + +## 5. GPU memory budget (3× 80GB in use; calibration gates final numbers) + +`gpu_memory_utilization` is a fraction of **total** GPU memory, and each +process accounts only for itself. Budget per GPU: + + Σ(slices on that GPU) + (CUDA ctx per process ~1–2GB) + boot/wake transients ≤ 0.88 + +Initial slices (recalibrated in Step 1): + +| Service | Weights on disk | TP / GPU | util slice | Notes | +|---|---|---|---|---| +| vllm-text | 43GB (FP8) | 2 / GPU0+1 (exclusive) | 0.85 | full 262K context; KV ~25GB — no 131K fallback needed | +| vllm-ocr | 1.8GB | 1 / GPU2 | 0.10 | small vision model | +| vllm-embed | 12GB (8B) | 1 / GPU2 | 0.25 | `--max-model-len 8192`; raisable at calibration | + +Per-GPU: GPU0 = GPU1 = 0.85 + 1 ctx ≈ 0.87; GPU2 = 0.10 + 0.25 = 0.35 + 2 ctx. +Disjoint pools — text never competes with the small services, and the small +services have an entire idle GPU of headroom. + +Sleeping (level 2) residual is ~2.5GB per **process** (E2E measured ~5GB +total for one TP=2 instance = 1 proc/GPU). Under v3.3 placement: GPU0/GPU1 +host 1 process each (text) → ~2.5GB residual; GPU2 hosts 2 (ocr + embed) +→ ~5GB. Measured in Step 1. + +## 6. Router design + +Python 3.12-slim container, **pinned** fastapi/httpx/uvicorn versions, +`restart: unless-stopped`. ~300 lines. + +**Single process, single event loop, two listening sockets** (public :8000 + +admin :8010): run two `uvicorn.Server` instances as coroutines on **one** +asyncio loop (uvicorn can't bind two ports in one worker, and two separate +processes would break the in-process wake locks — the admin listener MUST +share the same locks and depth tracking as the request path). + +### 6.1 Model registry (env or `models.json`) + +``` +text : { service: "vllm-text", model: "Qwen3.6-35B-A3B-FP8", aliases: ["qwen3.6-35b-a3b-fp8", ...] } +ocr : { service: "vllm-ocr", model: "OvisOCR2" } +embed: { service: "vllm-embed", model: "Qwen3-Embedding-8B" } +``` + +(No 0.6B alias — the 8B is the confirmed model (2026-08-14); an alias would +mislead clients into thinking they got a different model than they did.) + +Matching is case-insensitive. Unknown `model` → OpenAI-style 404 +(`model_not_found`), **no wake triggered**. + +### 6.2 Request path + +1. Resolve target service: + - Content-Type `application/json` → parse body, use `model` field; + missing `model` defaults to text **only for** `/v1/chat/completions` + and `/v1/completions`; other endpoints require an explicit model. + - Non-JSON bodies (e.g. `multipart/form-data` image uploads for OCR — + review C2) → **never fully parse**; extract only the small `model` + text field from the multipart parts (without reading file parts) and + forward the raw buffered body unchanged. A multipart request on + `/v1/chat/completions` with no resolvable `model` **defaults to OCR** + (that's the image-bearing path). + - `/v1/embeddings` → embed service regardless of body. +2. Check service state (cached briefly — don't add a backend round trip to + every request): + - **awake** → proxy immediately (httpx streaming passthrough, no buffering). + - **sleeping/offloading** → single-flight per service under an asyncio + lock; concurrent callers await the same event. Re-check `is_sleeping` + after acquiring the lock (another path may have just woken or slept it). +3. Wake sequence (proven in E2E on 2026-08-14): + `POST /wake_up` → `POST /collective_rpc {"method":"reload_weights"}` + → `POST /reset_prefix_cache` → poll `GET /health` until 200. +4. Proxy the held request. + +**Proxy timeouts (review M4):** httpx connect timeout 5s; read timeout +**900s** (matching today's nginx — capped, never disabled, so a wedged +backend stream can't hang forever) for both normal and streaming paths; +connection pool sized for ≥16 concurrent requests. A 5s default read timeout +would kill nearly every generation. + +**Other public endpoints:** +- `/v1/models` → aggregated list of all three registered models (clients + commonly list before calling). +- `/health` → router liveness + per-service state summary (uptime checks). +- `/metrics` → pass through from all services, or omit; decide at + implementation (was unauthenticated under nginx anyway). + +### 6.2.1 Depth-aware error semantics + +The router tracks each service's sleep depth from its own actions. +**Depth loss after router restart or direct `vllmctl` admin calls:** if +sleeping but depth unknown → treat as `offloaded` (conservative — worst case +the client waits slightly longer than needed, never gets a too-early retry). + +Hold-first policy: requests queue during wake and only error if the +depth-dependent hold deadline is exceeded. + +| Situation | Hold deadline | then Status | `Retry-After` | body `sleep_depth` | +|---|---|---|---|---| +| Wake from level 1 (sleep, ~1–6s) | 30s | 503 | `10` | `"sleeping"` | +| Wake from level 2 (offload, ~15–60s+) | 180s | 503 | `60` | `"offloaded"` | +| Container restarting (cold start = 2–10 min NFS load) | 300s | 503 | `600` | `"restarting"` | +| Unknown model | — | 404 | — | `model_not_found` | + +503 body (OpenAI-style, parseable): + +```json +{"error": {"type": "model_waking", "code": "model_waking", + "message": "Model 'OvisOCR2' is waking from offload; retry shortly", + "sleep_depth": "offloaded", "estimated_wake_seconds": 45}} +``` + +**Wake-sequence failure:** retry the sequence once, then the depth-aware 503. +Never proxy to a half-awake backend. All deadlines/estimates get re-tuned +from latencies logged during calibration (Step 1). + +### 6.3 Tiered idle management (replaces `vllmctl idle-watch`) + +| Tier | Trigger (idle) | Action | State after | Wake cost | +|---|---|---|---|---| +| Sleep | 15 min (`IDLE_SLEEP_MIN`) | `POST /sleep?level=1` | weights → host RAM | ~1–6s | +| Offload | 3 hr (`IDLE_OFFLOAD_MIN`) | `POST /sleep?level=2` | RAM freed; ~5GB/GPU ctx | ~15–60s (NFS) | + +Race-safety (review M1): the idle manager takes the **same per-service +asyncio lock** as the wake path and re-checks under the lock that +(a) active-request count == 0 and (b) last-activity is still past the +threshold. Last-activity is refreshed at request **completion** as well as +arrival, and an active-request counter guards long streaming generations — +a 20-minute stream must not trigger sleep mid-flight. + +**Verify at implementation:** `/sleep?level=2` on a service already sleeping +at level 1 must discard the offloaded weights (not no-op/error). Fallback: +wake → re-sleep at level 2. + +**Verify at implementation:** host RAM — level 1 holds a full weights copy +(text 43GB + ocr + embed ≈ 57GB total if all three nap simultaneously). +Check `free -g` headroom. + +**NFS risk (accepted, documented):** level-2 wake requires a healthy NFS +mount. If NFS is down, the model is unwakeable and the router 503s (with +`restarting`-grade Retry-After) until NFS returns. Mitigation option: cap the +text model at level 1 (weights stay in RAM). Decide after calibration. + +### 6.4 Router restart mid-wake (review M3) + +The wake lock is in-process and dies with the router. The new instance must +not fire a second concurrent wake at a backend already waking. Mitigation: +the wake sequence is effectively idempotent from the router's view (re-POST +`/wake_up` on an already-waking backend is safe; vLLM serializes it), and the +request path re-checks `is_sleeping` under the lock before acting. + +### 6.5 Admin API — separate listener, `127.0.0.1:8010` only + +- `GET /admin/status` — per-service: running?, sleeping?, depth, last activity, wake in progress +- `POST /admin/wake/{key}` / `POST /admin/sleep/{key}[?level=1|2]` — manual control (updates the router's depth tracking) + +### 6.6 Failure modes + +| Failure | Router behavior | +|---|---| +| vLLM container down/crashing | Detect unhealthy; hold up to 300s (`restart: unless-stopped` recovers it — but note cold start is 2–10 min for text); then 503 `Retry-After: 600` | +| Wake sequence fails | Retry once, then depth-aware 503 (6.2.1); never proxy half-awake | +| Router crash | `restart: unless-stopped`; vLLM services unaffected; at most in-flight requests lost; depth resets to conservative `offloaded` | +| Host reboot | **All three vLLM containers cold-start AWAKE** (vLLM has no boot-asleep) — see §7; text on GPU0/1 and small services on GPU2 initialize in parallel without racing; recovery takes minutes, router serves `restarting` 503s meanwhile | +| NFS outage | Level-2 unwakeable → 503s until NFS returns (§6.3) | +| Request without `model` field | Defaults to text **only** on chat/completions; elsewhere 400 | + +## 7. Compose changes + +- Add `router` service (build `./router/`): public `8000:8000`, admin + `127.0.0.1:8010:8010`, `depends_on` all vllm services + (`condition: service_started`). +- **Boot ordering (simplified in v3.3):** no staggering needed — text + initializes on GPU0/1 while OCR and embed initialize on GPU2, disjoint + memory pools, no race. All services use plain `depends_on: + service_started` (sleep makes healthchecks flappy, so never + `service_healthy`). The earlier wait-for-healthy entrypoint choreography + (review-2 issue #1) is obsolete under this placement. +- Split `vllm` into `vllm-text` / `vllm-ocr` / `vllm-embed` with per-model + args baked in; `MODEL_NAME`/`EXTRA_ARGS`/`TP_SIZE`/`MAX_MODEL_LEN` in + `.env` become obsolete for serving (kept only for rollback). +- Remove `nginx` service and `nginx.conf`. +- All services: `restart: unless-stopped`, `ipc: host`, + `NCCL_P2P_DISABLE=1`, explicit `CUDA_VISIBLE_DEVICES`. + +Per-model serving args (all include `--disable-custom-all-reduce +--enable-sleep-mode`, **explicit `--max-model-len` each** — review M2): + +- **text:** `--reasoning-parser qwen3 --enable-auto-tool-choice --tool-call-parser qwen3_coder --max-model-len 262144 --gpu-memory-utilization 0.85` +- **ocr:** `--max-model-len 32768 --gpu-memory-utilization 0.10` (+ vision args from OvisOCR2 model card, e.g. `--limit-mm-per-prompt`) +- **embed:** `--max-model-len 8192 --gpu-memory-utilization 0.25` + +## 8. `vllmctl` changes + +- `up [MODEL]` / `down [MODEL]` → HTTP to `127.0.0.1:8010` admin API (no + docker needed for routine control). +- `status` → `GET 127.0.0.1:8010/admin/status` pretty-printed. +- `pull` unchanged (docker, rare, manual). +- `idle-watch` removed (router owns idle management). + +## 9. Implementation order + +0. **Backups + cutover plan (review M5):** + `cp compose.yml compose.yml.v2.bak; cp vllmctl vllmctl.v2.bak`. + **Status 2026-08-14: `renbaibing` is now in the docker group** (works + directly in fresh login shells; use `sg docker -c "…"` in sessions + started before the change — `.runas.py`/`.user.env` are obsolete). + Old stack already brought down by user; GPUs idle. Rollback: restore + `.v2.bak` files, `docker compose up -d`. +1. **Calibration (gates everything):** bring up the three services with + initial slices (old stack stopped); measure awake footprints, sleeping + residuals, level-1 and level-2 wake latencies, host RAM cost; adjust + slices so each GPU ≤ 0.88 including contexts. Log latencies → tune the + §6.2.1 deadlines from data. +2. `router/` app (registry, wake single-flight, raw-body-safe routing, + streaming proxy with proper timeouts, tiered idle with lock, admin + listener) + compose rewiring; nginx removed. +3. `vllmctl` rewrite of up/down/status. +4. E2E verification (below). +5. Docs: README, TODO, memory. + +## 10. E2E test matrix + +1. All asleep → `POST /v1/chat/completions` (text) → 200 after wake; latency logged. +2. While text awake → `POST /v1/embeddings` → embed wakes independently; both serve. +3. OCR request with **`multipart/form-data`** body (not just JSON base64) → routes correctly, no body parsing error. +4. Idle tiers (shortened timers): 1 min → level-1 sleep (host RAM grows by weights); 3 min → level-2 offload (RAM freed); auto-wake from both. +5. Streaming request during wake → first chunk after wake, no buffering. +6. Unknown model name → 404 `model_not_found`, no wake triggered. +7. Dev endpoints (`/sleep`, `/wake_up`, `/collective_rpc`, `/reset_prefix_cache`) **and** `/admin/*` on public :8000 → 404; admin reachable only via 127.0.0.1:8010. +8. Kill vllm-text container → request → `restarting` 503 (`Retry-After: 600`) during recovery → 200 after restart completes. +9. Concurrent 10× requests while sleeping → exactly one wake sequence (router logs). +10. Depth-aware errors (hold deadline forced to 1s): level-1 wake → 503 `Retry-After: 10`, `"sleeping"`; level-2 → `Retry-After: 60`, `"offloaded"`; body parses with `estimated_wake_seconds`. +11. Idle-race: fire a request at the moment the idle timer expires → request served, no sleep mid-request (lock + active-counter verified). +12. Router restart mid-wake → no double-wake crash; request eventually served. +13. Two services waking concurrently (text + embed requested simultaneously) → both complete, no lock cross-talk. +14. Path traversal: `/v1/../sleep`, `/v1%2f..%2fsleep` → 404. +15. Host reboot (if feasible to test): parallel boot (text on GPU0/1, small services on GPU2 — no race), router serves `restarting` 503s, all three become available. + +## 11. Risks / open items + +- **GPU slice calibration** gates final numbers (§5, Step 1). v3.3 placement + (text exclusive on GPU0/1 at 0.85) removes the KV-starvation concern; + calibration confirms. +- **GPU2 trust:** single-GPU tests pass (2026-08-17, twice), but it has a + fault history (Aug 11 driver re-probe). Mitigation: E2E test 3+4 exercise + it before rollout; if GPU2 degrades in production, move OCR/embed back to + GPU0/1 with one env-var change each (slices re-tighten to v3.2 values). +- ~~Which embedding model~~ — **resolved 2026-08-14: Qwen3-Embedding-8B.** +- **Level 1→2 escalation** and **host RAM** (~57GB if all three nap) need verification (§6.3). +- **NFS outage strands level-2 models** — accepted risk or cap text at level 1 (§6.3). +- Wake latency must fit calling services' client timeouts — confirm their settings before rollout. diff --git a/.claude/memory/sleep-mode-implementation-plan.md b/.claude/memory/sleep-mode-implementation-plan.md new file mode 100644 index 0000000..9f1607a --- /dev/null +++ b/.claude/memory/sleep-mode-implementation-plan.md @@ -0,0 +1,783 @@ +# Implementation Plan: vLLM Sleep Mode + nginx Front Door + +**Date:** 2026-08-14 +**Status:** Revised (v2) — Critical issues from review addressed + +## Changelog (v2) + +**Fixed Critical Issues:** +1. ✅ Healthcheck dependency — changed to `service_started` (doesn't require healthy) +2. ✅ nginx upstream host — changed from `127.0.0.1:8001` to `vllm:8000` +3. ✅ Port binding clarified — vLLM: `127.0.0.1:8001:8000` (localhost only), nginx: external 8000 +4. ✅ Added `--enable-sleep-mode` flag to EXTRA_ARGS +5. ✅ Implemented idle watcher update with `/sleep` calls + +**Important Concerns Addressed:** +- Added note about `/metrics` exposure +- Clarified in-flight request handling during `/sleep` +- Updated wake time estimate to be more realistic +- Removed redundant `return 403` after `deny all` + +--- + +## Overview + +Transform the current vLLM setup from full container teardown (`docker compose down`) to vLLM Sleep Mode, with nginx as a security front door. + +### Current Architecture + +``` +Client → [vLLM Container :8000] + ↓ + down = docker compose down (full teardown) + up = docker compose up (cold start, 2-10 min) +``` + +### Target Architecture + +``` +Internet → [nginx :8000] → [vLLM :8000 (Docker network) ] + ↓ + [vLLM :8001 (localhost only) ] + ↓ + sleep = POST /sleep (model hibernates, GPU freed) + wake = POST /wake_up (model resumes, ~5-15s for 35B) +``` + +### Port Binding Clarification + +| Port | Access | Purpose | +|------|--------|---------| +| `8000` (nginx) | External (0.0.0.0:8000) | Public API, clients connect here | +| `8001` (vLLM) | Localhost only (127.0.0.1:8001) | Admin access for `vllmctl`, sleep/wake endpoints | +| `8000` (Docker network) | Internal | nginx → vLLM communication within Docker network | + +--- + +## Phase 1: Enable vLLM Sleep Mode + +### 1.1 Update `compose.yml` + +**Critical changes:** + +1. Add `VLLM_SERVER_DEV_MODE=1` to environment (required for Sleep Mode endpoints) +2. Add `--enable-sleep-mode` to EXTRA_ARGS (this was missing in v1!) +3. Expose vLLM on localhost only: `127.0.0.1:8001:8000` +4. Keep vLLM accessible on Docker network at port 8000 + +**Before:** +```yaml +services: + vllm: + image: vllm/vllm-openai:${VLLM_VERSION:-v0.27.1} + container_name: vllm + restart: unless-stopped + ports: + - "${VLLM_PORT:-8000}:8000" + volumes: + - ${MODEL_ROOT:-/data/home/renbaibing/huggingface}:/models:ro + ipc: host + gpus: all + environment: + - HF_TOKEN=${HF_TOKEN:-} + - MODEL_NAME=${MODEL_NAME} + - TP_SIZE=${TP_SIZE:-2} + - MAX_MODEL_LEN=${MAX_MODEL_LEN:-32768} + - GPU_MEM_UTIL=${GPU_MEM_UTIL:-} + - EXTRA_ARGS=${EXTRA_ARGS:-} + - NCCL_P2P_DISABLE=${NCCL_P2P_DISABLE:-1} + healthcheck: + test: ["CMD-SHELL", "python3 -c \"import urllib.request; urllib.request.urlopen('http://127.0.0.1:8000/health', timeout=5)\" || exit 1"] + interval: 15s + timeout: 10s + retries: 8 + start_period: 1200s +``` + +**After:** +```yaml +services: + vllm: + image: vllm/vllm-openai:${VLLM_VERSION:-v0.27.1} + container_name: vllm + restart: unless-stopped + + # Two port bindings: + # 1. Docker network port 8000 (internal, for nginx) + # 2. Localhost port 8001 (admin access for vllmctl) + ports: + - "127.0.0.1:8001:8000" # Admin access - localhost only + expose: + - "8000" # Docker network - for nginx + + volumes: + - ${MODEL_ROOT:-/data/home/renbaibing/huggingface}:/models:ro + ipc: host + gpus: all + + environment: + - VLLM_SERVER_DEV_MODE=1 # Required for Sleep Mode endpoints + - HF_TOKEN=${HF_TOKEN:-} + - MODEL_NAME=${MODEL_NAME} + - TP_SIZE=${TP_SIZE:-2} + - MAX_MODEL_LEN=${MAX_MODEL_LEN:-32768} + - GPU_MEM_UTIL=${GPU_MEM_UTIL:-} + # EXTRA_ARGS MUST include --enable-sleep-mode + - EXTRA_ARGS=${EXTRA_ARGS:- --enable-sleep-mode} + - NCCL_P2P_DISABLE=${NCCL_P2P_DISABLE:-1} + + # Healthcheck stays the same - container will be "unhealthy" during sleep + # This is OK - nginx uses service_started, not service_healthy + healthcheck: + test: ["CMD-SHELL", "python3 -c \"import urllib.request; urllib.request.urlopen('http://127.0.0.1:8000/health', timeout=5)\" || exit 1"] + interval: 15s + timeout: 10s + retries: 8 + start_period: 1200s +``` + +**Rationale:** +- `VLLM_SERVER_DEV_MODE=1` enables dev endpoints including `/sleep`, `/wake_up`, `/is_sleeping` +- `--enable-sleep-mode` in EXTRA_ARGS is REQUIRED for sleep functionality (this was the critical missing piece in v1) +- Port `127.0.0.1:8001` gives `vllmctl` localhost-only access to admin endpoints +- `expose: - "8000"` makes port 8000 available on Docker network for nginx +- Healthcheck will fail during sleep, but that's OK — nginx uses `service_started` condition + +### 1.2 Update `vllmctl` + +**First, update the API URL:** + +```bash +# Change at the top of vllmctl: +VLLM_PORT="$(env_get VLLM_PORT 8000)" +API="http://127.0.0.1:8001" # Direct to vLLM localhost port, bypasses nginx +``` + +**Modify `cmd_down()`:** + +```bash +cmd_down() { + say 'Putting vLLM to sleep (frees GPU memory, keeps server alive)...' + + # Try to sleep the model (level 2 = discard weights, minimal RAM) + if curl -fs -X POST "$API/sleep?level=2" >/dev/null 2>&1; then + ok "Model is sleeping. Server still running." + + # Wait a moment for sleep to complete + sleep 2 + + # Verify sleep state + local sleeping + sleeping="$(curl -fs "$API/is_sleeping" 2>/dev/null)" + if [ "$sleeping" = "true" ]; then + ok "Confirmed: Model is in sleep state." + else + warn "Sleep state unclear - check with: curl $API/is_sleeping" + fi + else + warn "Sleep request failed — server may not be ready. Falling back to full stop." + dcompose down -t 20 || return 1 + fi + + # Show GPU memory after sleep + local used + used="$(nvidia-smi --query-gpu=index,memory.used --format=csv,noheader,nounits 2>/dev/null | awk -F', *' '{printf "GPU%s:%sMiB ", $1, $2}')" + say "GPU memory now: $used" +} +``` + +**Rationale:** +- `level=2` discards weights entirely (minimal RAM usage) +- Wake time for 35B model: ~5-15s (more realistic than the 0.8-2.6s estimate) +- Adds sleep verification to confirm the operation worked +- Falls back to full stop if sleep fails + +**Modify `cmd_up()`:** + +```bash +cmd_up() { + local model="${1:-}" + if [ -n "$model" ]; then + if [ ! -d "$MODEL_ROOT/$model" ]; then + err "Model '$model' not found in $MODEL_ROOT" + err "Available:"; cmd_list + return 1 + fi + env_set MODEL_NAME "$model" + fi + model="$(env_get MODEL_NAME)" + [ -n "$model" ] || { err 'No MODEL_NAME set in .env'; return 1; } + + local st + st="$(container_state)" + + # Case 1: Container not running at all + case "$st" in + absent*|exited*|dead*) + say "Starting vLLM with model '${C_CYAN}$model${C_OFF}' (TP=$(env_get TP_SIZE 2), Sleep Mode enabled)…" + dcompose up -d || return 1 + wait_for_ready "$model" + return 0 + ;; + esac + + # Case 2: Container running - check if sleeping or different model + local sleeping loaded + sleeping="$(curl -fs "$API/is_sleeping" 2>/dev/null)" + loaded="$(loaded_model)" + + # Model already loaded and serving + if [ "$sleeping" != "true" ] && [ "$loaded" = "$model" ]; then + ok "Already serving '$model'." + return 0 + fi + + # Model is sleeping - wake it up + if [ "$sleeping" = "true" ]; then + say "Model is sleeping. Waking up '${C_CYAN}$model${C_OFF}'…" + # Level 2 sleep requires reload_weights after wake_up + curl -fs -X POST "$API/wake_up" >/dev/null 2>&1 || { + err "Wake-up request failed." + return 1 + } + # Reload weights (required for level 2 sleep) + curl -fs -X POST "$API/collective_rpc" \ + -H 'Content-Type: application/json' \ + -d '{"method":"reload_weights"}' >/dev/null 2>&1 || { + warn "reload_weights request failed - model may still wake up" + } + # Reset prefix cache (required for level 2 sleep) + curl -fs -X POST "$API/reset_prefix_cache" >/dev/null 2>&1 || { + warn "reset_prefix_cache failed - non-critical" + } + # Wait for model to be ready + wait_for_ready "$model" + ok "Model '$model' is awake and ready." + return 0 + fi + + # Different model loaded - need to switch + if [ -n "$loaded" ] && [ "$loaded" != "$model" ]; then + say "Switching from '${C_YELLOW}$loaded${C_OFF}' to '${C_CYAN}$model${C_OFF}'…" + say "Restarting container with new model (this will take a few minutes)…" + dcompose down -t 20 + sleep 2 + dcompose up -d || return 1 + wait_for_ready "$model" + ok "Now serving '$model'." + return 0 + fi + + # Shouldn't reach here, but handle gracefully + err "Unexpected state - container running but no model loaded. Try: ./vllmctl restart" + return 1 +} +``` + +**Rationale:** +- Now handles the sleeping state properly +- For level 2 sleep, calls the required sequence: `/wake_up` → `reload_weights` → `reset_prefix_cache` +- Keeps container restart for model switching (cleaner than in-place reload) +- Better error messages and state handling + +**Modify `cmd_status()`:** + +```bash +cmd_status() { + local st loaded cfg sleeping + st="$(container_state)" + cfg="$(env_get MODEL_NAME)" + + say "${C_CYAN}Container:${C_OFF} $st" + + # Check if sleeping first (loaded_model won't work when sleeping) + sleeping="$(curl -fs "$API/is_sleeping" 2>/dev/null)" + + if [ "$sleeping" = "true" ]; then + say "${C_CYAN}Serving:${C_OFF} ${C_DIM}model sleeping${C_OFF} (configured: $cfg)" + else + loaded="$(loaded_model)" + if [ -n "$loaded" ]; then + say "${C_CYAN}Serving:${C_OFF} ${C_GREEN}$loaded${C_OFF} at http://127.0.0.1:${VLLM_PORT}/v1" + else + say "${C_CYAN}Serving:${C_OFF} ${C_DIM}nothing loaded${C_OFF} (configured: $cfg)" + fi + fi + + say "${C_CYAN}GPU:${C_OFF}" + nvidia-smi --query-gpu=index,name,memory.used,memory.total,utilization.gpu \ + --format=csv,noheader 2>/dev/null | sed 's/^/ /' + + # External access (nginx) + say "${C_CYAN}Public API:${C_OFF} http://$(hostname -f | head -1):${VLLM_PORT}/v1" + say "${C_CYAN}Admin API:${C_OFF} $API (localhost only)" + + # Idle watcher status + local pidfile="$ROOT/.idle.pid" + if [ -f "$pidfile" ] && kill -0 "$(cat "$pidfile" 2>/dev/null)" 2>/dev/null; then + say "${C_CYAN}Idle-watch:${C_OFF} active (pid $(cat "$pidfile), $(cat "$ROOT/.idle.minutes" 2>/dev/null) min timeout)" + else + say "${C_CYAN}Idle-watch:${C_OFF} ${C_DIM}off${C_OFF} (enable: ./vllmctl idle-watch on [minutes])" + fi +} +``` + +**Modify `idle_loop()` — CRITICAL FIX:** + +```bash +idle_loop() { + local timeout_min="$1" + local timeout_s=$(( timeout_min * 60 )) + local last_active=0 + say "[idle-watch] started: will auto-sleep after ${timeout_min} min without requests (poll 30s)" + + # Use internal API for sleep calls + local SLEEP_API="http://127.0.0.1:8001" + + while true; do + sleep 30 + + # Check if container is running at all + local st + st="$(container_state)" + case "$st" in + absent*|exited*|dead*) + # Container not running - nothing to do + last_active=0 + continue + ;; + esac + + # Check if already sleeping + local sleeping + sleeping="$(curl -fs "$SLEEP_API/is_sleeping" 2>/dev/null)" + if [ "$sleeping" = "true" ]; then + # Already sleeping, nothing to do + continue + fi + + # Check for activity using metrics + local idle_s='' run_wait=0 + idle_s="$(curl -fs -m 5 "$SLEEP_API/metrics" 2>/dev/null \ + | awk '/^vllm:time_since_last_request_seconds/ {print $2; exit}')" + + if [ -z "$idle_s" ] || [ "$idle_s" = "+Inf" ] || [ "$idle_s" = "NaN" ]; then + # Fall back: count running/waiting requests + run_wait="$(curl -fs -m 5 "$SLEEP_API/metrics" 2>/dev/null \ + | awk '/^vllm:num_requests_(running|waiting)/ {s+=$2} END {print s+0}')" + if [ "${run_wait:-0}" != "0" ]; then + last_active="$(date +%s)" + continue + fi + [ "$last_active" = 0 ] && last_active="$(date +%s)" + idle_s=$(( $(date +%s) - last_active )) + fi + + # Integer compare (strip possible decimals) + local idle_i=${idle_s%%.*} + if [ -n "$idle_i" ] && [ "$idle_i" -ge "$timeout_s" ] 2>/dev/null; then + say "[idle-watch] idle for ${idle_i}s ≥ ${timeout_s}s → putting model to sleep" + if curl -fs -X POST "$SLEEP_API/sleep?level=2" >/dev/null 2>&1; then + say "[idle-watch] model is now sleeping at $(date '+%F %T')" + else + warn '[idle-watch] sleep request failed - check if server is responsive' + fi + last_active=0 + fi + done +} +``` + +**Rationale:** +- Uses internal API (`:8001`) for all calls +- Checks if container is running before attempting API calls +- Checks if already sleeping (avoid redundant sleep calls) +- Calls `/sleep?level=2` instead of `dcompose down` +- Graceful error handling if server is unresponsive + +### 1.3 Update `.env` + +No changes needed — everything is in `compose.yml` and `vllmctl`. + +--- + +## Phase 2: Add nginx Front Door + +### 2.1 Create `nginx.conf` + +```nginx +events { + worker_connections 1024; +} + +http { + # Rate limiting (optional, can be commented out) + limit_req_zone $binary_remote_addr zone=api_limit:10m rate=10r/s; + + # Upstream using Docker service name - CRITICAL FIX + upstream vllm { + server vllm:8000; # Docker service name, not 127.0.0.1 + keepalive 32; + } + + server { + listen 8000; + server_name _; + + # Block dev endpoints — deny all is sufficient + location /sleep { + deny all; + } + location /wake_up { + deny all; + } + location /is_sleeping { + deny all; + } + location /collective_rpc { + deny all; + } + location /reset_prefix_cache { + deny all; + } + + # Main API proxy + location / { + limit_req zone=api_limit burst=20 nodelay; + + proxy_pass http://vllm; + proxy_http_version 1.1; + + # Headers for WebSocket/streaming support + proxy_set_header Upgrade $http_upgrade; + proxy_set_header Connection "upgrade"; + proxy_set_header Host $host; + proxy_set_header X-Real-IP $remote_addr; + proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; + proxy_set_header X-Forwarded-Proto $scheme; + + # Timeouts for LLM inference (increased from v1) + proxy_read_timeout 900s; + proxy_send_timeout 900s; + proxy_connect_timeout 10s; + + # Disable buffering for streaming + proxy_buffering off; + } + + # Health endpoint (for health checks, allows unhealthy during sleep) + location /health { + proxy_pass http://vllm/health; + access_log off; + # Don't fail if unhealthy - vLLM can be sleeping + proxy_next_upstream error timeout http_502 http_503 http_504; + } + + # Metrics endpoint (WARNING: exposed without auth) + # Consider adding authentication if this becomes publicly accessible + location /metrics { + proxy_pass http://vllm/metrics; + } + } +} +``` + +**Rationale:** +- Fixed: `server vllm:8000` uses Docker service name (was `127.0.0.1:8001` in v1) +- Removed redundant `return 403` after `deny all` +- Increased timeouts to 900s for long-running inference +- Added `proxy_next_upstream` for health endpoint tolerance +- Added warning comment about `/metrics` exposure + +### 2.2 Update `compose.yml` + +Add nginx service with CRITICAL FIX to dependency: + +```yaml +services: + nginx: + image: nginx:alpine + container_name: vllm-nginx + restart: unless-stopped + ports: + - "${VLLM_PORT:-8000}:8000" + volumes: + - ./nginx.conf:/etc/nginx/nginx.conf:ro + # CRITICAL FIX: Use service_started, not service_healthy + # vLLM healthcheck fails during sleep, but that's OK + depends_on: + vllm: + condition: service_started # Changed from service_healthy + networks: + - default + + vllm: + # ... (as shown in Phase 1.1) +``` + +**Rationale:** +- **CRITICAL FIX**: Changed from `service_healthy` to `service_started` +- This allows nginx to start even when vLLM healthcheck fails (during sleep) +- nginx will still return 502 if vLLM is completely down, which is correct behavior + +### 2.3 Port Binding Summary + +After all changes, the port layout is: + +| From | To | Who Can Access | Purpose | +|-----|-----|----------------|---------| +| `0.0.0.0:8000` | nginx:8000 | Everyone | Public API | +| nginx | vllm:8000 | Docker only | nginx→vLLM | +| `127.0.0.1:8001` | vllm:8000 | Localhost only | Admin (`vllmctl`) | + +**Security Model:** +- External clients hit nginx on port 8000 +- nginx blocks `/sleep`, `/wake_up`, `/is_sleeping`, etc. +- `vllmctl` uses localhost:8001 to bypass nginx and access admin endpoints +- vLLM containers can talk to each other on Docker network at vllm:8000 + +--- + +## Phase 3: Testing & Validation + +### 3.1 Test Sleep Mode + +```bash +# 1. Start model +./vllmctl up Qwen3.6-35B-A3B + +# 2. Verify serving through nginx +curl http://localhost:8000/v1/models +# Should show: {"object":"list","data":[{"id":"Qwen3.6-35B-A3B",...}]} + +# 3. Put to sleep +./vllmctl down +# Should show: "Model is sleeping. Server still running." + +# 4. Check GPU memory freed +nvidia-smi +# GPU memory should be significantly lower + +# 5. Verify external API blocked - dev endpoints +curl http://localhost:8000/sleep +# Should return: 403 Forbidden + +# 6. Verify internal API works +curl http://localhost:8001/is_sleeping +# Should return: true + +# 7. Test public API during sleep +curl http://localhost:8000/v1/models +# Should return error or timeout (model is sleeping) + +# 8. Wake up +./vllmctl up +# Should wake up the existing model + +# 9. Verify serving again +curl http://localhost:8000/v1/models +# Should work again +``` + +### 3.2 Test Idle Watcher + +```bash +# Enable short idle timeout for testing +./vllmctl idle-watch on 1 + +# Wait 1 minute, then check status +./vllmctl status +# Should show: "model sleeping" + +# Test that wake-up works +./vllmctl up +./vllmctl status +# Should show: serving the model + +# Turn off when done +./vllmctl idle-watch off +``` + +### 3.3 Test Error Handling + +```bash +# 1. Stop vLLM container +docker compose stop vllm + +# 2. Try request through nginx +curl http://localhost:8000/v1/models +# Should get: 502 Bad Gateway + +# 3. Start vLLM again +docker compose start vllm + +# 4. Verify recovery +curl http://localhost:8000/v1/models +# Should work again +``` + +### 3.4 Test In-Flight Request Handling + +**Note:** vLLM Sleep Mode will allow in-flight requests to complete before sleeping. Test: + +```bash +# Start a long-running request in background +curl -N http://localhost:8000/v1/completions \ + -H "Content-Type: application/json" \ + -d '{"model":"Qwen3.6-35B-A3B","prompt":"Tell me a long story","max_tokens":500}' & +CURL_PID=$! + +# Immediately put to sleep +./vllmctl down + +# The request should complete (may take a moment) +wait $CURL_PID +echo "Request completed with exit code: $?" +``` + +--- + +## Phase 4: Optional Enhancements + +### 4.1 Protect /metrics Endpoint + +If `/metrics` exposure is a concern: + +```nginx +location /metrics { + # Optional: basic auth + # auth_basic "Metrics"; + # auth_basic_user_file /etc/nginx/.htpasswd; + + # Or restrict to localhost + # allow 127.0.0.1; + # deny all; + + proxy_pass http://vllm/metrics; +} +``` + +### 4.2 Friendly Error When Model Is Sleeping + +Add Lua-based checking (requires nginx with Lua module): + +```nginx +# This requires nginx-lua or OpenResty - can add later if needed +``` + +### 4.3 Request Logging + +Enable access logging for debugging: + +```nginx +http { + access_log /var/log/nginx/access.log; + error_log /var/log/nginx/error.log; + # ... +} +``` + +And mount log volume in compose.yml: + +```yaml +nginx: + volumes: + - ./nginx.conf:/etc/nginx/nginx.conf:ro + - ./logs/nginx:/var/log/nginx +``` + +--- + +## Migration Steps + +1. **Backup current setup:** + ```bash + cp compose.yml compose.yml.bak + cp vllmctl vllmctl.bak + ``` + +2. **Update `compose.yml`:** + - Add nginx service + - Update vLLM service with new ports and environment + - Add `--enable-sleep-mode` to EXTRA_ARGS + +3. **Create `nginx.conf`:** + - Copy the config above + - Adjust as needed + +4. **Update `vllmctl`:** + - Change API URL to `:8001` + - Modify `cmd_down()`, `cmd_up()`, `cmd_status()`, `idle_loop()` + - Use the code blocks from Phase 1.2 + +5. **Test:** + ```bash + docker compose down + docker compose up -d + ./vllmctl status + ``` + +6. **Rollback if needed:** + ```bash + docker compose down + cp compose.yml.bak compose.yml + cp vllmctl.bak vllmctl + docker compose up -d + ``` + +--- + +## Security Considerations + +1. **Dev endpoints are protected** — nginx blocks `/sleep`, `/wake_up`, `/is_sleeping`, `/collective_rpc`, `/reset_prefix_cache` +2. **Internal access only** — `vllmctl` uses port 8001 which is bound to localhost only +3. **Container isolation** — vLLM and nginx are in the same Docker network +4. **WARNING: /metrics is exposed** — Consider adding authentication if deployment becomes public +5. **VLLM_SERVER_DEV_MODE=1** — This enables dev endpoints; ensure nginx properly blocks them + +--- + +## Performance Impact + +- **nginx overhead**: Minimal (~1-2ms latency, sub-1% CPU) +- **Sleep Mode wake**: ~5-15s for level 2 with 35B model (realistic estimate) +- **Memory**: Level 2 sleep uses minimal CPU RAM, frees ~90% GPU memory +- **In-flight requests**: vLLM allows completion before sleep + +--- + +## Open Questions (Resolved) + +1. ✅ **Should idle watcher call `/sleep` directly?** + - **RESOLVED**: Yes, updated `idle_loop()` to call `/sleep?level=2` instead of `dcompose down` + +2. **Should we optimize model switching?** + - Current: Restart container on model switch (clean, reliable) + - Could optimize later: Use `/sleep` + `/wake_up` + `reload_weights` for faster switches + - Decision: Start with restart, optimize later if needed + +3. ⚠️ **What about `/metrics` endpoint?** + - Currently accessible through nginx (no auth) + - **Decision**: Leave open for now, add auth later if deployment becomes public + +--- + +## Estimated Effort (Updated) + +- Phase 1 (vLLM Sleep Mode): ~45 minutes (increased due to more comprehensive changes) +- Phase 2 (nginx): ~30 minutes +- Phase 3 (Testing): ~45 minutes (more thorough testing) +- **Total**: ~2 hours (updated from 1.5 hours) + +--- + +## Next Steps + +Once you approve this revised plan, I'll: +1. Create the modified files (`compose.yml`, `nginx.conf`, `vllmctl`) +2. Test in a worktree or provide step-by-step instructions +3. Document rollback procedure + +**All critical issues from the first review have been addressed.** + +Ready to proceed? diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..a70293a --- /dev/null +++ b/.env.example @@ -0,0 +1,9 @@ +# Copy to .env and fill in. Only these three are read by compose.yml; +# per-model serving args live in compose.yml itself. +HF_TOKEN= +VLLM_VERSION=v0.27.1 +MODEL_ROOT=/data/home/renbaibing/huggingface + +# Optional client-side convenience (open-deep-research etc.) +OPENAI_API_KEY=EMPTY +OPENAI_BASE_URL=http://127.0.0.1:8000/v1 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..9f3912f --- /dev/null +++ b/.gitignore @@ -0,0 +1,19 @@ +# secrets / local env — .env.example is the template +.env +.user.env + +# Claude local settings (memory/ is committed, settings are not) +.claude/settings.local.json + +# ephemera +*.bak +*.log +.idle.* +router-data/ +GPU-report.local.txt + +# python +__pycache__/ +*.pyc +.pytest_cache/ +.venv/ diff --git a/CALIBRATION.md b/CALIBRATION.md new file mode 100644 index 0000000..339d6cd --- /dev/null +++ b/CALIBRATION.md @@ -0,0 +1,266 @@ +# Step 1 calibration report — router front door plan v3.3 + +**Date:** 2026-08-17 +**Scope:** plan §5 (GPU budget), §7 (compose), §9 step 1 (calibration). +**State at end of calibration:** `vllm-text`, `vllm-ocr`, `vllm-embed` all +**Up and awake**, serving on debug ports 8001/8002/8003. Router not built +(another work stream); its service block is defined in compose but not started. + +## 0. Headline results + +| Metric | vllm-text (TP=2) | vllm-ocr (TP=1) | vllm-embed (TP=1) | +|---|---|---|---| +| GPUs | 0+1 | 2 | 2 | +| util slice (final) | 0.85 | 0.10 | 0.25 | +| Awake GPU footprint | 67.8–70.3 GB per GPU (0.83–0.86) | 7.9 GB (0.097) | 21.0 GB (0.256) | +| Sleeping GPU residual (per process) | 5.5–5.7 GB per GPU | 3.4 GB | ~1.0 GB | +| Sleep L1 latency (first/warm) | 6.7 s / ~1.5 s | 0.8 s / 0.1 s | 6.5 s / <0.1 s | +| Sleep L2 latency (from awake) | 0.2 s | 0.1 s | 0.04 s | +| Wake from L1 — `wake_up` only | **2.5–3.8 s** | ~0.3 s | ~1.1 s | +| Wake from L1 — full sequence (with reload) | 23.4 s | 0.7 s | 4.0 s | +| Wake from L2 — full sequence (mandatory) | 22.9 s | 0.6 s | 3.3 s | +| Host RAM cost of L1 sleep (first cycle) | 39.7 GB | 2.5 GB | 20.7 GB | +| Boot time (weights → serving) | ~5 min (warm cache) | ~2.5 min | ~1.5 min | + +**No slice changes were needed.** The initial v3.3 slices all fit (§2). + +## 1. Compose config actually used + +`/data/home/renbaibing/vllm/compose.yml` — four services, nginx removed. + +- **Common to all three vLLM services:** image `vllm/vllm-openai:v0.27.1`, + `${MODEL_ROOT}:/models:ro`, `ipc: host`, `gpus: all` + explicit + `CUDA_VISIBLE_DEVICES` (the plan's one-env-var rollback mechanism), + `VLLM_SERVER_DEV_MODE=1`, `NCCL_P2P_DISABLE=1`, `HF_TOKEN`, `restart: + unless-stopped`, loopback-only debug port `127.0.0.1:800X:8000`, + `--disable-custom-all-reduce --enable-sleep-mode`, explicit + `--max-model-len` each. No healthchecks anywhere (sleep makes `/health` + semantics misleading — see §5 — and `service_healthy` gating is forbidden + by plan §7). +- **vllm-text** (GPU 0,1; debug 8001): + `/models/Qwen3.6-35B-A3B-FP8 --served-model-name Qwen3.6-35B-A3B-FP8 + --tensor-parallel-size 2 --max-model-len 262144 --gpu-memory-utilization 0.85 + --reasoning-parser qwen3 --enable-auto-tool-choice --tool-call-parser qwen3_coder + --disable-custom-all-reduce --enable-sleep-mode` +- **vllm-ocr** (GPU 2; debug 8002): + `/models/OvisOCR2 --served-model-name OvisOCR2 --tensor-parallel-size 1 + --max-model-len 32768 --gpu-memory-utilization 0.10 --max-num-seqs 256 + --disable-custom-all-reduce --enable-sleep-mode` + (`--max-num-seqs 256` is the one added flag — see §3.) +- **vllm-embed** (GPU 2; debug 8003): + `/models/Qwen3-Embedding-8B --served-model-name Qwen3-Embedding-8B + --tensor-parallel-size 1 --max-model-len 8192 --gpu-memory-utilization 0.25 + --disable-custom-all-reduce --enable-sleep-mode` + (no task flag needed — auto-detected, see §3.) +- **router** (defined, NOT started): `build: ./router`, `8000:8000` public + + `127.0.0.1:8010:8010` admin, `depends_on` the three vLLM services with + `condition: service_started`. + +## 2. GPU budget verification (plan §5) + +Per-GPU sums, all awake (measured `nvidia-smi`, 81920 MiB per A800): + +| GPU | Contents | Slice sum | Measured used | Fraction | +|---|---|---|---|---| +| 0 | text (TP worker 0) | 0.85 + ctx | 67779–70249 MiB | 0.827–0.857 | +| 1 | text (TP worker 1) | 0.85 + ctx | 67779–70249 MiB | 0.827–0.857 | +| 2 | ocr 0.10 + embed 0.25 | 0.35 + 2 ctx | 28929 MiB | 0.353 | + +All ≤ 0.88 target. GPU0/1 peak at fresh boot 70249 MiB (0.857) — inside +budget with ~11.7 GB headroom. Sleeping (all processes L1/L2): GPU0/1 → +5.7 GB each, GPU2 → 4.5 GB total. + +Text KV cache (from logs): **Available KV cache 41.72 GiB per GPU, GPU KV +cache 4,309,772 tokens, "Maximum concurrency for 262,144 tokens per request: +16.44x"**, CUDA graphs 1.62 GiB, weights 17.48 GiB/GPU (FP8). Far above the +~20 GB concern threshold — the plan's KV-starvation worry is fully retired: +full 262K context served with 16x concurrency. (vLLM notes util 0.85 with +CUDA-graph profiling ≈ 0.8316 effective; no action needed.) + +OCR KV: 2.11 GiB → 173,056 tokens → 5.28x at 32K. Embed: pooling model +(KV small); max_model_len 8192 as planned. + +No OOM at any point; no slice adjustments made. + +## 3. Model-specific flags discovered + +**OvisOCR2 (vllm-ocr):** +- Architecture `Qwen3_5ForConditionalGeneration` is natively registered in + v0.27.1 → **no `--trust-remote-code` needed** (left out). +- dtype: `auto` resolves to bfloat16 from config; no flag needed. +- Vision limits: defaults fine (tested with a 640x160 PNG via + `image_url` base64 — correct OCR text back). No `--limit-mm-per-prompt`. +- The README's `gdn_prefill_backend="triton"` (written for vllm 0.22.1) was + NOT needed; v0.27.1 resolves the GDN backend automatically (`auto`). +- **Required addition: `--max-num-seqs 256`.** Without it boot fails: + `max_num_seqs (1024) exceeds available Mamba cache blocks (316)` — the + hybrid GDN model needs one Mamba cache block per decode sequence and the + 0.10 slice only fits 316. 256 < 316; OCR concurrency of 256 is far beyond + any realistic page-parsing load. (Alternative would have been raising util; + unnecessary.) + +**Qwen3-Embedding-8B (vllm-embed):** +- **No flag needed.** v0.27.1 auto-detects: "Found pooling configuration" + (reads sentence-transformers `modules.json`), logs + `--runner auto → pooling`, `--convert auto → embed`, and builds + `PoolerConfig(seq_pooling_type='LAST', use_activation=True)` — i.e. exactly + Qwen3-Embedding semantics (last-token pooling + normalized output). + Verified: `/v1/embeddings` returns dim=4096, norm=1.0000. +- Note for the router team: v0.27.1 has **no `--task` flag** (it was replaced + by `--runner`/`--convert`). If explicitness is ever wanted, the equivalent + is `--runner pooling --convert embed`; do NOT pass `--task embed` (unknown + arg → boot failure). + +**Qwen3.6-35B-A3B-FP8 (vllm-text):** booted with plan args verbatim; +fp8 block quantization auto-detected (`quantization=fp8`), runs on A800 +(Ampere) via the fp8/marlin path, no flags required. + +## 4. Sleep/wake measurements (per service) + +All numbers from the loopback debug ports; wake "full sequence" = +`POST /wake_up` → `POST /collective_rpc {"method":"reload_weights"}` → +`POST /reset_prefix_cache` → `GET /health` 200. + +### vllm-text +| Step | Latency | +|---|---| +| sleep L1 (first, allocates pinned pool) | 6.69 s | +| sleep L1 (warm, pool reused) | ~1.5 s | +| GPU after L1 | 5477 MiB per GPU (from 70249) | +| RAM delta after first L1 | −39.7 GB (weights 35 GB + buffers) | +| wake L1, `wake_up` only | 2.47 s (earlier run 3.7 s) | +| wake L1, full sequence | 23.42 s (reload_weights alone 20.94 s) | +| sleep L2 from awake | 0.22 s; GPU 5717 MiB; RAM −0.1 GB | +| wake L2, full sequence | 22.94 s (wake_up 1.75 + reload 21.17) | +| L1→L2 escalation | HTTP 200 in 0.01 s, **no RAM freed**, re-wake 23.7 s | + +### vllm-ocr +| Step | Latency | +|---|---| +| sleep L1 | 0.82 s; GPU 8305→3459 MiB; RAM −2.5 GB | +| wake L1 full sequence | 0.68 s | +| sleep L2 | 0.11 s; GPU 3459 MiB | +| wake L2 full sequence | 0.63 s | +| L1→L2 escalation | 200 in 0.01 s, RAM retained, re-wake 0.67 s | + +### vllm-embed +| Step | Latency | +|---|---| +| sleep L1 | 6.51 s (D2H copy); GPU 20996→~1050 MiB (GPU2 total 8981 with OCR awake); RAM −20.7 GB | +| wake L1 full sequence | 4.00 s (wake_up 1.06 + reload 2.93) | +| sleep L2 | 0.04 s | +| wake L2 full sequence | 3.31 s (reload 3.08 s from page-cached NFS) | +| L1→L2 escalation | 200 in 0.01 s, RAM retained, re-wake 4.00 s | + +### Host RAM (plan §6.3 verification) +Host total **1082 GB** (`free -g`: 1007 GiB total column, 903 GiB available +with all three awake and page cache warm). Worst case all three nap at L1 +simultaneously (first cycle after boot): 39.7 + 20.7 + 2.5 ≈ **63 GB** of +pinned host memory (plan estimated ~57 GB). With ~900 GB still available, +simultaneous L1 naps are entirely safe. Note: the pinned pools are reused +across sleep cycles — after the first L1→wake cycle, subsequent L1 sleeps +barely move `free` (allocator reuse), so don't alarm if the second nap shows +a ~0 delta. + +## 5. Router-relevant behavior findings (important) + +1. **`/health` returns 200 while a service is SLEEPING** (and `/v1/models` + too). Health is a liveness check of the API server, NOT an awake check. + The router must gate on **`/is_sleeping`**. +2. **A request sent to a sleeping service does not error — it HANGS** + (queued behind the paused scheduler; verified >2 min). Never proxy to a + backend without first checking `is_sleeping`. +3. **Requests admitted between `wake_up` and `reload_weights` get HTTP 200 + with GARBAGE content** (verified twice: OCR and text, reasoning field full + of `!!!!...`). The full wake sequence must COMPLETE before the held + request is proxied. "Never proxy to a half-awake backend" is not + theoretical — it produces silently-wrong 200s. +4. **L1 wake does NOT need `reload_weights`.** After `wake_up` alone, + text produced bit-identical output to the reloaded state (same prompt, + temperature 0, identical reasoning text; correct answers 17*23=391, + 6*7=42). `reload_weights` after L1 costs ~21 s for text and adds nothing. + → Router rule: **L1 wake = `wake_up` (+ optional reset_prefix_cache), + ~2.5–4 s; L2 wake = full sequence with `reload_weights`, ~23 s (text).** + If depth is unknown (router restart), use the full sequence. +5. **`POST /sleep?level=2` on a service already at level 1 is a well-behaved + NO-OP** (plan §6.3 "verify"): HTTP 200 in ~0.01 s, `is_sleeping` stays + true, **the L1 host-RAM copy is retained** (allocator code: level-2 + `sleep(offload_tags=())` never touches existing `cpu_backup_tensor`s), + and the next wake is still RAM-fast. It does NOT free host RAM. + → Tier escalation must be implemented as **wake (cheap, from RAM) then + `sleep?level=2` from awake**, not as a direct L1→L2 call. +6. **Concurrent `wake_up`s are safe**: 5 parallel `POST /wake_up` at a + sleeping text service all returned 200 in ~4.03 s each (serialized by the + engine). Combined with finding 3, this confirms plan §6.4's + idempotent-re-wake assumption. +7. `reset_prefix_cache` is instant (≤0.01 s) and harmless (prefix caching is + disabled for the hybrid text model anyway; enabled for embed). +8. `/sleep` also accepts `?mode=` (default `abort` discards in-flight + requests at sleep time) and level 0 (scheduler pause only, no memory + change); `wake_up` accepts `?tags=` filtering. Not needed by the router + today, but available. + +## 6. Recommended final slices (unchanged from v3.3) + +| Service | GPU(s) | util | Notes | +|---|---|---|---| +| vllm-text | 0,1 (TP=2) | **0.85** | KV 41.7 GiB/GPU, 16.4x @ 262K — ample | +| vllm-ocr | 2 | **0.10** | with `--max-num-seqs 256` | +| vllm-embed | 2 | **0.25** | 21 GB measured; ~15–16 GB is weights | + +GPU2 has ~50 GB of headroom if embed ever needs `--max-model-len` above 8192. + +## 7. Recommended router deadlines (plan §6.2.1) + +Measured: L1 wake (wake_up only) worst case ~4 s; L2 wake (full sequence) +text ~23 s warm — all files page-cached; a true cold-NFS L2 wake has not been +measured (NFS streams at ~1–6 GB/s here depending on cache; 37.5 GB cold +could add ~30–60 s). Boot from scratch: text ~5 min warm, 2–10 min cold +(NOTES-2026-08-13). + +| Situation | Plan value | Recommendation | +|---|---|---| +| Hold, wake from L1 | 30 s | **keep 30 s** (comfortable; measured ≤4 s). Could drop to 15 s. `Retry-After: 10` fine. `estimated_wake_seconds`: **5** | +| Hold, wake from L2 | 180 s | **keep 180 s** (measured 23 s warm; covers ~2 min cold-NFS margin). `Retry-After: 60` fine. `estimated_wake_seconds`: **30** (60 if NFS suspected cold) | +| Hold, container restarting | 300 s | **raise to 600 s** — cold text boot can reach 10 min; with 300 s the client gets a 503 exactly while recovery is still in progress, then must re-poll anyway. Keep `Retry-After: 600` | + +Also recommend the router implement the depth-aware fast path from §5.4 +(skip `reload_weights` for known-L1 wakes) — it turns the common case +(15-min idle tier) into a ~3 s wait for the text model instead of ~23 s. + +## 8. Incidents and resolutions during calibration + +- **Qwen3.6-35B-A3B-FP8 on disk was corrupt/incomplete** (16 layer files + wrong-sized, 18 files missing/zero-byte, tokenizer/index empty) — caused + by six stale, hung `hf download` processes from 2026-08-14 that had raced + each other on the same `--local-dir`. Resolution: sha256-verified every + file against the HF manifest, salvaged 14 intact files, re-downloaded the + remaining 24.8 GB (HuggingFace via proxy was ~1–3 MB/s; ModelScope's + mirror of the same public repo, direct + anonymous, averaged 3–6 MB/s; + every byte sha256-verified against the official HF manifest — ModelScope + was only a transport). New verified copy now lives at + `/data/home/renbaibing/huggingface/Qwen3.6-35B-A3B-FP8`; the corrupt + original is preserved at `.../Qwen3.6-35B-A3B-FP8.corrupt-20260817` + (36 GB — safe to delete once trusted). + Two stale-download casualties were also found and fixed: + `Qwen3-Embedding-8B/model-00003-of-00004.safetensors` was truncated + (re-fetched, verified) and OvisOCR2 verified fully intact. + The six stale `hf download` processes are STILL hung (killing them was + outside this session's permissions) — an operator should kill them; their + target dir no longer exists, so they are harmless but should not be + restarted as-is. +- **vllm-embed boot crash** (`SafetensorError: incomplete metadata`) — the + truncated shard above; resolved by re-download. +- **vllm-ocr boot crash** (`max_num_seqs (1024) exceeds available Mamba + cache blocks (316)`) — resolved with `--max-num-seqs 256` (kept util 0.10). + +## 9. Final functional verification (all three awake) + +- text: `POST /v1/chat/completions` → correct answers ("Paris", 391, 42), + reasoning parser active, ~146–157 tok/s decode. Tool-call/reasoning args + accepted at boot. +- ocr: chat completion with base64 PNG image → correct text extraction + ("vLLM calibration 2026-08-17 / GPU2 sleep-mode test image"), 0.4 s. +- embed: `POST /v1/embeddings` → 4096-dim unit-norm vectors. +- All three: `GET /health` 200, `GET /v1/models` lists the right model, + `GET /is_sleeping` false. diff --git a/GPU-report.txt b/GPU-report.txt new file mode 100644 index 0000000..76a085b --- /dev/null +++ b/GPU-report.txt @@ -0,0 +1,69 @@ +GPU/IOMMU fault report — v2, 2026-08-17 +(supersedes the 2026-08-14 report; the fault has MOVED, not healed) + +> Three NVIDIA A800-SXM4-80GB on CPU root ports (no PCIe switch): +> GPU0 = 0000:3d:00.0 serial 1324522063353 +> GPU1 = 0000:63:00.0 serial 1324522063362 +> GPU2 = 0000:ab:00.0 serial 1324522063359 +> Driver 595.71.05, kernel 6.8.0-137-generic (changed from 6.8.0-117 at the +> Aug-14 reboot), VBIOS 92.00.A4.00.02 on all three. + +PROBLEM + +GPU<->GPU P2P DMA corrupts data. Corruption always follows the SOURCE GPU: + + Source GPU | 2026-08-13 | 2026-08-17 (after reboot + kernel change) + -----------+------------+------------- + GPU0 | CORRUPT | CORRUPT + GPU1 | clean | CORRUPT + GPU2 | CORRUPT | clean + +The Aug-14 reboot did not fix anything: GPU0 has been a corrupt source both +times, and the failure relocated from GPU2's DMA path to GPU1's. Reproducible +test: /data/home/renbaibing/vllm/diag/p2p_check.py (run inside any CUDA +container, e.g. vllm/vllm-openai:v0.27.1). Result: 16,777,215 of 16,777,216 +elements wrong on the corrupt paths, bit-stable across runs. + +ROOT-CAUSE ASSESSMENT (platform translation path, NOT bad GPUs) + +- All hardware counters clean on all three GPUs: ECC volatile AND aggregate + zero, zero remapped rows, zero retired pages, links x16 Gen4. +- Per-GPU H2D/D2H copies and matmul PASS on all three GPUs. Only + *peer*-addressed DMA corrupts. +- Deterministic all-elements-wrong corruption = data written to wrong + addresses (translation fault), not random bit errors from failing silicon. +- VT-d is ENABLED: 14 active DMAR units, /proc/cmdline has no "iommu=pt" — + every peer write crosses VT-d translation. There are no PCIe switches, so + all P2P funnels through the root complex + IOMMU. +- Kernel evidence, all pointing at GPU0 (3d:00) as the initiator: + - Jul 6: Xid 31 CE2 MMU faults (FAULT_PDE, VIRT_WRITE) on 3d:00 + - Aug 13: DMAR: [DMA Write NO_PASID] device [3d:00.0] fault reason 0x71 + "Present bit in first-level paging entry is clear" (hundreds) +- All NVLink links report inActive on all three SXM4 boards (topo shows + PCIe-only: NODE/SYS). SXM4 boards exist for NVLink — links being down at + 3 GPUs populated is itself worth investigating. + +ASK (in priority order) + +1. Decisive test: add "intel_iommu=on iommu=pt" to the kernel command line, + reboot, re-run diag/p2p_check.py. All 6 ordered pairs must print OK. + Clean result = root cause proven (VT-d peer-DMA translation). +2. Investigate why NVLink is inactive (SBIOS setting / baseboard population / + link training). If NVLink comes up, NCCL stops using the PCIe P2P path + entirely and this bug stops mattering for our workload. +3. Check ACS on the three GPU root ports; consider SBIOS/BMC firmware and a + driver branch matched to kernel 6.8.0-137 (driver + kernel both changed at + the same reboot that relocated the fault). +4. Only if corruption survives iommu=pt: slot-swap two GPUs. Corruption + following the serial = bad board; following the slot = bad fabric. + Serials are recorded above for that comparison. + +IMPACT + +Tensor-parallel inference is unusable across GPUs without the software +workaround (NCCL_P2P_DISABLE=1 + --disable-custom-all-reduce, i.e. NCCL +stages through host memory). Measured cost at our TP=2 workload: none +(~146 tok/s). Single-GPU workloads are unaffected and verified clean. +Please re-run diag/p2p_check.py after ANY reboot or admin change and record +the matrix; also re-record "nvidia-smi --query-gpu=index,serial,pci.bus_id" +to track board vs slot. diff --git a/NOTES-2026-08-13.md b/NOTES-2026-08-13.md new file mode 100644 index 0000000..cc9180e --- /dev/null +++ b/NOTES-2026-08-13.md @@ -0,0 +1,210 @@ +# vLLM rework — 2026-08-13 + +Complete record of what was changed on this box today, why, the diagnosis of +the broken inference, and what the system admin needs to know about the GPUs. + +--- + +## 1. Starting point + +- One compose-managed container `vllm` (image `vllm/vllm-openai:latest`, + actually **v0.22.0**, built 2026-05-29), serving + `Qwen3.6-35B-A3B` (67 GB, hybrid MoE + linear-attention "GDN" model) from + NFS (`/data/home/renbaibing/huggingface`) with `--tensor-parallel-size 2` + on GPUs 0+1, port 8000, 262144 context. +- The compose file hardcoded that one model (volume + full command), so + switching models meant hand-editing YAML, and the container sat on both + GPUs permanently (`restart: unless-stopped`) even though requests are rare. +- User `renbaibing` has no docker socket access; docker rights live with + user `x640` (sudo group). Credentials for that were provided in `.user.env`. + +## 2. What was built + +### Files + +| File | Purpose | +|---|---| +| `compose.yml` | Rewritten: model-agnostic. Mounts the whole model root `/data/home/renbaibing/huggingface` read-only at `/models`; serves `/models/${MODEL_NAME}`. Image tag, TP size, context length, GPU mem fraction and extra flags all env-driven. Adds a docker healthcheck on `/health`. Includes the P2P workaround env (`NCCL_P2P_DISABLE=1`, see §4). | +| `.env` | All serving knobs: `MODEL_NAME`, `MODEL_ROOT`, `VLLM_PORT`, `TP_SIZE`, `MAX_MODEL_LEN`, `GPU_MEM_UTIL`, `EXTRA_ARGS`, `VLLM_VERSION` (image pin), `NCCL_P2P_DISABLE`. | +| `vllmctl` | Management CLI (bash). Subcommands below. Runs docker as `x640` via the pty helper. | +| `.runas.py` | Helper that runs a command as `x640` (`su` + `sudo -S`) over a pty, answering the password prompts from `.user.env`, keeping prompt text out of command output. | +| `.user.env` | `DOCKER_USER` / `DOCKER_USER_PASSWORD` for x640. **chmod 600.** | +| `README.md` | Usage guide. | +| `diag/p2p_check.py` | Pairwise GPU P2P corruption test (the decisive hardware test — §4). | +| `diag/gpu_integrity.py` | Per-GPU H2D/D2H + matmul sanity. | +| `diag/weight_check.py` | Full model tensor finiteness scan. | +| `idle-watch.log`, `.idle.pid`, `.idle.minutes` | Runtime state of the idle watcher. | + +### `vllmctl` subcommands + +``` +list models on disk + what is loaded + sizes +up [MODEL] load/serve MODEL (recreates container when switching) +down stop container → GPU memory freed in ~1 s +restart [MODEL] down + up +status container/health, served model, GPU memory, watcher +logs [-f] [N] container logs +pull HF_REPO [NAME] download a model into MODEL_ROOT (uses HF_TOKEN from .env) +idle-watch on [MIN] daemon: auto-`down` after MIN min without requests (default 15) +idle-watch off|status +``` + +The idle watcher polls `/metrics` every 30 s (`vllm:num_requests_running` / +`num_requests_waiting`, falling back to its own activity tracking on older +vLLM builds that lack `time_since_last_request_seconds`). It keeps watching +after an unload, so later manual `up`s are guarded too. + +### Docker access + +Every docker command runs as `x640` via `sudo -S` driven through the pty +helper (plain `sudo` refuses without a terminal; `sudo -S` works). If an +admin ever runs `sudo usermod -aG docker renbaibing`, the helper can be +retired by changing `dexec()` in `vllmctl` to a plain call. + +## 3. Upgrade v0.22.0 → v0.27.1 (pinned) + +- `compose.yml` now uses `vllm/vllm-openai:${VLLM_VERSION:-v0.27.1}`; + `.env` pins `VLLM_VERSION=v0.27.1` (image pulled 2026-08-13, 30.8 GB). +- Reason: v0.22's vendored FLA/GDN Triton kernels for this hybrid model have + known bug classes upstream, and v0.27.1's hybrid-cache handling is much + further along. (The NaN we saw turned out to be hardware — §4 — but the + newer image is the right baseline anyway.) +- Rollback: set `VLLM_VERSION=latest` (the old v0.22 image is still on disk) + and `./vllmctl restart`. +- Flag note: `--enable-auto-tool-choice`, `--tool-call-parser qwen3_coder`, + `--reasoning-parser qwen3` all still accepted by v0.27.1. + `VLLM_DISABLE_CUSTOM_ALL_REDUCE` is **not** a valid env var in v0.27 — + use the `--disable-custom-all-reduce` flag (in `EXTRA_ARGS`) instead. + +## 4. The broken-inference diagnosis (root cause: GPU P2P DMA corruption) + +### Symptoms after recreating the container + +- Every request returned `!!!!!…` garbage; logits contained NaN + (`Out of range float values are not JSON compliant: nan`). +- Generation crawled at ~0.25 tok/s with GPUs at "100 % util" but only + ~50 W draw (spin-wait, not compute). +- On v0.27.1, startup warmup hung indefinitely; worker processes eventually + died silently (zombies) while their GPU kernels kept spinning. + +### Hypotheses ruled out + +| Hypothesis | Evidence against | +|---|---| +| My compose changes broke it | Same image sha (`0fec7ec5…`), identical final argv in logs, same model path | +| Corrupt weights on NFS | `diag/weight_check.py`: all 26 shards, 1045 tensors, zero NaN/Inf | +| GPU compute/H2D-D2H faults | `diag/gpu_integrity.py`: all 3 GPUs pass copies + matmul; volatile ECC counters 0; no retired pages | +| vLLM compiled-kernel bug | `--enforce-eager` produced identical garbage | +| Host load / NFS | Host idle (load ~1 on 128 cores), NFS read 6 GB/s cached | + +### The decisive test — `diag/p2p_check.py` + +Host→GPU→GPU→host round-trip per ordered pair (16 M floats each): + +| Transfer | Result | +|---|---| +| GPU0→GPU1 | **corrupt — 16,777,215 / 16,777,216 elements wrong** | +| GPU1→GPU0 | clean | +| GPU1→GPU2 | clean | +| GPU2→GPU1 | **corrupt (all elements)** | +| GPU0→GPU2 | **corrupt (all elements)** | +| GPU2→GPU0 | **corrupt (all elements)** | + +Tensor parallel does an all-reduce across GPU0↔GPU1 on every layer — with +that path corrupt, activations turn to NaN and collectives hang. That single +finding explains the garbage, the 0.25 tok/s spin, and the wedged warmups. + +### Corroborating kernel-log evidence (`sudo dmesg -T`) + +- **2026-07-06**: cluster of `NVRM: Xid (PCI:0000:3d:00): 31 … MMU Fault: + ENGINE CE2 … FAULT_PDE ACCESS_TYPE_VIRT_WRITE` — copy-engine DMA faults on + **GPU0**. This coincides with when the model stopped giving good replies + (last good use: July). +- **2026-08-13 09:58:30** (during vLLM warmup): `DMAR: [DMA Write NO_PASID] + Request device [3d:00.0] fault addr 0xccfff000 [fault reason 0x71] SM: + Present bit in first-level paging entry is clear` (+ repeats, "122 + callbacks suppressed") — IOMMU rejecting GPU0 DMA writes. +- **2026-07-28**: `nvidia-persistenced` (re)started — someone already + serviced the NVIDIA stack after the July faults. +- **2026-08-11**: `nvidia 0000:ab:00.0: Using 47-bit DMA addresses` — GPU2 + was re-probed. +- `nvidia-smi nvlink -e`: *all NVLink links inActive* on all three A800s + (SXM4 boards — NVLink should normally be up). +- `nvidia-smi topo -m`: GPU0↔GPU1 connected via `NODE` (PCIe through host + bridges), no `NV#` links. + +### Workaround in place + +In `compose.yml` / `.env`: + +``` +NCCL_P2P_DISABLE=1 +EXTRA_ARGS=... --disable-custom-all-reduce +``` + +i.e. NCCL is forbidden from using direct GPU↔GPU DMA (it stages reductions +through host shared memory) and vLLM's own P2P-based custom all-reduce is +off. Performance impact is negligible here: measured **~146 tok/s** decode +with the workaround vs 0.25 tok/s with corruption. + +**Do not remove these two settings until the admin has fixed/re-verified the +machines's P2P paths** (re-run `diag/p2p_check.py` after any reboot/repair; +all pairs must report `OK`). + +## 5. Current state (end of day) + +- Container `vllm`: **v0.27.1**, healthy, serving `Qwen3.6-35B-A3B` on + GPUs 0+1, port 8000. Verified: correct completions, correct chat + + reasoning output, ~146 tok/s. +- Idle watcher **armed at 15 minutes** (auto-unloads the model when nobody + is using it; `./vllmctl idle-watch off` to disable). +- GPU2 remains untouched/free (4 MiB used). +- Weights verified intact; image pinned; rollback path documented (§3). + +## 6. For the system admin — GPU/IOMMU fault report + +> Three NVIDIA A800-SXM4-80GB (PCI `3d:00`, `63:00`, `ab:00`), driver +> 595.71.05, host up since ~2026-06-02. +> +> **Problem:** GPU↔GPU P2P DMA corrupts data on 4 of 6 ordered pairs +> (only transfers *from* GPU1 are clean). Reproducible test: +> `/data/home/renbaibing/vllm/diag/p2p_check.py` (needs torch + GPUs, e.g. +> run inside any CUDA container). Result today: 16,777,215 of 16,777,216 +> elements wrong on 0→1, 2→1, 0→2, 2→0. +> +> **Kernel evidence:** +> - Jul 6: repeated `Xid 31` CE2 MMU faults (FAULT_PDE, VIRT_WRITE) on GPU0 +> (`0000:3d:00`), pids 1461777–1538048. +> - Aug 13: `DMAR: [DMA Write NO_PASID] Request device [3d:00.0] … +> fault reason 0x71: SM: Present bit in first-level paging entry is clear` +> (hundreds of faults, "122 callbacks suppressed"). +> - All NVLinks report `inActive` (`nvidia-smi nvlink -e`); topology shows +> PCIe-only (`NODE`) between GPUs. +> - Aug 11: GPU2 (`ab:00`) was re-probed by the driver. +> - Jul 28: nvidia-persistenced was (re)started. +> +> **Impact:** tensor-parallel inference across GPUs is unusable without +> disabling P2P; we work around it in software (host-staged NCCL), at no +> measurable speed loss for our workload, but the underlying fault remains. +> +> **Ask:** investigate IOMMU/VT-d state and the PCIe fabric for GPU0 +> (and GPU2), check whether NVLink should be active on these boards, +> consider a reboot to clear IOMMU state, then re-run the p2p_check script. +> If corruption persists after reboot, it points at hardware (PCIe path or +> GPU0 itself). + +## 7. Loose ends / ideas + +- If an admin adds `renbaibing` to the `docker` group, simplify `vllmctl` + (`dexec()` → plain call) and delete `.runas.py`/`.user.env`. +- The idle watcher is a plain `nohup` process; it does **not** survive a + host reboot. If wanted permanently, wrap it in a systemd user unit. +- A wake-on-request reverse proxy (port 8000 → auto-`up` on first request) + would make unloads fully transparent to clients; not built today since + requests are rare and manual `up` takes a few minutes anyway (NFS weight + load dominates). +- `Qwen3.6-35B-A3B` officially targets Hopper-class GPUs per vLLM recipes; + it works on these A800s but has no bundled A800 MoE tuning configs + ("Using default MoE config" warning) — expect merely good, not peak, speed. +- `friendly_hertz` / `funny_shamir` are two ancient OnlyOffice containers — + unrelated, left untouched. diff --git a/NOTES-2026-08-17-gpu2-recheck.md b/NOTES-2026-08-17-gpu2-recheck.md new file mode 100644 index 0000000..f60537c --- /dev/null +++ b/NOTES-2026-08-17-gpu2-recheck.md @@ -0,0 +1,146 @@ +# GPU re-check — 2026-08-17 + +Follow-up to `NOTES-2026-08-13.md`. The admin reportedly fixed GPU2; re-ran +the same diagnostics (`diag/p2p_check.py`, `diag/gpu_integrity.py`, inside +`vllm/vllm-openai:v0.27.1`). + +## Result + +**Per-GPU integrity: all three GPUs PASS** (H2D/D2H copy fidelity, matmul +vs double-precision reference, 20× random-copy stress). + +**P2P matrix: still 4/6 paths corrupt — the fault MOVED, it did not heal.** + +| Path | 2026-08-13 | 2026-08-17 (2 runs, identical) | +|---|---|---| +| 0→1 | CORRUPT | CORRUPT | +| 0→2 | CORRUPT | CORRUPT | +| 1→0 | clean | **CORRUPT** | +| 1→2 | clean | **CORRUPT** | +| 2→0 | CORRUPT | **OK** | +| 2→1 | CORRUPT | **OK** | + +Before: only GPU1-sourced transfers were clean. Now: only GPU2-sourced +transfers are clean; everything leaving GPU0/GPU1 corrupts +(~16.7M/16.7M elements wrong, stable across runs). + +## Interpretation + +- GPU2 in isolation is healthy, and its DMA engine is now the most + trustworthy on the box. +- The PCIe/IOMMU P2P fault persists host-wide; whatever was changed + relocated the corrupting paths instead of fixing them. +- `NCCL_P2P_DISABLE=1` + `--disable-custom-all-reduce` remains mandatory for + any multi-GPU (TP>1) workload. Host-staged copies (D2H→H2D) verify clean + on all GPUs, which is why TP=2 with P2P disabled has been working. + +## Consequences for the vLLM stack + +- Single-GPU services (OCR, embedding) may run on **GPU2** safely — they + never issue P2P transfers; weights load H2D and outputs return D2H, both + verified clean. +- Recommended placement (plan v3.3): text TP=2 on GPU0+GPU1 (exclusive), + OCR + embed on GPU2. + +## For the admin + +The 2026-08-13 report asked to investigate IOMMU/VT-d and the PCIe fabric. +The new matrix shows the corrupting paths moved from {0→1, 0→2, 2→1, 2→0} +to {0→1, 0→2, 1→0, 1→2} — i.e. **all transfers sourced from GPU0 and GPU1 +now corrupt, while GPU2-sourced transfers became clean**. Reproducible: +two consecutive runs of `diag/p2p_check.py`, identical results. + +--- + +## Root-cause analysis (later on 2026-08-17) + +**The "GPU2 fixed, others broken" reading is an illusion. The GPUs are fine; +the host's peer-DMA translation path is what's broken.** + +The two matrices re-sorted by SOURCE GPU (corruption always follows the +source — `p2p_check.py` uses torch `.to()`, which issues a source-side +copy-engine DMA): + +| Source GPU | 2026-08-13 | 2026-08-17 | +|---|---|---| +| GPU0 (`3d:00`) | CORRUPT | CORRUPT | +| GPU1 (`63:00`) | clean | **CORRUPT** | +| GPU2 (`ab:00`) | CORRUPT | **clean** | + +- GPU0 is a corrupt source BOTH times. Nothing healed; the Aug-14 reboot + (which also bumped kernel 6.8.0-117 → 6.8.0-137) rebuilt the platform state + and relocated the failure from GPU2's DMA path to GPU1's. Restarts will keep + "moving" the fault like this — that is the signature of a software/platform + state problem, not dying silicon. +- Corruption signature is deterministic mistranslation, not hardware noise: + 16,777,215 / 16,777,216 elements wrong, bit-stable across runs. Flaky + hardware gives random bit errors; wrong-everywhere means wrong addresses. +- Re-verified 2026-08-17: zero ECC (volatile AND aggregate), zero remapped + rows, zero retired pages on all three GPUs; all links x16 Gen4; per-GPU + H2D/D2H + matmul pass on all three. +- **VT-d confirmed ENABLED**: 14 active DMAR units under `/sys/class/iommu`, + and `/proc/cmdline` has NO `iommu=pt` — every peer write crosses VT-d + translation. +- No PLX switches (`lspci -t`): each GPU sits on its own CPU root port, so + all P2P funnels through the root complex + IOMMU. +- The kernel evidence from 08-13 points the same way, and only ever at GPU0: + Jul 6 `Xid 31` CE2 MMU faults (VIRT_WRITE) on `3d:00`; Aug 13 `DMAR + [DMA Write] fault reason 0x71` on `3d:00`. + +**Conclusion:** platform fault in peer-DMA translation — VT-d peer-to-peer +mappings interacting with driver 595.71.05 / kernel 6.8.0-137 / SBIOS ACS +state. Reboots rebuild per-device translation contexts, which is why the bad +set moves between GPUs. + +### GPU serials (baseline for board-vs-slot tracking) + +| Bus | Serial | VBIOS | +|---|---|---| +| `0000:3d:00.0` | 1324522063353 | 92.00.A4.00.02 | +| `0000:63:00.0` | 1324522063362 | 92.00.A4.00.02 | +| `0000:ab:00.0` | 1324522063359 | 92.00.A4.00.02 | + +If a slot-swap test is ever done: corruption following the **serial** = bad +board; following the **bus/slot** = bad fabric/root port. + +## Interconnect primer: NVLink vs PCIe P2P (why this box is stuck on fallback #3) + +NVLink is dedicated GPU-to-GPU interconnect hardware (~400 GB/s per GPU on +A800), built into the SXM4 module design — in multi-GPU HGX servers it is +wired to NVSwitch chips on the baseboard. PCIe P2P DMA (~25 GB/s practical, +Gen4 x16) is the standard fallback and is *supposed to be fully correct* — +it works on the vast majority of GPU servers. + +This box is unusual twice over: + +1. **All NVLink links `inActive` on all three A800-SXM4s.** SXM4 exists for + NVLink, so "inActive" means not-up, not absent. Candidate causes: only 3 + GPUs populated in a 4/8-GPU baseboard (incomplete NVLink topology), NVLink + disabled in SBIOS, links not training (bridge/baseboard seating), or a + driver reporting quirk. **If NVLink can be brought up, NCCL uses it and + the corrupt PCIe path becomes irrelevant.** +2. **The PCIe P2P path itself corrupts** — the actual fault diagnosed above. + +So NCCL on this box sits on the third fallback: + +| Path | Bandwidth | Status on this box | +|---|---|---| +| NVLink | ~400 GB/s | links down | +| PCIe P2P DMA | ~25 GB/s | corrupts data (the fault) | +| Host-staged D2H→CPU→H2D | ~12 GB/s effective | ✅ forced by `NCCL_P2P_DISABLE=1` | + +Host-staged works because per-GPU H2D/D2H is verified clean on all three +GPUs — only *peer*-addressed DMA mistranslates. Cost at TP=2 is negligible +here (~146 tok/s measured). + +## Fix ladder for the admin (decisive test first) + +1. **GRUB: `intel_iommu=on iommu=pt`** → reboot → re-run `diag/p2p_check.py`; + all 6 ordered pairs must print `OK`. A clean result proves the root cause + is VT-d peer-DMA translation. This is the decisive experiment. +2. **Bring NVLink up** (SBIOS/baseboard investigation) — sidesteps PCIe P2P + entirely even if (1) cannot fully fix it. +3. Check ACS on the three root ports; SBIOS/BMC firmware updates; driver + branch matched to kernel 6.8.0-137. +4. Only if corruption survives `iommu=pt`: suspect PCIe fabric hardware → + slot-swap two GPUs and compare against the serial table above. diff --git a/README.md b/README.md new file mode 100644 index 0000000..700d757 --- /dev/null +++ b/README.md @@ -0,0 +1,156 @@ +# vLLM model server — router front door, auto wake-on-request + +> History: 2026-08-13 rework + GPU-fault diagnosis (`NOTES-2026-08-13.md`), +> GPU2 re-check + root cause (`NOTES-2026-08-17-gpu2-recheck.md`), +> measured GPU/wake numbers (`CALIBRATION.md`). +> Hardware sanity scripts live in **`diag/`** (`p2p_check.py` is the key one). + +Three models are **always available** at one OpenAI-compatible endpoint. +Calling services never need to know whether a model is loaded — the router +front door wakes it on request: + +``` +http://:8000/v1/... +``` + +Everything runs inside the Docker stack. No host-side cron, watchers, or +docker commands on the request path. + +## Models + +| Type | Model name (use in `model` field) | Service / GPUs | +|---|---|---| +| Text / chat | `Qwen3.6-35B-A3B-FP8` (aliases: `qwen`, `text`, …) | `vllm-text`, TP=2, GPU0+1 | +| OCR (vision) | `OvisOCR2` (aliases: `ocr`, `ovis`) | `vllm-ocr`, GPU2 | +| Embeddings | `Qwen3-Embedding-8B` (aliases: `embed`, …) | `vllm-embed`, GPU2 | + +`GET /v1/models` lists all three. Model matching is case-insensitive; an +unknown name returns a 404 `model_not_found` (no wake is triggered). + +## What a request sees + +- **Awake model** → proxied immediately (streaming passes through unbuffered). +- **Sleeping model** → the request is *held* while the router wakes it + (single-flight: 10 concurrent requests → one wake), then proxied. +- **Wake exceeds the hold deadline** → `503` with a depth-calibrated + `Retry-After` header AND an OpenAI-style JSON body: + +| Waking from | Typical wake | Hold deadline | `Retry-After` | body `sleep_depth` | +|---|---|---|---|---| +| level-1 sleep (weights in RAM) | 2.5–4 s (text), <1 s (small) | 30 s | `10` | `"sleeping"` | +| level-2 offload (weights on NFS) | ~23 s (text), 1–4 s (small) | 180 s | `60` | `"offloaded"` | +| container restart / cold boot | 2–10 min (text, NFS) | 600 s | `600` | `"restarting"` | + +```json +{"error": {"type": "model_waking", "code": "model_waking", + "message": "Model '…' is waking from offload; retry shortly", + "sleep_depth": "offloaded", "estimated_wake_seconds": 60}} +``` + +## Idle management (automatic) + +| After idle | Action | Effect | +|---|---|---| +| 15 min | sleep (level 1) | GPU freed; weights parked in host RAM | +| 3 h | offload (level 2) | host RAM freed too; ~5 GB/GPU CUDA context remains | + +First request after either tier wakes the model transparently. Timers: +`IDLE_SLEEP_MIN` / `IDLE_OFFLOAD_MIN` env on the router service. + +## Quick start + +```bash +curl http://127.0.0.1:8000/v1/chat/completions -H 'Content-Type: application/json' \ + -d '{"model":"Qwen3.6-35B-A3B-FP8","messages":[{"role":"user","content":"hi"}]}' + +curl http://127.0.0.1:8000/v1/embeddings -H 'Content-Type: application/json' \ + -d '{"model":"Qwen3-Embedding-8B","input":"hello"}' +``` + +### vllmctl + +```bash +./vllmctl status # per-model state table (sleep depth, activity) +./vllmctl up [MODEL] # force-wake (HTTP to the admin API — no docker) +./vllmctl sleep [MODEL] # level-1 sleep now +./vllmctl down [MODEL|all] # level-2 offload now +./vllmctl list # models on disk + served +./vllmctl logs [-f] [N] [SVC] +./vllmctl restart SVC +./vllmctl pull # download a new model (needs docker) +``` + +`up`/`down`/`sleep`/`status` are pure HTTP against the admin listener; +routine control never needs docker. + +## Architecture + +``` +callers ──► :8000 ──► router (FastAPI, ./router) [only public ingress] + ├─► vllm-text GPU0+1 (TP=2, util 0.85, 262K ctx) + ├─► vllm-ocr GPU2 (TP=1, util 0.10) + └─► vllm-embed GPU2 (TP=1, util 0.25) +127.0.0.1:8010 ──► router admin API (vllmctl; not reachable off-host) +127.0.0.1:8001-8003 ──► direct debug access to each vLLM (local users only) +``` + +- Each vLLM service is always running with sleep mode + (`--enable-sleep-mode`, `VLLM_SERVER_DEV_MODE=1`). The router owns + wake/sleep/idle, model routing, streaming proxy, and endpoint hiding + (dev endpoints and `/admin/*` 404 on the public port; paths are + normalized against traversal). +- The router is a single process with two listeners (public + admin) on one + event loop, so admin calls share the request path's locks and depth state. +- Wake-intent is persisted (named volume `router-state`), so a router + restart mid-wake conservatively completes `reload_weights` before serving + — an interrupted level-2 wake can never serve garbage. +- All services: `restart: unless-stopped` (host reboots self-heal; requests + get `"restarting"` 503s during the gap). + +### Router tunables + +All env-driven with sane defaults — see the table at the top of +`router/config.py` (registry/aliases, listeners, proxy timeouts, wake hold +deadlines, idle tiers, state file). Unit tests: `router/tests/` +(91 tests, no GPU needed). + +## Hardware constraints (do not remove) + +- GPU P2P DMA **corrupts data** on this host (4/6 pairs; see NOTES). + `NCCL_P2P_DISABLE=1` + `--disable-custom-all-reduce` are mandatory on + every vLLM service. Host-staged copies are verified clean. +- Text spans GPU0+GPU1 exclusively; OCR and embed live on GPU2 (single-GPU + services never issue P2P). Moving a service = edit its + `CUDA_VISIBLE_DEVICES` in `compose.yml`, then re-tighten memory slices + (per-GPU sums ≤ 0.88 including CUDA contexts). +- Measured footprints in `CALIBRATION.md`. + +## How docker access works + +`renbaibing` is in the `docker` group (since 2026-08-14) — `docker` and +`docker compose` run directly, no sudo. Shells opened **before** the group +change need `sg docker -c "…"` or a fresh login. The old `.runas.py` / +`.user.env` mechanism is obsolete. + +## Operations + +```bash +docker compose up -d # whole stack +docker compose restart router # front door only (state survives) +docker logs vllm-router # structured wake/sleep/idle events +curl 127.0.0.1:8010/admin/status +``` + +The pre-router (nginx) stack was removed 2026-08-17 (before `git init`) — +it survives only as a design doc +(`.claude/memory/sleep-mode-implementation-plan.md`). Rollback points are +git commits from 2026-08-17 onward. + +## Notes + +- `.env` provides `VLLM_VERSION`, `MODEL_ROOT`, `HF_TOKEN` (used by + compose) plus client-side vars (`OPENAI_BASE_URL` etc.). All per-model + serving args live in `compose.yml`. +- A `.corrupt-20260817` copy of the FP8 model sits next to the repaired + original — delete it once the repaired copy is trusted (43 GB). +- Design/decision record: `.claude/memory/router-front-door-plan.md` (v3.3). diff --git a/TODO.md b/TODO.md new file mode 100644 index 0000000..0aedd23 --- /dev/null +++ b/TODO.md @@ -0,0 +1,53 @@ +# Router Front Door — implementation record + +**Completed:** 2026-08-17 · **Status:** implemented, E2E-verified, in service ✅ + +Supersedes the Sleep Mode + nginx plan (v2, 2026-08-14) — that stack worked +but did not wake models on request; the front door's real job. Design record: +`.claude/memory/router-front-door-plan.md` (v3.3). + +## What was built + +| Piece | What it does | +|---|---| +| `router/` (FastAPI) | Only public ingress (:8000). Model routing, wake-on-request (single-flight), streaming proxy, tiered idle sleep/offload, depth-aware 503s, dev-endpoint hiding, persisted wake-intent recovery. Admin API on 127.0.0.1:8010. | +| `compose.yml` | 4 services: `vllm-text` (TP=2 GPU0+1, util 0.85, 262K ctx), `vllm-ocr` (GPU2, 0.10, `--max-num-seqs 256`), `vllm-embed` (GPU2, 0.25), `router`. Named volume `router-state` for wake-intent. nginx removed. | +| `vllmctl` | `status/up/down/sleep/list/logs/restart/pull`. Routine control is pure HTTP to the admin API — no docker, no credentials. | +| `CALIBRATION.md` | Measured GPU footprints, sleep residuals, wake latencies (L1 fast path ~2.5–4 s text; L2 ~23 s), host RAM, backend quirks (`/health` lies while asleep; L1 wake needs `wake_up` only; L1→L2 re-sleep is a no-op → wake-then-sleep). | + +## Verification trail + +- Plan review: 2 independent rounds, all critical/major issues fixed. +- Router unit tests: 91 passed (no GPU needed) — `router/tests/`. +- E2E: 15-test matrix on real hardware — 14 PASS, 1 FAIL (router restart in + the wake_up→reload_weights window served garbage). +- Fix: persisted wake-intent + startup recovery; re-verified PASS on the + real stack (503 `router_shutting_down` for in-flight, clean content after + recovery, backend logs show reload precedes any proxy). +- Deployment fix: state volume switched from bind mount (uid mismatch, + silently unwritable) to a named volume seeded from image ownership. + +## Known limitations (accepted) + +- Host RAM grows by the weights size while a model naps at level 1 + (~63 GB if all three nap; host has 1082 GB). +- An NFS outage makes level-2-offloaded models unwakeable until NFS returns + (router 503s meanwhile). Optional mitigation: cap text at level 1. +- A stream already in flight when the router shuts down is truncated + (client sees EOF, not a 503) — unavoidable once bytes are sent. +- vLLM's chat endpoint is JSON-only: multipart is routed correctly but real + OCR calls use JSON + base64 images. +- Host reboot: all models cold-boot awake (vLLM cannot boot asleep); + requests get `"restarting"` 503s for the load minutes. + +## Housekeeping + +- [ ] Delete `Qwen3.6-35B-A3B-FP8.corrupt-20260817` (43 GB, in MODEL_ROOT) + once the repaired FP8 copy is trusted. +- [x] Old-stack leftovers removed 2026-08-17 before `git init`: nginx.conf, + `*.v2.bak`, `.runas.py`, `.user.env`, idle-watcher files. The v2 stack + survives only as design docs (`sleep-mode-implementation-plan.md`). +- [ ] Confirm calling services' client timeouts ≥ ~60 s (offload wake). +- [ ] Ask the admin to run the GPU fix ladder (`iommu=pt` decisive test) — + see `NOTES-2026-08-17-gpu2-recheck.md`. If P2P is ever fixed, + `NCCL_P2P_DISABLE` can be revisited (perf only; not required). diff --git a/claude_code_env.sh b/claude_code_env.sh new file mode 100644 index 0000000..de40f60 --- /dev/null +++ b/claude_code_env.sh @@ -0,0 +1,205 @@ +#!/bin/bash + +set -euo pipefail + +# ======================== +# 常量定义 +# ======================== +SCRIPT_NAME=$(basename "$0") +NODE_MIN_VERSION=18 +NODE_INSTALL_VERSION=22 +NVM_VERSION="v0.40.3" +CLAUDE_PACKAGE="@anthropic-ai/claude-code" +CONFIG_DIR="$HOME/.claude" +CONFIG_FILE="$CONFIG_DIR/settings.json" +API_BASE_URL="https://open.bigmodel.cn/api/anthropic" +API_KEY_URL="https://open.bigmodel.cn/usercenter/proj-mgmt/apikeys" +API_TIMEOUT_MS=3000000 + +# ======================== +# 工具函数 +# ======================== + +log_info() { + echo "🔹 $*" +} + +log_success() { + echo "✅ $*" +} + +log_error() { + echo "❌ $*" >&2 +} + +ensure_dir_exists() { + local dir="$1" + if [ ! -d "$dir" ]; then + mkdir -p "$dir" || { + log_error "Failed to create directory: $dir" + exit 1 + } + fi +} + +# ======================== +# Node.js 安装函数 +# ======================== + +install_nodejs() { + local platform=$(uname -s) + + case "$platform" in + Linux|Darwin) + log_info "Installing Node.js on $platform..." + + # 安装 nvm + log_info "Installing nvm ($NVM_VERSION)..." + curl -s https://raw.githubusercontent.com/nvm-sh/nvm/"$NVM_VERSION"/install.sh | bash + + # 加载 nvm + log_info "Loading nvm environment..." + \. "$HOME/.nvm/nvm.sh" + + # 安装 Node.js + log_info "Installing Node.js $NODE_INSTALL_VERSION..." + nvm install "$NODE_INSTALL_VERSION" + + # 验证安装 + node -v &>/dev/null || { + log_error "Node.js installation failed" + exit 1 + } + log_success "Node.js installed: $(node -v)" + log_success "npm version: $(npm -v)" + ;; + *) + log_error "Unsupported platform: $platform" + exit 1 + ;; + esac +} + +# ======================== +# Node.js 检查函数 +# ======================== + +check_nodejs() { + if command -v node &>/dev/null; then + current_version=$(node -v | sed 's/v//') + major_version=$(echo "$current_version" | cut -d. -f1) + + if [ "$major_version" -ge "$NODE_MIN_VERSION" ]; then + log_success "Node.js is already installed: v$current_version" + return 0 + else + log_info "Node.js v$current_version is installed but version < $NODE_MIN_VERSION. Upgrading..." + install_nodejs + fi + else + log_info "Node.js not found. Installing..." + install_nodejs + fi +} + +# ======================== +# Claude Code 安装 +# ======================== + +install_claude_code() { + if command -v claude &>/dev/null; then + log_success "Claude Code is already installed: $(claude --version)" + else + log_info "Installing Claude Code..." + npm install -g "$CLAUDE_PACKAGE" || { + log_error "Failed to install claude-code" + exit 1 + } + log_success "Claude Code installed successfully" + fi +} + +configure_claude_json(){ + node --eval ' + const os = require("os"); + const fs = require("fs"); + const path = require("path"); + + const homeDir = os.homedir(); + const filePath = path.join(homeDir, ".claude.json"); + if (fs.existsSync(filePath)) { + const content = JSON.parse(fs.readFileSync(filePath, "utf-8")); + fs.writeFileSync(filePath, JSON.stringify({ ...content, hasCompletedOnboarding: true }, null, 2), "utf-8"); + } else { + fs.writeFileSync(filePath, JSON.stringify({ hasCompletedOnboarding: true }, null, 2), "utf-8"); + }' +} + +# ======================== +# API Key 配置 +# ======================== + +configure_claude() { + log_info "Configuring Claude Code..." + echo " You can get your API key from: $API_KEY_URL" + read -s -p "🔑 Please enter your ZHIPU API key: " api_key + echo + + if [ -z "$api_key" ]; then + log_error "API key cannot be empty. Please run the script again." + exit 1 + fi + + ensure_dir_exists "$CONFIG_DIR" + + # 写入配置文件 + node --eval ' + const os = require("os"); + const fs = require("fs"); + const path = require("path"); + + const homeDir = os.homedir(); + const filePath = path.join(homeDir, ".claude", "settings.json"); + const apiKey = "'"$api_key"'"; + + const content = fs.existsSync(filePath) + ? JSON.parse(fs.readFileSync(filePath, "utf-8")) + : {}; + + fs.writeFileSync(filePath, JSON.stringify({ + ...content, + env: { + ANTHROPIC_AUTH_TOKEN: apiKey, + ANTHROPIC_BASE_URL: "'"$API_BASE_URL"'", + API_TIMEOUT_MS: "'"$API_TIMEOUT_MS"'", + CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC: 1 + } + }, null, 2), "utf-8"); + ' || { + log_error "Failed to write settings.json" + exit 1 + } + + log_success "Claude Code configured successfully" +} + +# ======================== +# 主流程 +# ======================== + +main() { + echo "🚀 Starting $SCRIPT_NAME" + + check_nodejs + install_claude_code + configure_claude_json + configure_claude + + echo "" + log_success "🎉 Installation completed successfully!" + echo "" + echo "🚀 You can now start using Claude Code with:" + echo " claude" +} + +main "$@" diff --git a/compose.yml b/compose.yml new file mode 100644 index 0000000..1b8d810 --- /dev/null +++ b/compose.yml @@ -0,0 +1,119 @@ +# vLLM serving stack — router front door (plan v3.3, 2026-08-17). +# +# nginx is GONE: the FastAPI router (./router) is the only public ingress +# (host :8000), with a separate admin listener on 127.0.0.1:8010. +# Each model is a dedicated always-running vLLM service with Sleep Mode +# enabled; the router wakes them on request. +# +# Hardware constraints (NOTES-2026-08-13.md, NOTES-2026-08-17-gpu2-recheck.md): +# * GPU P2P DMA corrupts data host-wide (4/6 paths). NCCL_P2P_DISABLE=1 +# + --disable-custom-all-reduce are MANDATORY on every service. +# * text runs TP=2 on GPU0+GPU1 (exclusive); OCR and embed run TP=1 on +# GPU2 (single-GPU services never issue P2P transfers). +# * Moving a service to another GPU = edit its CUDA_VISIBLE_DEVICES (then +# re-tighten slices). +# +# No healthchecks on purpose: sleep mode makes /health flap, and nothing here +# may use a service_healthy dependency (plan §7). + +services: + vllm-text: + image: vllm/vllm-openai:${VLLM_VERSION:-v0.27.1} + container_name: vllm-text + restart: unless-stopped + ports: + - "127.0.0.1:8001:8000" # direct debug access (local only) + volumes: + - ${MODEL_ROOT:-/data/home/renbaibing/huggingface}:/models:ro + ipc: host + gpus: all + environment: + VLLM_SERVER_DEV_MODE: "1" # sleep-mode endpoints (/sleep /wake_up /is_sleeping ...) + HF_TOKEN: ${HF_TOKEN:-} + NCCL_P2P_DISABLE: "1" # broken P2P DMA on this host (see header) + CUDA_VISIBLE_DEVICES: "0,1" + command: > + /models/Qwen3.6-35B-A3B-FP8 + --served-model-name Qwen3.6-35B-A3B-FP8 + --tensor-parallel-size 2 + --max-model-len 262144 + --gpu-memory-utilization 0.85 + --reasoning-parser qwen3 + --enable-auto-tool-choice + --tool-call-parser qwen3_coder + --disable-custom-all-reduce + --enable-sleep-mode + + vllm-ocr: + image: vllm/vllm-openai:${VLLM_VERSION:-v0.27.1} + container_name: vllm-ocr + restart: unless-stopped + ports: + - "127.0.0.1:8002:8000" # direct debug access (local only) + volumes: + - ${MODEL_ROOT:-/data/home/renbaibing/huggingface}:/models:ro + ipc: host + gpus: all + environment: + VLLM_SERVER_DEV_MODE: "1" + HF_TOKEN: ${HF_TOKEN:-} + NCCL_P2P_DISABLE: "1" + CUDA_VISIBLE_DEVICES: "2" + command: > + /models/OvisOCR2 + --served-model-name OvisOCR2 + --tensor-parallel-size 1 + --max-model-len 32768 + --gpu-memory-utilization 0.10 + --max-num-seqs 256 + --disable-custom-all-reduce + --enable-sleep-mode + + vllm-embed: + image: vllm/vllm-openai:${VLLM_VERSION:-v0.27.1} + container_name: vllm-embed + restart: unless-stopped + ports: + - "127.0.0.1:8003:8000" # direct debug access (local only) + volumes: + - ${MODEL_ROOT:-/data/home/renbaibing/huggingface}:/models:ro + ipc: host + gpus: all + environment: + VLLM_SERVER_DEV_MODE: "1" + HF_TOKEN: ${HF_TOKEN:-} + NCCL_P2P_DISABLE: "1" + CUDA_VISIBLE_DEVICES: "2" + command: > + /models/Qwen3-Embedding-8B + --served-model-name Qwen3-Embedding-8B + --tensor-parallel-size 1 + --max-model-len 8192 + --gpu-memory-utilization 0.25 + --disable-custom-all-reduce + --enable-sleep-mode + + router: + build: ./router + container_name: vllm-router + restart: unless-stopped + ports: + - "8000:8000" # public API — only public ingress + - "127.0.0.1:8010:8010" # admin API (vllmctl / debugging), local only + volumes: + # Wake-intent state: lets a restarted router tell "awake and ready" from + # "awake because the previous one died between wake_up and reload_weights" + # (plan 6.4 / E2E case 12). Written atomically, a few hundred bytes. + # Named volume (not a bind mount): Docker seeds it with the image's + # /state ownership, so the non-root router user can always write it. + - router-state:/state + depends_on: + vllm-text: + condition: service_started + vllm-ocr: + condition: service_started + vllm-embed: + condition: service_started + +volumes: + router-state: diff --git a/diag/gpu_integrity.py b/diag/gpu_integrity.py new file mode 100644 index 0000000..9fe7321 --- /dev/null +++ b/diag/gpu_integrity.py @@ -0,0 +1,33 @@ +#!/usr/bin/env python3 +"""Per-GPU basic integrity: H2D/D2H copy fidelity, matmul correctness, +repeated-copy stress. Does NOT test P2P — use p2p_check.py for that. + +Run: docker exec vllm python3 /path/to/gpu_integrity.py +""" +import torch + +print("torch", torch.__version__, "cuda", torch.version.cuda, "devs", torch.cuda.device_count()) +for dev in range(torch.cuda.device_count()): + try: + torch.cuda.set_device(dev) + x = torch.randint(0, 2**31 - 1, (128 * 1024 * 1024,), dtype=torch.int32) + y = x.to(f"cuda:{dev}") + z = y.cpu() + copy_ok = torch.equal(x, z) + a = torch.randn(2048, 2048, dtype=torch.float32) + b = torch.randn(2048, 2048, dtype=torch.float32) + ref = (a.double() @ b.double()).float() + c = (a.to(f"cuda:{dev}") @ b.to(f"cuda:{dev}")).cpu() + err = (c - ref).abs().max().item() + stress_ok = True + for _ in range(20): + s = torch.randint(0, 2**31 - 1, (16 * 1024 * 1024,), dtype=torch.int32) + if not torch.equal(s, s.to(f"cuda:{dev}").cpu()): + stress_ok = False + break + print(f"GPU{dev}: copy_ok={copy_ok} matmul_max_err={err:.3e} " + f"stress_ok={stress_ok} name={torch.cuda.get_device_name(dev)}") + del x, y, z, a, b, c + torch.cuda.empty_cache() + except Exception as e: + print(f"GPU{dev}: ERROR {type(e).__name__}: {e}") diff --git a/diag/p2p_check.py b/diag/p2p_check.py new file mode 100644 index 0000000..03b0080 --- /dev/null +++ b/diag/p2p_check.py @@ -0,0 +1,58 @@ +#!/usr/bin/env python3 +"""GPU P2P transfer integrity check. + +Tests every ordered GPU pair: host -> src GPU -> dst GPU -> host, and +compares against the host copy. ANY mismatch means the P2P DMA path is +corrupting data and tensor-parallel inference MUST run with P2P disabled +(NCCL_P2P_DISABLE=1 and, for vLLM, --disable-custom-all-reduce). + +Run somewhere with torch + all GPUs visible, e.g.: + docker exec vllm python3 /models/../diag/p2p_check.py +or docker run --rm --gpus all -v $PWD/diag:/diag vllm/vllm-openai:v0.27.1 \ + python3 /diag/p2p_check.py +""" +import sys +import torch + + +def main(): + n = torch.cuda.device_count() + print(f"torch={torch.__version__} gpus={n}") + if n < 2: + print("need >= 2 GPUs"); return 2 + for i in range(n): + for j in range(n): + if i != j: + print(f"p2p {i}->{j} supported: {torch.cuda.can_device_access_peer(i, j)}") + x = torch.arange(16 * 1024 * 1024, dtype=torch.float32, device="cpu") + ref = x.clone() + bad_pairs = [] + for src in range(n): + for dst in range(n): + if src == dst: + continue + try: + a = x.to(f"cuda:{src}") + b = a.to(f"cuda:{dst}") + c = b.cpu() + torch.cuda.synchronize() + n_bad = int((c != ref).sum().item()) + status = "OK" if n_bad == 0 else f"CORRUPT ({n_bad}/{ref.numel()} elements wrong)" + print(f"{src}->{dst}: {status}") + if n_bad: + bad_pairs.append((src, dst)) + del a, b, c + torch.cuda.empty_cache() + except Exception as e: + print(f"{src}->{dst}: FAILED {type(e).__name__} {str(e)[:160]}") + bad_pairs.append((src, dst)) + if bad_pairs: + print(f"\nRESULT: {len(bad_pairs)} corrupt/failing P2P paths: {bad_pairs}") + print("Keep NCCL_P2P_DISABLE=1 (+ --disable-custom-all-reduce for vLLM).") + return 1 + print("\nRESULT: all P2P paths clean.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/diag/weight_check.py b/diag/weight_check.py new file mode 100644 index 0000000..a7dd874 --- /dev/null +++ b/diag/weight_check.py @@ -0,0 +1,43 @@ +#!/usr/bin/env python3 +"""Verify every tensor of a safetensors model is finite (no NaN/Inf) and +sane. Usage: python3 weight_check.py /models/ + +Run inside the vllm container (it has torch + safetensors): + docker cp diag/weight_check.py vllm:/tmp/ && \ + docker exec vllm python3 /tmp/weight_check.py /models/Qwen3.6-35B-A3B +""" +import glob +import os +import sys + +import torch +from safetensors import safe_open + + +def main(): + model_dir = sys.argv[1] if len(sys.argv) > 1 else "/models" + files = sorted(glob.glob(os.path.join(model_dir, "*.safetensors"))) + if not files: + print(f"no safetensors found in {model_dir}") + return 2 + total = bad = 0 + for f in files: + with safe_open(f, framework="pt", device="cpu") as st: + for k in st.keys(): + t = st.get_tensor(k) + total += 1 + if torch.is_floating_point(t): + if not torch.isfinite(t).all(): + print(f"BAD {os.path.basename(f)}::{k} " + f"nan={int(torch.isnan(t).sum())} inf={int(torch.isinf(t).sum())}", + flush=True) + bad += 1 + elif t.dtype != torch.bool and (t.abs() > 1e6).any(): + print(f"ODD {os.path.basename(f)}::{k} max={t.abs().max().item()}", + flush=True) + print(f"DONE shards={len(files)} tensors={total} bad={bad}") + return 1 if bad else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/router/Dockerfile b/router/Dockerfile new file mode 100644 index 0000000..9940314 --- /dev/null +++ b/router/Dockerfile @@ -0,0 +1,28 @@ +# vllm-router: single process, one event loop, two sockets (public :8000, +# admin :8010). No GPU, no model weights -- pure HTTP front door. +FROM python:3.12-slim + +# Non-root runtime user. +RUN groupadd --system --gid 10001 router \ + && useradd --system --uid 10001 --gid router --home-dir /app router + +WORKDIR /app + +# Deps first so code changes don't bust the layer. +COPY requirements.txt ./ +RUN pip install --no-cache-dir -r requirements.txt + +COPY app.py config.py services.py routing.py admin_api.py ./ + +# Wake-intent state dir: a named volume seeded from this ownership, so the +# non-root runtime user can always write it (bind mounts inherit host uids). +RUN mkdir -p /state && chown router:router /state + +USER router +EXPOSE 8000 8010 + +# Liveness only (a sleeping backend is normal, not an outage). +HEALTHCHECK --interval=30s --timeout=5s --start-period=15s --retries=3 \ + CMD ["python", "-c", "import urllib.request;urllib.request.urlopen('http://127.0.0.1:8000/health', timeout=3)"] + +CMD ["python", "app.py"] diff --git a/router/admin_api.py b/router/admin_api.py new file mode 100644 index 0000000..1009fd5 --- /dev/null +++ b/router/admin_api.py @@ -0,0 +1,143 @@ +"""Admin listener (the :8010 socket) -- `vllmctl` / debugging surface. + +This app is mounted on a *separate socket*, never on the public one (plan +review C5: a public /admin/sleep would be a trivial remote DoS). It shares +the exact same ServiceManager, locks and depth tracking as the request path. + +The socket binds 0.0.0.0 *inside the container*; the host-side "localhost +only" restriction comes from compose publishing `127.0.0.1:8010:8010`. +""" + +from __future__ import annotations + +import logging +import time +from typing import Any + +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse + +from config import DEPTH_AWAKE +from services import ServiceManager + +log = logging.getLogger("vllm_router.admin") + + +def build_admin_app(manager: ServiceManager) -> FastAPI: + cfg = manager.cfg + app = FastAPI(title="vllm-router admin", version="1.0.0", + docs_url=None, redoc_url=None, openapi_url=None) + + def service_or_404(key: str) -> tuple[Any, JSONResponse | None]: + # Accepts the service key ("text") or the model name / alias + # ("OvisOCR2", "ocr", "qwen3.6-35b-a3b-fp8"), case-insensitively. + resolved = key.strip().lower() + if resolved not in manager.services: + resolved = cfg.index.get(key.strip().lower(), "") + svc = manager.services.get(resolved) + if svc is None: + body = {"error": {"message": f"Unknown service or model '{key}'.", + "type": "invalid_request_error", "code": "not_found", + "known": sorted(manager.services)}} + return None, JSONResponse(body, status_code=404) + return svc, None + + @app.get("/admin/status") + async def admin_status() -> JSONResponse: + status = await manager.status(live_probe=True) + return JSONResponse({ + "router": { + "uptime_s": status["uptime_s"], + "public_port": cfg.public_port, + "admin_port": cfg.admin_port, + "idle": { + "enabled": cfg.idle_enabled, + "sleep_min": cfg.idle_sleep_min, + "offload_min": cfg.idle_offload_min, + }, + }, + "services": status["services"], + "api": "http://127.0.0.1:%d/v1" % cfg.public_port, + }) + + @app.post("/admin/wake/{key}") + async def admin_wake(key: str) -> JSONResponse: + svc, err = service_or_404(key) + if err is not None: + return err + outcome = await manager.ensure_awake(svc.cfg.key) + svc.last_activity = time.monotonic() # idle clock restarts + body = { + "service": svc.cfg.key, + "model": svc.cfg.model, + "ok": outcome.ok, + "depth": outcome.depth, + "reason": outcome.reason, + "wake_in_progress": svc.wake_in_progress, + } + if outcome.latency_s is not None: + body["latency_s"] = round(outcome.latency_s, 2) + if not outcome.ok: + policy = cfg.policy(outcome.depth) + body["retry_after_s"] = policy.retry_after_s + body["estimated_wake_seconds"] = policy.est_wake_s + body["error"] = {"type": "model_waking", "code": "model_waking", + "sleep_depth": outcome.depth, + "estimated_wake_seconds": policy.est_wake_s} + if outcome.detail: + body["error"]["message"] = outcome.detail + log.info("admin_wake service=%s ok=false depth=%s reason=%s", + svc.cfg.key, outcome.depth, outcome.reason) + return JSONResponse(body, status_code=503, + headers={"Retry-After": str(policy.retry_after_s)}) + log.info("admin_wake service=%s ok=true depth=%s reason=%s", + svc.cfg.key, outcome.depth, outcome.reason) + return JSONResponse(body) + + @app.post("/admin/sleep/{key}") + async def admin_sleep(key: str, request: Request) -> JSONResponse: + svc, err = service_or_404(key) + if err is not None: + return err + level = 1 + if request.query_params.get("level"): + try: + level = int(request.query_params["level"]) + except ValueError: + return JSONResponse( + {"error": {"message": "level must be 1 or 2", + "type": "invalid_request_error"}}, status_code=400) + if level not in (1, 2): + return JSONResponse( + {"error": {"message": "level must be 1 or 2", + "type": "invalid_request_error"}}, status_code=400) + result = await manager.sleep_service(svc.cfg.key, level, reason="admin") + status_code = 200 if result.get("ok") else 409 + result["level_requested"] = level + result["model"] = svc.cfg.model + log.info("admin_sleep service=%s level=%d ok=%s reason=%s", + svc.cfg.key, level, result.get("ok"), result.get("reason")) + return JSONResponse(result, status_code=status_code) + + @app.get("/health") + @app.get("/admin/health", include_in_schema=False) + async def admin_health() -> JSONResponse: + return JSONResponse({"status": "ok", + "awake": [k for k, s in manager.services.items() + if s.depth == DEPTH_AWAKE], + "wake_recovery_pending": [k for k, s in manager.services.items() + if s.pending_reload]}) + + @app.exception_handler(Exception) + async def internal_error(_request: Request, exc: Exception) -> JSONResponse: + if manager.shutting_down: + return JSONResponse( + {"error": {"type": "router_shutting_down", "code": "router_shutting_down", + "message": "Router is shutting down; retry shortly"}}, + status_code=503, headers={"Retry-After": "5"}) + return JSONResponse( + {"error": {"type": "internal_error", "code": "internal_error", + "message": f"Unhandled router error: {type(exc).__name__}"}}, + status_code=500) + + return app diff --git a/router/app.py b/router/app.py new file mode 100644 index 0000000..23da301 --- /dev/null +++ b/router/app.py @@ -0,0 +1,187 @@ +"""vllm-router -- single-process front door for the vLLM serving stack. + +Plan v3.3 section 6. One process, ONE asyncio event loop, TWO listening +sockets: + + * public :8000 -> OpenAI API (/v1/*), /health, (/metrics) + * admin :8010 -> /admin/* (for `vllmctl`; published by compose as + `127.0.0.1:8010:8010`, so it is host-loopback only. Inside the + container it binds 0.0.0.0 because docker cannot publish to a + container-loopback bind -- see config.ROUTER_ADMIN_HOST.) + +Two uvicorn.Server instances are run as coroutines on the same loop because +uvicorn cannot bind two ports in one worker, and two *processes* would break +the in-process wake locks -- the admin listener must share the very same +ServiceManager, asyncio locks and depth tracking as the request path. + +What it owns: model-name -> service mapping, wake-on-request (single-flight, +depth-aware hold), streaming proxy, tiered idle sleep, dev-endpoint hiding +(404 for anything outside the allowlist, matched on the *normalized* path). + +Configuration (environment; authoritative list + defaults in config.py): + + registry VLLM_TEXT_MODEL=Qwen3.6-35B-A3B-FP8 VLLM_OCR_MODEL=OvisOCR2 + VLLM_EMBED_MODEL=Qwen3-Embedding-8B (+ *_SERVICE, *_ALIASES) + listeners ROUTER_PUBLIC_PORT=8000 ROUTER_ADMIN_PORT=8010 + ROUTER_ADMIN_HOST=127.0.0.1 + proxy ROUTER_CONNECT_TIMEOUT=5 ROUTER_READ_TIMEOUT=900 + ROUTER_POOL_SIZE=32 ROUTER_MAX_BODY_BYTES=134217728 + ROUTER_STATE_CACHE_TTL=2 (keep short: /health returns 200 on a + sleeping backend, so readiness is gated on /is_sleeping) + wake WAKE_HOLD_SLEEP_S=30 WAKE_HOLD_OFFLOAD_S=180 + WAKE_HOLD_RESTART_S=600 (text cold boot measured ~10 min), + WAKE_RETRY_AFTER_{SLEEP,OFFLOAD,RESTART}=10/60/600, + WAKE_ATTEMPTS=2, WAKE_HTTP_TIMEOUT=300 + idle IDLE_SLEEP_MIN=15 (level 1) IDLE_OFFLOAD_MIN=180 (level 2) + IDLE_POLL_SECONDS=30 IDLE_ENABLED=1 + +Run: python app.py (see Dockerfile; `pytest` from router/tests) +""" + +from __future__ import annotations + +import asyncio +import contextlib +import logging +import signal +import sys +from pathlib import Path + +# Sibling modules (config/services/routing/admin_api) -- makes the app work +# both as `python app.py` and when imported by the test suite. +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +import uvicorn # noqa: E402 + +from admin_api import build_admin_app # noqa: E402 +from config import Config, load_config # noqa: E402 +from routing import build_public_app # noqa: E402 +from services import ServiceManager # noqa: E402 + +log = logging.getLogger("vllm_router") + +ROUTES = ("GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS") + + +def setup_logging(level: str = "info") -> None: + logging.basicConfig( + level=getattr(logging, level.upper(), logging.INFO), + format="%(asctime)s %(levelname)-7s %(name)s %(message)s", + datefmt="%Y-%m-%dT%H:%M:%S", + ) + for noisy in ("httpx", "httpcore", "uvicorn.access"): + logging.getLogger(noisy).setLevel(logging.WARNING) + + +@contextlib.contextmanager +def _no_signal_capture() -> "contextlib.AbstractContextManager[None]": + """Replaces uvicorn.Server.capture_signals (see serve()).""" + yield + + +def build_apps(cfg: Config | None = None, transport=None): + """Build the public + admin apps around one shared ServiceManager. + + `transport` is an httpx transport override used by the test suite to stub + the vLLM backends without any network, GPU or docker. + """ + cfg = cfg or load_config() + manager = ServiceManager(cfg, transport=transport) + public_app = build_public_app(manager) + admin_app = build_admin_app(manager) + + @contextlib.asynccontextmanager + async def lifespan(_app): + manager.start() # idempotent: called by both listeners + yield + await manager.stop() # idempotent as well + + # Both apps share one manager; the lifespan is attached post-construction + # so there is exactly one definition of it. + public_app.router.lifespan_context = lifespan + admin_app.router.lifespan_context = lifespan + return public_app, admin_app, manager + + +def _make_server(app, host: str, port: int, cfg: Config) -> uvicorn.Server: + return uvicorn.Server(uvicorn.Config( + app, + host=host, + port=port, + log_level=cfg.log_level, + access_log=cfg.access_log, + # Long generations must not be cut off, but shutdown must terminate. + timeout_graceful_shutdown=10, + # The app does its own (structured) request logging. + timeout_keep_alive=75, + )) + + +async def serve(cfg: Config | None = None) -> None: + cfg = cfg or load_config() + setup_logging(cfg.log_level) + public_app, admin_app, manager = build_apps(cfg) + + servers = [ + _make_server(public_app, cfg.public_host, cfg.public_port, cfg), + _make_server(admin_app, cfg.admin_host, cfg.admin_port, cfg), + ] + + # uvicorn's signal capture is per-Server; with two servers the inner + # capture would restore its own handler and swallow SIGTERM for the other + # listener. Disable it and own the signals here (falling back gracefully + # on uvicorn versions without that method). + for server in servers: + if hasattr(server, "capture_signals"): + server.capture_signals = _no_signal_capture + + loop = asyncio.get_running_loop() + stopping = asyncio.Event() + + def _request_shutdown() -> None: + if stopping.is_set(): + return + stopping.set() + log.info("router_stop signal_received") + # Flag first, so requests dying in the shutdown window answer with a + # depth-aware 503 instead of a bare 500 (E2E case 12). + manager.begin_shutdown() + for server in servers: + server.should_exit = True + + for sig in (signal.SIGINT, signal.SIGTERM): + with contextlib.suppress(NotImplementedError, ValueError, RuntimeError): + loop.add_signal_handler(sig, _request_shutdown) + + async def _propagate_exit() -> None: + """Belt and braces: if only one server saw the shutdown, tell the other.""" + while not stopping.is_set(): + if any(s.should_exit for s in servers): + for s in servers: + s.should_exit = True + await asyncio.sleep(0.1) + + watchdog = asyncio.create_task(_propagate_exit(), name="exit-watchdog") + log.info("router_start public=%s:%d admin=%s:%d services=%s", + cfg.public_host, cfg.public_port, cfg.admin_host, cfg.admin_port, + ",".join(f"{s.key}={s.service}/{s.model}" for s in cfg.services.values())) + try: + await asyncio.gather(*(server.serve() for server in servers)) + finally: + watchdog.cancel() + with contextlib.suppress(asyncio.CancelledError): + await watchdog + await manager.stop() + log.info("router_stopped") + + +def main() -> int: + try: + asyncio.run(serve()) + except KeyboardInterrupt: # pragma: no cover - interactive runs + return 0 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/router/config.py b/router/config.py new file mode 100644 index 0000000..9bef658 --- /dev/null +++ b/router/config.py @@ -0,0 +1,293 @@ +"""Router configuration: model registry + tunables. + +Everything is environment-driven with the defaults from plan v3.3 (see +`.claude/memory/router-front-door-plan.md` sections 4/6). The `Config` +object is a plain mutable dataclass so tests (and calibration) can override +a single knob without re-reading the environment. + +Environment variables (all optional): + + Registry + VLLM_TEXT_SERVICE default vllm-text docker service name (== DNS name) + VLLM_TEXT_MODEL default Qwen3.6-35B-A3B-FP8 + VLLM_TEXT_ALIASES comma list, default "qwen,qwen3,qwen3.6,qwen3.6-35b, + qwen3.6-35b-a3b,text,chat,default" + VLLM_OCR_SERVICE default vllm-ocr + VLLM_OCR_MODEL default OvisOCR2 + VLLM_OCR_ALIASES default "ocr,ovis,ovisocr,ovis-ocr" + VLLM_EMBED_SERVICE default vllm-embed + VLLM_EMBED_MODEL default Qwen3-Embedding-8B + VLLM_EMBED_ALIASES default "embed,embedding,embeddings,qwen3-embedding, + qwen3-embedding-8b" + VLLM_BACKEND_PORT default 8000 (container port each vllm service listens on) + + Listeners + ROUTER_PUBLIC_HOST default 0.0.0.0 + ROUTER_PUBLIC_PORT default 8000 (the only public ingress) + ROUTER_ADMIN_HOST default 0.0.0.0 (inside the container; the host-side + 127.0.0.1 restriction comes from compose's + "127.0.0.1:8010:8010" port mapping -- docker cannot + publish to a container-loopback bind. Set 127.0.0.1 + when running the router on bare metal.) + ROUTER_ADMIN_PORT default 8010 (vllmctl / debugging) + + Proxy (plan 6.2 "Proxy timeouts") + ROUTER_CONNECT_TIMEOUT default 5 seconds + ROUTER_READ_TIMEOUT default 900 seconds (capped, never disabled) + ROUTER_POOL_SIZE default 32 (>= 16 concurrent requests) + ROUTER_PROBE_TIMEOUT default 3 seconds for /is_sleeping + /health + ROUTER_MAX_BODY_BYTES default 134217728 (128 MiB request body cap) + + Wake-state cache + ROUTER_STATE_CACHE_TTL default 2.0 seconds (must stay short: /health + returns 200 on a sleeping backend and a proxied + request to one HANGS, so a stale "awake" must not + bypass the is_sleeping gate for long) + + Depth-aware hold/503 policy (plan 6.2.1; restart deadline from calibration) + WAKE_HOLD_SLEEP_S default 30 hold deadline, wake from level 1 + (measured wake 2.5-3.8s) + WAKE_HOLD_OFFLOAD_S default 180 hold deadline, wake from level 2 + WAKE_HOLD_RESTART_S default 600 hold deadline, container restarting + (measured text cold boot up to + ~10 min from NFS) + WAKE_RETRY_AFTER_SLEEP default 10 Retry-After header value + WAKE_RETRY_AFTER_OFFLOAD default 60 + WAKE_RETRY_AFTER_RESTART default 600 + WAKE_EST_SLEEP_S default 6 body estimated_wake_seconds + WAKE_EST_OFFLOAD_S default 60 + WAKE_EST_RESTART_S default 600 + WAKE_ATTEMPTS default 2 (sequence retried once, plan 6.2.1) + WAKE_HTTP_TIMEOUT default 300 per-call timeout inside the sequence + WAKE_HEALTH_POLL_S default 2.0 /health poll interval while waking + + Tiered idle (plan 6.3) + IDLE_SLEEP_MIN default 15 minutes -> POST /sleep?level=1 + IDLE_OFFLOAD_MIN default 180 minutes -> POST /sleep?level=2 + IDLE_POLL_SECONDS default 30 idle-manager scan interval + IDLE_ENABLED default 1 (0 disables the background tiering) + + Misc + ROUTER_LOG_LEVEL default info + ROUTER_METRICS default 0 (1 = also expose /metrics publicly) + ROUTER_ACCESS_LOG default 0 (uvicorn access log; off by default, + the app logs one structured line per request) + ROUTER_STATE_FILE default /state/router-state.json -- persisted wake + intent + depth, so a restarted router can tell + "awake and ready" from "awake because the previous + router died between wake_up and reload_weights". + Compose mounts the named volume `router-state` + at /state (seeded with image ownership so the + non-root user can write it). + Set empty to disable persistence. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field, replace + + +def _str(name: str, default: str) -> str: + v = os.environ.get(name) + return default if v is None or v == "" else v + + +def _int(name: str, default: int) -> int: + try: + return int(_str(name, str(default))) + except ValueError: + return default + + +def _float(name: str, default: float) -> float: + try: + return float(_str(name, str(default))) + except ValueError: + return default + + +def _bool(name: str, default: bool) -> bool: + return _str(name, "1" if default else "0").strip().lower() in ("1", "true", "yes", "on") + + +def _csv(name: str, default: str) -> tuple[str, ...]: + return tuple(p.strip() for p in _str(name, default).split(",") if p.strip()) + + +@dataclass(frozen=True) +class ServiceConfig: + """One registry entry: key -> (docker service, served model name).""" + + key: str + service: str + model: str + aliases: tuple[str, ...] = () + + @property + def base_url(self) -> str: + return f"http://{self.service}:{BACKEND_PORT}" + + +# Read once at import; the base_url property depends on it. +BACKEND_PORT = _int("VLLM_BACKEND_PORT", 8000) + + +@dataclass +class DepthPolicy: + """Hold/retry parameters for one sleep depth (plan 6.2.1 table).""" + + depth: str + hold_s: float + retry_after_s: int + est_wake_s: int + phrase: str # "... is waking from ; retry shortly" + + def error_body(self, model: str) -> dict: + return { + "error": { + "type": "model_waking", + "code": "model_waking", + "message": f"Model '{model}' is waking from {self.phrase}; retry shortly", + "sleep_depth": self.depth, + "estimated_wake_seconds": self.est_wake_s, + } + } + + +@dataclass +class Config: + services: dict[str, ServiceConfig] + # name (lower-cased) -> key, for model resolution + index: dict[str, str] = field(default_factory=dict) + + public_host: str = "0.0.0.0" + public_port: int = 8000 + # NOTE: 0.0.0.0 is required *inside the container* -- docker port + # publishing cannot reach a container-loopback bind. The "localhost only" + # guarantee of plan section 4 comes from compose mapping the port as + # `127.0.0.1:8010:8010` (host side). On bare metal set + # ROUTER_ADMIN_HOST=127.0.0.1. + admin_host: str = "0.0.0.0" + admin_port: int = 8010 + + connect_timeout: float = 5.0 + read_timeout: float = 900.0 + pool_size: int = 32 + probe_timeout: float = 3.0 + max_body_bytes: int = 128 * 1024 * 1024 + + state_cache_ttl: float = 2.0 + + hold_sleep_s: float = 30.0 + hold_offload_s: float = 180.0 + hold_restart_s: float = 600.0 + retry_after_sleep: int = 10 + retry_after_offload: int = 60 + retry_after_restart: int = 600 + est_sleep_s: int = 6 + est_offload_s: int = 60 + est_restart_s: int = 600 + wake_attempts: int = 2 + wake_http_timeout: float = 300.0 + wake_health_poll_s: float = 2.0 + + idle_sleep_min: float = 15.0 + idle_offload_min: float = 180.0 + idle_poll_s: float = 30.0 + idle_enabled: bool = True + + metrics_enabled: bool = False + log_level: str = "info" + access_log: bool = False + # Wake-intent state file (bind-mounted so it survives a router restart). + state_file: str | None = "/state/router-state.json" + + def policy(self, depth: str) -> DepthPolicy: + if depth == DEPTH_SLEEPING: + return DepthPolicy(depth, self.hold_sleep_s, self.retry_after_sleep, + self.est_sleep_s, "sleep") + if depth == DEPTH_RESTARTING: + return DepthPolicy(depth, self.hold_restart_s, self.retry_after_restart, + self.est_restart_s, "a container restart") + # DEPTH_OFFLOADED and anything unknown: be conservative (plan 6.2.1). + return DepthPolicy(DEPTH_OFFLOADED, self.hold_offload_s, self.retry_after_offload, + self.est_offload_s, "offload") + + +# Sleep depths tracked by the router. +DEPTH_AWAKE = "awake" +DEPTH_SLEEPING = "sleeping" # /sleep?level=1 - weights moved to host RAM +DEPTH_OFFLOADED = "offloaded" # /sleep?level=2 - RAM freed, weights on NFS +DEPTH_RESTARTING = "restarting" # container down / cold-starting + + +def load_config() -> Config: + services: dict[str, ServiceConfig] = { + "text": ServiceConfig( + key="text", + service=_str("VLLM_TEXT_SERVICE", "vllm-text"), + model=_str("VLLM_TEXT_MODEL", "Qwen3.6-35B-A3B-FP8"), + aliases=_csv("VLLM_TEXT_ALIASES", + "qwen,qwen3,qwen3.6,qwen3.6-35b,qwen3.6-35b-a3b,text,chat,default"), + ), + "ocr": ServiceConfig( + key="ocr", + service=_str("VLLM_OCR_SERVICE", "vllm-ocr"), + model=_str("VLLM_OCR_MODEL", "OvisOCR2"), + aliases=_csv("VLLM_OCR_ALIASES", "ocr,ovis,ovisocr,ovis-ocr"), + ), + "embed": ServiceConfig( + key="embed", + service=_str("VLLM_EMBED_SERVICE", "vllm-embed"), + model=_str("VLLM_EMBED_MODEL", "Qwen3-Embedding-8B"), + aliases=_csv("VLLM_EMBED_ALIASES", + "embed,embedding,embeddings,qwen3-embedding,qwen3-embedding-8b"), + ), + } + + index: dict[str, str] = {} + for cfg in services.values(): + names = {cfg.model.lower(), cfg.key.lower(), *(a.lower() for a in cfg.aliases)} + for name in names: + index[name] = cfg.key + + return Config( + services=services, + index=index, + public_host=_str("ROUTER_PUBLIC_HOST", "0.0.0.0"), + public_port=_int("ROUTER_PUBLIC_PORT", 8000), + admin_host=_str("ROUTER_ADMIN_HOST", "0.0.0.0"), + admin_port=_int("ROUTER_ADMIN_PORT", 8010), + connect_timeout=_float("ROUTER_CONNECT_TIMEOUT", 5.0), + read_timeout=_float("ROUTER_READ_TIMEOUT", 900.0), + pool_size=_int("ROUTER_POOL_SIZE", 32), + probe_timeout=_float("ROUTER_PROBE_TIMEOUT", 3.0), + max_body_bytes=_int("ROUTER_MAX_BODY_BYTES", 128 * 1024 * 1024), + state_cache_ttl=_float("ROUTER_STATE_CACHE_TTL", 2.0), + hold_sleep_s=_float("WAKE_HOLD_SLEEP_S", 30.0), + hold_offload_s=_float("WAKE_HOLD_OFFLOAD_S", 180.0), + hold_restart_s=_float("WAKE_HOLD_RESTART_S", 600.0), + retry_after_sleep=_int("WAKE_RETRY_AFTER_SLEEP", 10), + retry_after_offload=_int("WAKE_RETRY_AFTER_OFFLOAD", 60), + retry_after_restart=_int("WAKE_RETRY_AFTER_RESTART", 600), + est_sleep_s=_int("WAKE_EST_SLEEP_S", 6), + est_offload_s=_int("WAKE_EST_OFFLOAD_S", 60), + est_restart_s=_int("WAKE_EST_RESTART_S", 600), + wake_attempts=max(1, _int("WAKE_ATTEMPTS", 2)), + wake_http_timeout=_float("WAKE_HTTP_TIMEOUT", 300.0), + wake_health_poll_s=_float("WAKE_HEALTH_POLL_S", 2.0), + idle_sleep_min=_float("IDLE_SLEEP_MIN", 15.0), + idle_offload_min=_float("IDLE_OFFLOAD_MIN", 180.0), + idle_poll_s=_float("IDLE_POLL_SECONDS", 30.0), + idle_enabled=_bool("IDLE_ENABLED", True), + metrics_enabled=_bool("ROUTER_METRICS", False), + log_level=_str("ROUTER_LOG_LEVEL", "info").lower(), + access_log=_bool("ROUTER_ACCESS_LOG", False), + state_file=_str("ROUTER_STATE_FILE", "/state/router-state.json") or None, + ) + + +def clone(cfg: Config) -> Config: + """Copy for tests.""" + return replace(cfg, services=dict(cfg.services), index=dict(cfg.index)) diff --git a/router/pytest.ini b/router/pytest.ini new file mode 100644 index 0000000..59a0ff6 --- /dev/null +++ b/router/pytest.ini @@ -0,0 +1,5 @@ +[pytest] +testpaths = tests +asyncio_mode = auto +filterwarnings = + error::DeprecationWarning:tests.* diff --git a/router/requirements.txt b/router/requirements.txt new file mode 100644 index 0000000..d902a79 --- /dev/null +++ b/router/requirements.txt @@ -0,0 +1,10 @@ +# Pinned (plan section 6). Direct dependencies only; the transitive set +# (starlette, pydantic, anyio, ...) is resolved by pip at build time and +# recorded by the image build. +# +# fastapi == 0.141.1 (current stable, verified with starlette 1.6) +# uvicorn == 0.52.3 (two Server instances on one asyncio loop) +# httpx == 0.28.1 (AsyncClient streaming proxy) +fastapi==0.141.1 +uvicorn==0.52.3 +httpx==0.28.1 diff --git a/router/routing.py b/router/routing.py new file mode 100644 index 0000000..c76d6d8 --- /dev/null +++ b/router/routing.py @@ -0,0 +1,536 @@ +"""Public listener: path allowlist, model resolution, streaming proxy. + +Security shape (plan section 6): + * the *normalized* path decides what is served, so `/v1/../sleep` and + `/v1%2f..%2fsleep` both collapse to `/sleep` and get a 404; + * `/admin/*` and every vLLM dev endpoint are simply not in the allowlist; + * the admin listener lives on another socket (see admin_api.py). +""" + +from __future__ import annotations + +import json +import logging +import posixpath +import re +import time +import urllib.parse +from typing import Any + +import httpx +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse, Response, StreamingResponse +from starlette.background import BackgroundTask +from starlette.datastructures import MutableHeaders +from starlette.middleware.body_limit import RequestBodyLimitMiddleware + +from config import ( + DEPTH_AWAKE, + DEPTH_OFFLOADED, + DEPTH_RESTARTING, + DEPTH_SLEEPING, + Config, +) +from services import ServiceManager, ServiceState + +log = logging.getLogger("vllm_router.http") + +# Endpoints whose bodies legitimately omit `model` (plan 6.2 / 6.6). +DEFAULT_TEXT_PATHS = frozenset({"/v1/chat/completions", "/v1/completions"}) +# A multipart request on the chat path is the image-bearing (OCR) path. +MULTIPART_DEFAULT_PATH = "/v1/chat/completions" + +# Headers that must never be forwarded in either direction. +HOP_BY_HOP = frozenset({ + "connection", "keep-alive", "proxy-authenticate", "proxy-authorization", + "te", "trailers", "transfer-encoding", "upgrade", +}) +REQUEST_ONLY_STRIP = HOP_BY_HOP | {"host", "content-length", "expect"} +RESPONSE_STRIP = HOP_BY_HOP | {"content-length"} + +_BOUNDARY_RE = re.compile(r'boundary="?([^";,]+)"?', re.IGNORECASE) +_NAME_RE = re.compile(r'name="((?:[^"\\]|\\.)*)"', re.IGNORECASE) +# A text field we are willing to read fully (and re-send) while sniffing. +_MAX_MODEL_FIELD_BYTES = 4096 + + +# -------------------------------------------------------------------------- +# errors (OpenAI-shaped so existing clients can parse them) +# -------------------------------------------------------------------------- +def _error(status: int, message: str, *, err_type: str = "invalid_request_error", + code: str | None = None, param: str | None = None, + extra: dict[str, Any] | None = None, headers: dict[str, str] | None = None + ) -> JSONResponse: + err: dict[str, Any] = {"message": message, "type": err_type} + if param is not None: + err["param"] = param + if code is not None: + err["code"] = code + if extra: + err.update(extra) + return JSONResponse({"error": err}, status_code=status, headers=headers) + + +def not_found_response(message: str = "Not found") -> JSONResponse: + """Uniform 404 for everything outside the allowlist (paths *and* models).""" + return _error(404, message, code="not_found") + + +def shutting_down_response(detail: str = "router is shutting down") -> JSONResponse: + """503 for in-flight work that dies because the router is stopping.""" + return JSONResponse( + {"error": { + "type": "router_shutting_down", + "code": "router_shutting_down", + "message": f"Request aborted: {detail}; retry shortly", + "retry_after_seconds": 5, + }}, + status_code=503, + headers={"Retry-After": "5", "connection": "close"}, + ) + + +_SHUTDOWN_BODY = json.dumps({ + "error": { + "type": "router_shutting_down", + "code": "router_shutting_down", + "message": "Request aborted: the router is shutting down; retry shortly", + "retry_after_seconds": 5, + } +}).encode() + + +class ShutdownGuardMiddleware: + """Turns shutdown-induced cancellation of an in-flight request into a + depth-aware 503 instead of uvicorn's bare 21-byte 500 (E2E case 12). + + uvicorn cancels pending ASGI tasks once `timeout_graceful_shutdown` + expires; that CancelledError is a BaseException, so Starlette's + ServerErrorMiddleware (and every FastAPI exception_handler) ignores it and + uvicorn answers with a plain-text 500. This is the only layer that can + intercept it. A response whose body has already started streaming cannot + be un-sent: those connections are simply closed (client sees a truncated + stream), which is the best available behaviour. + """ + + def __init__(self, app, manager: ServiceManager) -> None: + self.app = app + self.manager = manager + + async def __call__(self, scope, receive, send) -> None: + if scope["type"] != "http": + await self.app(scope, receive, send) + return + started = False + + async def send_wrapper(message) -> None: + nonlocal started + if message["type"] == "http.response.start": + started = True + await send(message) + + try: + await self.app(scope, receive, send_wrapper) + except BaseException as exc: # noqa: BLE001 - CancelledError is one + if started or not self.manager.shutting_down: + raise + log.info("abort_shutdown method=%s path=%s error=%s", + scope.get("method"), scope.get("path"), type(exc).__name__) + try: + await send({"type": "http.response.start", "status": 503, "headers": [ + (b"content-type", b"application/json"), + (b"content-length", str(len(_SHUTDOWN_BODY)).encode("latin-1")), + (b"retry-after", b"5"), + (b"connection", b"close"), + ]}) + await send({"type": "http.response.body", "body": _SHUTDOWN_BODY}) + except Exception: # pragma: no cover - client may be gone already + pass + + +def model_not_found_response(model: str) -> JSONResponse: + return _error( + 404, + f"The model `{model}` does not exist or is not served by this endpoint.", + param="model", + code="model_not_found", + ) + + +def waking_response(cfg: Config, svc: ServiceState, depth: str, + detail: str = "") -> JSONResponse: + """Depth-aware 503 (plan 6.2.1).""" + policy = cfg.policy(depth) + body = policy.error_body(svc.cfg.model) + if detail: + body["error"]["message"] += f" ({detail})" + return JSONResponse( + body, + status_code=503, + headers={ + "Retry-After": str(policy.retry_after_s), + "x-vllm-router": "model-waking", + }, + ) + + +# -------------------------------------------------------------------------- +# path handling +# -------------------------------------------------------------------------- +def normalize_raw_path(raw: bytes | str) -> str: + """Decode %xx and collapse `.`/`..` segments BEFORE any allowlist check.""" + if isinstance(raw, (bytes, bytearray)): + raw = bytes(raw).decode("latin-1") + path = urllib.parse.unquote(raw) + path = posixpath.normpath(path) + if not path.startswith("/"): + path = "/" + path + return path + + +def classify_public_path(norm: str, cfg: Config) -> str | None: + """Returns the handler kind, or None for 404.""" + if norm == "/health": + return "health" + if norm == "/metrics" and cfg.metrics_enabled: + return "metrics" + if norm == "/v1/models": + return "models" + if norm.startswith("/v1/") and ".." not in norm.split("/") and "\\" not in norm: + return "api" + return None + + +# -------------------------------------------------------------------------- +# model resolution (plan 6.2 step 1) +# -------------------------------------------------------------------------- +def extract_multipart_model(body: bytes, content_type: str) -> str | None: + """Pull a small `model` text field out of a multipart body. + + Deliberately hand-rolled: Starlette's form parser would consume (and + buffer/spool) every file part. We walk the boundary delimiters, look at + part headers only, and never slice a file part's payload. + """ + match = _BOUNDARY_RE.search(content_type) + if match is None: + return None + delimiter = b"--" + match.group(1).encode("latin-1") + pos = 0 + while True: + start = body.find(delimiter, pos) + if start < 0: + return None + line_end = body.find(b"\r\n", start + len(delimiter)) + if line_end < 0: + return None + next_delim = body.find(delimiter, line_end) + if next_delim < 0: + return None + headers = body[line_end + 2:next_delim] + head, sep, _ = headers.partition(b"\r\n\r\n") + if not sep: + pos = next_delim + continue + disposition = b"" + for line in head.split(b"\r\n"): + if line.lower().startswith(b"content-disposition:"): + disposition = line + break + if b"filename=" in disposition.lower(): + pos = next_delim # file part: leave it untouched + continue + name_match = _NAME_RE.search(disposition.decode("latin-1", "replace")) + if name_match is None or name_match.group(1).lower() != "model": + pos = next_delim + continue + value = headers[len(head) + 4:].rstrip(b"\r\n") + if len(value) > _MAX_MODEL_FIELD_BYTES: + return None + return value.decode("utf-8", "replace").strip() or None + + +def resolve_model_name(cfg: Config, name: str) -> str | None: + """Registry lookup (case-insensitive over model names, aliases and keys).""" + return cfg.index.get(name.strip().lower()) + + +async def resolve_target(manager: ServiceManager, request: Request, norm_path: str + ) -> tuple[ServiceState | None, Response | None]: + """Decide which service a request belongs to. Returns (service, error); + exactly one is None. Never wakes anything when the model is unknown.""" + cfg = manager.cfg + + # 1. /v1/embeddings always goes to the embed service. + if norm_path == "/v1/embeddings": + return manager.services["embed"], None + + # 2. Model named in the path (GET /v1/models/{id}). + if norm_path.startswith("/v1/models/") and request.method in ("GET", "HEAD"): + key = resolve_model_name(cfg, urllib.parse.unquote(norm_path[len("/v1/models/"):])) + if key is None: + return None, model_not_found_response(norm_path[len("/v1/models/"):]) + return manager.services[key], None + + content_type = request.headers.get("content-type", "") + is_json = content_type.split(";")[0].strip().lower() in ("application/json",) or \ + content_type.strip().lower().endswith("+json") + is_multipart = content_type.split(";")[0].strip().lower().startswith("multipart/") + + if is_json: + raw = await request.body() + try: + payload = json.loads(raw or b"{}") + except (ValueError, UnicodeDecodeError): + return None, _error(400, "Request body is not valid JSON.", code="invalid_json") + if not isinstance(payload, dict): + return None, _error(400, "Request body must be a JSON object.", code="invalid_json") + model = payload.get("model") + if isinstance(model, str) and model.strip(): + key = resolve_model_name(cfg, model) + if key is None: + return None, model_not_found_response(model) + return manager.services[key], None + if norm_path in DEFAULT_TEXT_PATHS: + return manager.services["text"], None + return None, _error( + 400, + "Missing required parameter: 'model'.", + param="model", + code="missing_model", + ) + + if is_multipart: + model = extract_multipart_model(await request.body(), content_type) + if model is not None: + key = resolve_model_name(cfg, model) + if key is None: + return None, model_not_found_response(model) + return manager.services[key], None + if norm_path == MULTIPART_DEFAULT_PATH: + return manager.services["ocr"], None # image-bearing path + if norm_path in DEFAULT_TEXT_PATHS: + return manager.services["text"], None + return None, _error( + 400, + "Missing required parameter: 'model'.", + param="model", + code="missing_model", + ) + + # Any other body shape (empty, text/plain, ...): only the two chat paths + # have a defensible default. + if norm_path in DEFAULT_TEXT_PATHS: + return manager.services["text"], None + return None, _error( + 400, + "Missing required parameter: 'model'.", + param="model", + code="missing_model", + ) + + +# -------------------------------------------------------------------------- +# proxy +# -------------------------------------------------------------------------- +def _forwarded_request_headers(request: Request) -> list[tuple[str, str]]: + return [ + (k, v) for k, v in request.headers.items() + if k.lower() not in REQUEST_ONLY_STRIP + ] + + +async def _open_upstream(manager: ServiceManager, request: Request, svc: ServiceState, + norm_path: str) -> httpx.Response: + query = request.scope.get("query_string", b"").decode("latin-1") + url = svc.cfg.base_url + norm_path + (f"?{query}" if query else "") + upstream_req = manager.proxy_client.build_request( + request.method, + url, + headers=_forwarded_request_headers(request), + content=await request.body(), + ) + # Stream: the response body is handed to the client chunk by chunk and is + # never buffered here (SSE safe, 900s read timeout upstream). + return await manager.proxy_client.send(upstream_req, stream=True) + + +async def _proxy_response(manager: ServiceManager, request: Request, svc: ServiceState, + norm_path: str) -> Response: + upstream = await _open_upstream(manager, request, svc, norm_path) + headers = MutableHeaders() + for key, value in upstream.headers.items(): + if key.lower() not in RESPONSE_STRIP: + headers.append(key, value) + + async def finish() -> None: + # Runs when the last chunk has been sent (or the client hung up) -- + # this is what keeps the idle manager from sleeping mid-stream. + try: + await upstream.aclose() + finally: + manager.end_request(svc) + + return StreamingResponse( + upstream.aiter_raw(), + status_code=upstream.status_code, + headers=headers, + background=BackgroundTask(finish), + ) + + +async def forward(manager: ServiceManager, request: Request, svc: ServiceState, + norm_path: str) -> Response: + """Ensure the target is awake, then proxy. The request stays counted as + active for the whole life of the response body.""" + cfg = manager.cfg + manager.begin_request(svc) + handed_off = False + try: + for attempt in (1, 2): + outcome = await manager.ensure_awake(svc.cfg.key) + if not outcome.ok: + log.info("reject service=%s status=503 depth=%s reason=%s", + svc.cfg.key, outcome.depth, outcome.reason) + detail = "" if outcome.reason == "hold_timeout" else outcome.detail + return waking_response(cfg, svc, outcome.depth, detail) + try: + response = await _proxy_response(manager, request, svc, norm_path) + handed_off = True + return response + except httpx.TransportError as exc: + # Backend went away between the health check and the proxy call. + svc.reachable = False + svc.depth = None + svc.state_checked_at = float("-inf") + svc.last_error = f"proxy: {type(exc).__name__}" + log.warning("proxy_transport_error service=%s attempt=%d error=%s", + svc.cfg.key, attempt, type(exc).__name__) + if attempt == 2: + return waking_response(cfg, svc, DEPTH_RESTARTING, + "backend connection failed") + finally: + if not handed_off: + manager.end_request(svc) + return waking_response(cfg, svc, DEPTH_RESTARTING, "unreachable") # pragma: no cover + + +# -------------------------------------------------------------------------- +# public endpoints +# -------------------------------------------------------------------------- +def _service_summary(manager: ServiceManager, live: bool) -> dict[str, Any]: + out: dict[str, Any] = {} + for svc in manager.services.values(): + out[svc.cfg.key] = { + "service": svc.cfg.service, + "model": svc.cfg.model, + "reachable": svc.reachable, + "sleeping": None if svc.depth is None else svc.depth != DEPTH_AWAKE, + "depth": svc.depth, + "wake_in_progress": svc.wake_in_progress, + "active_requests": svc.active_requests, + } + return out + + +def build_public_app(manager: ServiceManager) -> FastAPI: + cfg = manager.cfg + app = FastAPI( + title="vllm-router", + version="1.0.0", + docs_url=None, + redoc_url=None, + openapi_url=None, + ) + # FastAPI does not forward **extra to Starlette, so the request-body cap + # (ROUTER_MAX_BODY_BYTES) is wired as middleware explicitly. + app.add_middleware(RequestBodyLimitMiddleware, max_body_size=cfg.max_body_bytes) + # Added last => outermost. Must sit outside everything because the + # exception it converts is a BaseException that Starlette's + # ServerErrorMiddleware deliberately does not catch. + app.add_middleware(ShutdownGuardMiddleware, manager=manager) + + @app.exception_handler(Exception) + async def internal_error(_request: Request, exc: Exception) -> JSONResponse: + """Anything unhandled becomes a parseable JSON error, never a bare + 21-byte "Internal Server Error". During shutdown that is a 503 with + Retry-After (E2E case 12).""" + if manager.shutting_down: + log.info("abort_shutdown method=%s path=%s error=%s", + _request.method, _request.url.path, type(exc).__name__) + return shutting_down_response(f"{type(exc).__name__} during shutdown") + log.exception("internal_error error=%s", type(exc).__name__) + return _error(500, f"Unhandled router error: {type(exc).__name__}", + err_type="internal_error", code="internal_error", + headers={"Retry-After": "1"}) + + @app.get("/health") + async def health() -> JSONResponse: + """Router liveness + cached per-service summary. Always 200: a sleeping + backend is normal, not an outage (docker healthchecks must not flap).""" + return JSONResponse({ + "status": "ok", + "router": { + "uptime_s": round(time.monotonic() - manager.started_at, 1), + "public_port": cfg.public_port, + "admin_port": cfg.admin_port, + "services": len(manager.services), + }, + "services": _service_summary(manager, live=False), + }) + + @app.get("/v1/models") + async def models() -> JSONResponse: + return JSONResponse({ + "object": "list", + "data": [ + {"id": svc.cfg.model, "object": "model", "created": 0, + "owned_by": "vllm-router"} + for svc in manager.services.values() + ], + }) + + # NOTE: /metrics is deliberately NOT registered as an explicit route -- the + # catch-all below re-checks it against the allowlist so it 404s whenever + # ROUTER_METRICS is off. + async def metrics() -> Response: + lines = [ + "# HELP vllm_router_service_depth 0=awake 1=sleeping 2=offloaded 3=restarting", + "# TYPE vllm_router_service_depth gauge", + ] + rank = {DEPTH_AWAKE: 0, DEPTH_SLEEPING: 1, DEPTH_OFFLOADED: 2, DEPTH_RESTARTING: 3} + for svc in manager.services.values(): + lines.append(f'vllm_router_service_depth{{service="{svc.cfg.key}"}} ' + f'{rank.get(svc.depth, 2)}') + lines.append(f'vllm_router_active_requests{{service="{svc.cfg.key}"}} ' + f'{svc.active_requests}') + return Response("\n".join(lines) + "\n", media_type="text/plain; version=0.0.4") + + async def gate(request: Request) -> Response: + raw = request.scope.get("raw_path") or request.url.path.encode() + norm = normalize_raw_path(raw) + kind = classify_public_path(norm, cfg) + if kind is None: + log.info("reject_path method=%s path=%s", request.method, norm) + return not_found_response() + if kind == "health": + return await health() + if kind == "models": + return await models() + if kind == "metrics": + return await metrics() + + svc, error = await resolve_target(manager, request, norm) + if error is not None: + log.info("reject_model method=%s path=%s status=%s", + request.method, norm, error.status_code) + return error + return await forward(manager, request, svc, norm) + + # Registered last: explicit routes above win, everything else lands here + # and is re-checked against the allowlist on the *normalized* path. + app.router.add_route( + "/{path:path}", + gate, + methods=["GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"], + ) + return app diff --git a/router/services.py b/router/services.py new file mode 100644 index 0000000..21396ae --- /dev/null +++ b/router/services.py @@ -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, + } diff --git a/router/tests/conftest.py b/router/tests/conftest.py new file mode 100644 index 0000000..e471a27 --- /dev/null +++ b/router/tests/conftest.py @@ -0,0 +1,263 @@ +"""Shared fixtures + a fake vLLM backend transport (no GPU / docker / network). + +The router's outbound httpx clients are built with an injectable transport +(`ServiceManager(cfg, transport=...)`), so a single fake serves all three +backends and records every call the router makes. +""" + +from __future__ import annotations + +import asyncio +import json +import sys +import urllib.parse +from types import SimpleNamespace +from pathlib import Path + +import httpx +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from app import build_apps # noqa: E402 +from config import clone, load_config # noqa: E402 +from services import ServiceManager # noqa: E402 + + +class FakeVLLM(httpx.AsyncBaseTransport): + """Stands in for vllm-text / vllm-ocr / vllm-embed. + + Endpoints: /health, /is_sleeping, /wake_up, /collective_rpc, + /reset_prefix_cache, /sleep?level=, plus /v1/* which echoes what it saw. + """ + + def __init__(self) -> None: + self.calls: list[str] = [] + self.services: dict[str, dict] = { + "vllm-text": self._svc(), + "vllm-ocr": self._svc(), + "vllm-embed": self._svc(), + } + self.stream_chunks: list[bytes] = [] + self.stream_delay: float = 0.0 # per-chunk delay for slow/SSE streams + + @staticmethod + def _svc() -> dict: + return { + "reachable": True, + "healthy": True, + "sleeping": False, + "wake_delay": 0.0, + "wake_fails": 0, # number of /wake_up calls to fail (500) + "sleep_fails": 0, # number of /sleep calls to fail (500) + "wake_up_calls": 0, + "reload_calls": 0, + "reset_calls": 0, + "sleep_calls": 0, + "api_calls": 0, + "api_delay": 0.0, + } + + # -- helpers ----------------------------------------------------------- + @staticmethod + def _response(status: int = 200, *, payload: bytes = b"", content_type: str | None = None): + """Build a *streamable* response. + + `httpx.Response(200, json=...)` would mark the body as already + consumed, which the router's streaming proxy would rightly reject. + """ + async def gen(): + if payload: + yield payload + + headers = {"content-type": content_type} if content_type else {} + return httpx.Response(status, content=gen(), headers=headers) + + def count(self, needle: str) -> int: + return sum(1 for call in self.calls if needle in call) + + def last(self, needle: str) -> int: + """Index of the last matching recorded call (-1 if none).""" + for i in range(len(self.calls) - 1, -1, -1): + if needle in self.calls[i]: + return i + return -1 + + def set_sleeping(self, key: str, sleeping: bool, level: int = 0) -> None: + svc = self.services[f"vllm-{key}"] + svc["sleeping"] = sleeping + svc["level"] = level + + # -- transport --------------------------------------------------------- + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + host = request.url.host or "" + svc = self.services.get(host) + if svc is None or not svc["reachable"]: + raise httpx.ConnectError(f"connection refused: {host}", request=request) + + path = request.url.path + params = request.url.params + # record method + full target (path + query) so tests can assert on + # things like "level=1" + target = path + (f"?{params}" if str(params) else "") + self.calls.append(f"{request.method} {host}{target}") + + if path == "/is_sleeping": + # Optional hook: report `true` for the next N polls even after a + # wake-up, to model a backend whose wake is still in flight. + if svc.get("hold_sleeping_polls", 0) > 0: + svc["hold_sleeping_polls"] -= 1 + return self._response(200, payload=b'{"is_sleeping": true}', + content_type="application/json") + return self._response(200, payload=json.dumps( + {"is_sleeping": bool(svc["sleeping"])}).encode(), + content_type="application/json") + + if path == "/health": + # NOTE: like the real vLLM, /health answers 200 even while the + # model is asleep -- readiness must be gated on /is_sleeping. + if not svc["healthy"]: + raise httpx.ConnectError(f"unhealthy: {host}", request=request) + return self._response(200) + + if path == "/wake_up": + svc["wake_up_calls"] += 1 + if svc["wake_fails"] > 0: + svc["wake_fails"] -= 1 + return self._response(500, payload=b"wake failed") + if svc.get("wake_flip_midway") and svc["wake_delay"]: + # Model the real hazard: is_sleeping flips to false *during* + # /wake_up, while reload_weights has not happened yet. A + # router that dies in this window leaves a backend that looks + # perfectly awake and serves garbage (E2E case 12). + await asyncio.sleep(svc["wake_delay"] / 2) + svc["sleeping"] = False + await asyncio.sleep(svc["wake_delay"] / 2) + return self._response(200) + await asyncio.sleep(svc["wake_delay"]) + svc["sleeping"] = False + return self._response(200) + + if path == "/collective_rpc": + svc["reload_calls"] += 1 + return self._response(200) + + if path == "/reset_prefix_cache": + svc["reset_calls"] += 1 + return self._response(200) + + if path == "/sleep": + svc["sleep_calls"] += 1 + if svc["sleep_fails"] > 0: + svc["sleep_fails"] -= 1 + return self._response(500, payload=b"sleep failed") + svc["sleeping"] = True + svc["level"] = int(params.get("level", "1")) + return self._response(200) + + if path.startswith("/v1/"): + svc["api_calls"] += 1 + if svc.get("api_delay"): + await asyncio.sleep(svc["api_delay"]) + ctype = request.headers.get("content-type", "") + if "text/event-stream" in ctype or self.stream_chunks: + async def gen(): + for chunk in self.stream_chunks: + await asyncio.sleep(self.stream_delay) + yield chunk + return httpx.Response(200, content=gen(), + headers={"content-type": "text/event-stream"}) + body = (request.content or b"").decode("utf-8", "replace") + return self._response( + 200, + payload=json.dumps({ + "echo_path": path, + "echo_method": request.method, + "echo_body": body, + "echo_content_type": ctype, + "echo_query": str(request.url.params), + "service": host, + }).encode(), + content_type="application/json", + ) + + return self._response(404, payload=json.dumps( + {"detail": f"no route {path}"}).encode(), content_type="application/json") + + +@pytest.fixture +def backend() -> FakeVLLM: + return FakeVLLM() + + +@pytest.fixture +def cfg(): + config = clone(load_config()) + config.state_cache_ttl = 0.0 # never trust cache -> deterministic + config.wake_health_poll_s = 0.005 + config.idle_enabled = False # idle is exercised via idle_tick() + config.hold_sleep_s = 5.0 + config.hold_offload_s = 5.0 + config.hold_restart_s = 5.0 + return config + + +@pytest.fixture +async def stack(cfg, backend): + public_app, admin_app, manager = build_apps(cfg, transport=backend) + try: + yield SimpleNamespace(public=public_app, admin=admin_app, manager=manager, + cfg=cfg, backend=backend) + finally: + await manager.proxy_client.aclose() + await manager.ctrl_client.aclose() + + +@pytest.fixture +async def pub(stack): + transport = httpx.ASGITransport(app=stack.public) + async with httpx.AsyncClient(transport=transport, base_url="http://router.test") as client: + yield client + + +@pytest.fixture +async def adm(stack): + transport = httpx.ASGITransport(app=stack.admin) + async with httpx.AsyncClient(transport=transport, base_url="http://router.test") as client: + yield client + + +async def raw_asgi(app, method: str, raw_path: str, *, headers=None, body: bytes = b""): + """Call the ASGI app with a hand-built scope so the *raw* (percent-encoded) + path reaches the router exactly as a hostile client would send it.""" + scope = { + "type": "http", "asgi": {"version": "3.0", "spec_version": "2.3"}, + "method": method, + "path": urllib.parse.unquote(raw_path), + "raw_path": raw_path.encode("latin-1"), + "query_string": b"", + "headers": [(k.lower().encode(), v.encode()) for k, v in (headers or [])], + "http_version": "1.1", "scheme": "http", + "server": ("router.test", 80), "client": ("127.0.0.1", 1234), + "root_path": "", + } + sent: list[dict] = [] + incoming: list[dict] = [{"type": "http.request", "body": body, "more_body": False}] + closed = asyncio.Event() + + async def receive(): + # Behaves like a real server: hand over the body once, then wait for + # the client to hang up (Starlette's disconnect listener relies on it). + if incoming: + return incoming.pop(0) + await closed.wait() + return {"type": "http.disconnect"} + + async def send(message): + sent.append(message) + + await app(scope, receive, send) + closed.set() + status = sent[0]["status"] + payload = b"".join(m.get("body", b"") for m in sent if m["type"] == "http.response.body") + return status, payload diff --git a/router/tests/test_admin_api.py b/router/tests/test_admin_api.py new file mode 100644 index 0000000..1fb5a97 --- /dev/null +++ b/router/tests/test_admin_api.py @@ -0,0 +1,120 @@ +"""Admin API surface (this is what `vllmctl` talks to).""" + +from __future__ import annotations + +import time + +from config import DEPTH_AWAKE, DEPTH_OFFLOADED + + +async def test_status_shape(adm, backend): + backend.set_sleeping("ocr", True, level=1) + r = await adm.get("/admin/status") + assert r.status_code == 200 + body = r.json() + assert set(body["services"]) == {"text", "ocr", "embed"} + text = body["services"]["text"] + for field in ("service", "model", "base_url", "reachable", "sleeping", + "depth", "wake_in_progress", "active_requests", + "last_activity_ago_s", "last_wake_latency_s", "last_error"): + assert field in text, field + assert text["model"] == "Qwen3.6-35B-A3B-FP8" + assert text["base_url"] == "http://vllm-text:8000" + assert body["services"]["ocr"]["sleeping"] is True + assert body["services"]["ocr"]["depth"] == "offloaded" # unknown level + assert "http://127.0.0.1:8000/v1" in body["api"] + + +async def test_wake_endpoint(adm, backend): + backend.set_sleeping("embed", True) + r = await adm.post("/admin/wake/embed") + assert r.status_code == 200 + body = r.json() + assert body["ok"] is True + assert body["depth"] == "awake" + assert backend.services["vllm-embed"]["wake_up_calls"] == 1 + + +async def test_wake_endpoint_reports_503_when_it_cannot_wake(adm, backend, stack): + backend.set_sleeping("text", True) + backend.services["vllm-text"]["wake_fails"] = 5 + r = await adm.post("/admin/wake/text") + assert r.status_code == 503 + assert r.headers["retry-after"] == "60" + body = r.json() + assert body["ok"] is False + assert body["error"]["sleep_depth"] == "offloaded" + assert body["error"]["estimated_wake_seconds"] == 60 + + +async def test_wake_accepts_model_name_and_case(adm, backend): + backend.set_sleeping("ocr", True) + r = await adm.post("/admin/wake/OvisOCR2") + assert r.status_code == 200 + assert backend.services["vllm-ocr"]["wake_up_calls"] == 1 + + +async def test_wake_unknown_key_is_404(adm): + assert (await adm.post("/admin/wake/nope")).status_code == 404 + + +async def test_sleep_level1_and_level2(adm, backend, stack): + r = await adm.post("/admin/sleep/text?level=1") + assert r.status_code == 200 + assert r.json()["ok"] is True + assert "level=1" in backend.calls[backend.last("POST vllm-text/sleep")] + assert stack.manager.services["text"].depth == "sleeping" + + r = await adm.post("/admin/sleep/text?level=2") + assert r.status_code == 200 + assert "level=2" in backend.calls[backend.last("POST vllm-text/sleep")] + assert stack.manager.services["text"].depth == DEPTH_OFFLOADED + + +async def test_sleep_defaults_to_level1(adm, backend): + r = await adm.post("/admin/sleep/embed") + assert r.status_code == 200 + assert "level=1" in backend.calls[backend.last("POST vllm-embed/sleep")] + + +async def test_sleep_rejects_bad_level(adm): + assert (await adm.post("/admin/sleep/text?level=3")).status_code == 400 + assert (await adm.post("/admin/sleep/text?level=abc")).status_code == 400 + + +async def test_sleep_unknown_key_is_404(adm, backend): + assert (await adm.post("/admin/sleep/nope")).status_code == 404 + assert backend.calls == [] + + +async def test_sleep_idempotent_when_already_at_depth(adm, backend, stack): + await adm.post("/admin/sleep/ocr?level=1") + backend.calls.clear() + r = await adm.post("/admin/sleep/ocr?level=1") + assert r.status_code == 200 + assert r.json()["already"] is True + assert backend.count("POST vllm-ocr/sleep") == 0 # no second POST /sleep + + +async def test_admin_wake_resets_the_idle_clock(adm, backend, stack): + backend.set_sleeping("text", True) + text = stack.manager.services["text"] + text.last_activity -= 10_000 + before = time.monotonic() - text.last_activity + assert (await adm.post("/admin/wake/text")).status_code == 200 + assert text.depth == DEPTH_AWAKE + assert (time.monotonic() - text.last_activity) < before + + +async def test_admin_health(adm): + r = await adm.get("/health") + assert r.status_code == 200 + assert r.json()["status"] == "ok" + + +async def test_admin_health_alias(adm): + """The documented admin surface says /admin/health.""" + assert (await adm.get("/admin/health")).status_code == 200 + body = (await adm.get("/admin/health")).json() + assert body["status"] == "ok" + assert "awake" in body diff --git a/router/tests/test_depth_errors.py b/router/tests/test_depth_errors.py new file mode 100644 index 0000000..a5bb161 --- /dev/null +++ b/router/tests/test_depth_errors.py @@ -0,0 +1,207 @@ +"""Depth-aware 503 semantics (plan 6.2.1) and never-proxy-half-awake.""" + +from __future__ import annotations + +import asyncio + +import httpx + +from conftest import raw_asgi + + +async def _sleep_at_level(stack, level: int) -> None: + result = await stack.manager.sleep_service("text", level, reason="test") + assert result["ok"], result + + +async def test_503_from_level1_sleep(pub, backend, stack): + await _sleep_at_level(stack, 1) + stack.cfg.hold_sleep_s = 0.05 + backend.services["vllm-text"]["wake_delay"] = 0.4 + r = await pub.post("/v1/chat/completions", json={"model": "text"}) + assert r.status_code == 503 + assert r.headers["retry-after"] == "10" + err = r.json()["error"] + assert err["code"] == "model_waking" + assert err["type"] == "model_waking" + assert err["sleep_depth"] == "sleeping" + assert isinstance(err["estimated_wake_seconds"], int) + assert err["estimated_wake_seconds"] == 6 + assert "Qwen3.6-35B-A3B-FP8" in err["message"] + + +async def test_503_from_level2_offload(pub, backend, stack): + await _sleep_at_level(stack, 2) + stack.cfg.hold_offload_s = 0.05 + backend.services["vllm-text"]["wake_delay"] = 0.4 + r = await pub.post("/v1/chat/completions", json={"model": "text"}) + assert r.status_code == 503 + assert r.headers["retry-after"] == "60" + err = r.json()["error"] + assert err["sleep_depth"] == "offloaded" + assert err["estimated_wake_seconds"] == 60 + + +async def test_503_when_container_restarting(pub, backend, stack): + backend.services["vllm-text"]["reachable"] = False + stack.cfg.hold_restart_s = 0.05 + r = await pub.post("/v1/chat/completions", json={"model": "text"}) + assert r.status_code == 503 + assert r.headers["retry-after"] == "600" + err = r.json()["error"] + assert err["sleep_depth"] == "restarting" + assert err["estimated_wake_seconds"] == 600 + + +async def test_unknown_depth_is_conservative_offloaded(pub, backend, stack): + """Router restarted / slept behind our back: depth unknown -> offloaded.""" + backend.set_sleeping("ocr", True, level=1) # actually only level 1 asleep + stack.cfg.hold_offload_s = 0.05 + backend.services["vllm-ocr"]["wake_delay"] = 0.3 + r = await pub.post("/v1/chat/completions", json={"model": "ocr"}) + assert r.status_code == 503 + err = r.json()["error"] + # conservative: worst-case depth, longest client wait + assert err["sleep_depth"] == "offloaded" + assert r.headers["retry-after"] == "60" + + +async def test_wake_sequence_retried_once_then_503(pub, backend, stack): + backend.set_sleeping("text", True) + backend.services["vllm-text"]["wake_fails"] = 2 + r = await pub.post("/v1/chat/completions", json={"model": "text"}) + assert r.status_code == 503 + assert r.headers["retry-after"] == "60" + err = r.json()["error"] + assert err["code"] == "model_waking" + assert err["sleep_depth"] == "offloaded" + # exactly one retry (plan 6.2.1) + assert backend.count("POST vllm-text/wake_up") == 2 + # never proxied to a half-awake backend + assert backend.count("POST vllm-text/v1/") == 0 + + +async def test_never_proxy_before_health_ok(pub, backend): + """The /v1 call must happen after the wake sequence, never before it. + + Requests admitted between wake_up and reload_weights return 200 + garbage, + so the whole sequence has to finish first. + """ + backend.set_sleeping("text", True) + backend.services["vllm-text"]["wake_delay"] = 0.1 + r = await pub.post("/v1/chat/completions", json={"model": "text"}) + assert r.status_code == 200 + svc = backend.services["vllm-text"] + assert svc["wake_up_calls"] == 1 + assert svc["reload_calls"] == 1 + assert svc["reset_calls"] == 1 + calls = backend.calls + idx = {needle: backend.last(needle) for needle in ( + "POST vllm-text/wake_up", + "POST vllm-text/collective_rpc", + "POST vllm-text/reset_prefix_cache", + "POST vllm-text/v1/chat/completions", + )} + assert idx["POST vllm-text/v1/chat/completions"] > idx["POST vllm-text/reset_prefix_cache"] + assert idx["POST vllm-text/reset_prefix_cache"] > idx["POST vllm-text/collective_rpc"] + assert idx["POST vllm-text/collective_rpc"] > idx["POST vllm-text/wake_up"] + assert calls[-1] == "POST vllm-text/v1/chat/completions" + + +async def test_level1_wake_uses_the_fast_path(pub, backend, stack): + """From level-1 sleep, /wake_up ALONE is enough (calibration 2026-08-17): + bit-identical output at temp 0, and ~20s cheaper than the reload sequence + (23.4s -> 2.5-3.8s on the text model).""" + assert (await stack.manager.sleep_service("text", 1))["ok"] is True + assert stack.manager.services["text"].depth == "sleeping" + backend.calls.clear() + svc = backend.services["vllm-text"] + for field in ("wake_up_calls", "reload_calls", "reset_calls"): + svc[field] = 0 + + r = await pub.post("/v1/chat/completions", json={"model": "text"}) + assert r.status_code == 200 + assert svc["wake_up_calls"] == 1 # wake_up only ... + assert svc["reload_calls"] == 0 # ... no reload_weights ... + assert svc["reset_calls"] == 0 # ... and no prefix-cache reset + assert backend.count("POST vllm-text/v1/chat/completions") == 1 + + +async def test_unknown_depth_uses_the_full_sequence(pub, backend, stack): + """Router restarted / slept out of band: unknown depth -> conservative + level-2 treatment, full sequence.""" + backend.set_sleeping("text", True) # router depth stays unknown + r = await pub.post("/v1/chat/completions", json={"model": "text"}) + assert r.status_code == 200 + svc = backend.services["vllm-text"] + assert (svc["wake_up_calls"], svc["reload_calls"], svc["reset_calls"]) == (1, 1, 1) + + +async def test_readiness_is_gated_on_is_sleeping_not_health(pub, backend, stack): + """/health answers 200 on a sleeping backend, so /is_sleeping is the gate; + nothing is proxied while the backend still reports is_sleeping=true (a + request sent to a sleeping backend hangs instead of erroring).""" + assert (await stack.manager.sleep_service("text", 2))["ok"] is True + backend.calls.clear() + backend.services["vllm-text"]["hold_sleeping_polls"] = 3 # wake "in flight" + + r = await pub.post("/v1/chat/completions", json={"model": "text"}) + assert r.status_code == 200 + calls = backend.calls + probes = [i for i, c in enumerate(calls) if c == "GET vllm-text/is_sleeping"] + assert len(probes) >= 4 # kept polling until it flipped + assert calls[-1] == "POST vllm-text/v1/chat/completions" # proxied last + assert backend.services["vllm-text"]["api_calls"] == 1 + + +async def test_backend_dies_between_health_and_proxy(pub, backend, stack): + """Transport error mid-proxy -> service marked restarting -> depth 503.""" + text = backend.services["vllm-text"] + + original = backend.handle_async_request + + async def flaky(request: httpx.Request) -> httpx.Response: + if request.url.path.startswith("/v1/"): + raise httpx.ConnectError("backend gone", request=request) + return await original(request) + + backend.handle_async_request = flaky + stack.cfg.hold_restart_s = 0.05 + r = await pub.post("/v1/chat/completions", json={"model": "text"}) + assert r.status_code == 503 + assert r.headers["retry-after"] == "600" + assert r.json()["error"]["sleep_depth"] == "restarting" + assert text["wake_up_calls"] == 0 + + +async def test_wake_after_level2_runs_full_sequence(pub, backend): + backend.set_sleeping("embed", True, level=2) + r = await pub.post("/v1/embeddings", json={"input": "hi"}) + assert r.status_code == 200 + svc = backend.services["vllm-embed"] + assert svc["wake_up_calls"] == 1 + assert svc["reload_calls"] == 1 # reload_weights is mandatory after L2 + assert svc["reset_calls"] == 1 # prefix cache reset too + assert svc["api_calls"] == 1 + + +async def test_error_body_is_openai_shaped(pub, backend, stack): + await _sleep_at_level(stack, 1) + stack.cfg.hold_sleep_s = 0.01 + backend.services["vllm-text"]["wake_delay"] = 0.2 + r = await pub.post("/v1/chat/completions", json={"model": "text"}) + payload = r.json() + assert set(payload) == {"error"} + assert set(payload["error"]) == { + "type", "code", "message", "sleep_depth", "estimated_wake_seconds" + } + assert r.headers["content-type"].startswith("application/json") + + +async def test_depth_survives_raw_traversal_requests(stack, backend): + """Traversal requests are rejected before any backend contact.""" + backend.set_sleeping("text", True) + for raw in ("/v1/../sleep", "/v1%2f..%2fsleep", "//sleep", "/v1/../../wake_up"): + status, _ = await raw_asgi(stack.public, "POST", raw) + assert status == 404, raw + assert backend.calls == [] diff --git a/router/tests/test_idle.py b/router/tests/test_idle.py new file mode 100644 index 0000000..4194d7e --- /dev/null +++ b/router/tests/test_idle.py @@ -0,0 +1,195 @@ +"""Tiered idle management: thresholds, lock races, active-request guard.""" + +from __future__ import annotations + +import asyncio +import time + +from config import DEPTH_AWAKE, DEPTH_OFFLOADED, DEPTH_SLEEPING + + +def _age(svc, seconds: float) -> None: + svc.last_activity = time.monotonic() - seconds + + +async def test_idle_level1_after_threshold(stack, backend): + stack.cfg.idle_sleep_min = 15 / 60.0 # 15 s in "minutes" + stack.cfg.idle_offload_min = 180 / 60.0 + text = stack.manager.services["text"] + text.depth = DEPTH_AWAKE + _age(text, 20) + + await stack.manager.idle_tick() + + assert backend.count("POST vllm-text/sleep") == 1 + assert "level=1" in backend.calls[backend.last("POST vllm-text/sleep")] + assert text.depth == DEPTH_SLEEPING + + +async def test_idle_escalates_to_level2(stack, backend): + stack.cfg.idle_sleep_min = 15 / 60.0 + stack.cfg.idle_offload_min = 180 / 60.0 + text = stack.manager.services["text"] + text.depth = DEPTH_AWAKE + _age(text, 400) + + await stack.manager.idle_tick() + assert "level=2" in backend.calls[backend.last("POST vllm-text/sleep")] + assert text.depth == DEPTH_OFFLOADED + + +async def test_level1_service_escalates_to_level2(stack, backend): + """Already napping (level 1) and still idle -> escalate to offload. + + A direct POST /sleep?level=2 on a level-1-sleeping backend is a + well-behaved NO-OP that retains the host-RAM copy (calibration + 2026-08-17), so the escalation must wake into RAM first and then offload. + """ + stack.cfg.idle_sleep_min = 1 / 60.0 + stack.cfg.idle_offload_min = 5 / 60.0 + text = stack.manager.services["text"] + assert (await stack.manager.sleep_service("text", 1))["ok"] is True + backend.calls.clear() + svc = backend.services["vllm-text"] + svc["wake_up_calls"] = svc["sleep_calls"] = 0 + _age(text, 400) + + await stack.manager.idle_tick() + assert text.depth == DEPTH_OFFLOADED + assert svc["wake_up_calls"] == 1 # wake into RAM ... + assert svc["sleep_calls"] == 1 # ... then offload + assert backend.last("POST vllm-text/sleep?level=2") > backend.last("POST vllm-text/wake_up") + + +async def test_offload_from_awake_is_direct(stack, backend): + """Only from depth 'awake' can level 2 be entered directly.""" + backend.services["vllm-embed"]["wake_up_calls"] = 0 + result = await stack.manager.sleep_service("embed", 2) + assert result["ok"] is True + assert backend.services["vllm-embed"]["wake_up_calls"] == 0 + assert backend.count("POST vllm-embed/sleep?level=2") == 1 + assert stack.manager.services["embed"].depth == DEPTH_OFFLOADED + + +async def test_escalation_failure_is_reported(stack, backend): + assert (await stack.manager.sleep_service("ocr", 1))["ok"] is True + backend.services["vllm-ocr"]["wake_fails"] = 1 + backend.calls.clear() + result = await stack.manager.sleep_service("ocr", 2) + assert result["ok"] is False + assert result["reason"] == "escalation_failed" + assert backend.count("POST vllm-ocr/sleep") == 0 + assert stack.manager.services["ocr"].depth == DEPTH_SLEEPING + + +async def test_offloaded_service_is_left_alone(stack, backend): + stack.cfg.idle_offload_min = 1 / 60.0 + text = stack.manager.services["text"] + text.depth = DEPTH_OFFLOADED + _age(text, 10_000) + await stack.manager.idle_tick() + assert backend.calls == [] + + +async def test_active_requests_block_idle_sleep(stack, backend): + stack.cfg.idle_sleep_min = 1 / 60.0 + text = stack.manager.services["text"] + text.depth = DEPTH_AWAKE + _age(text, 600) + + manager = stack.manager + manager.begin_request(text) # long generation in flight + try: + await manager.idle_tick() + assert backend.calls == [] + finally: + manager.end_request(text) + + +async def test_stream_completes_then_idle_can_sleep(stack, backend, pub): + """end_request (background task of the streamed response) re-opens sleep.""" + stack.cfg.idle_sleep_min = 1 / 60.0 + text = stack.manager.services["text"] + r = await pub.post("/v1/chat/completions", json={"model": "text"}) + assert r.status_code == 200 + assert text.active_requests == 0 # released after the body drained + + _age(text, 600) + await stack.manager.idle_tick() + assert backend.count("POST vllm-text/sleep") == 1 + + +async def test_race_last_activity_refreshed_under_lock(stack, backend): + """A request landed between the threshold check and the locked re-check.""" + stack.cfg.idle_sleep_min = 1 / 60.0 # threshold = 60 s + text = stack.manager.services["text"] + text.depth = DEPTH_AWAKE + _age(text, 600) # what the idle scan saw + + # ... and then a request arrived, refreshing last_activity *now* + text.last_activity = time.monotonic() + result = await stack.manager.sleep_service("text", 1, reason="idle", min_idle_s=60.0) + assert result["ok"] is False + assert result["reason"] == "activity_resumed" + assert backend.count("POST vllm-text/sleep") == 0 + + +async def test_race_active_request_seen_under_lock(stack, backend): + text = stack.manager.services["text"] + text.depth = DEPTH_AWAKE + _age(text, 600) + text.active_requests = 2 # arrived while we were scanning + result = await stack.manager.sleep_service("text", 1, reason="idle", min_idle_s=60.0) + assert result["ok"] is False + assert result["reason"] == "active_requests" + assert backend.calls == [] + text.active_requests = 0 + + +async def test_request_at_the_moment_the_timer_expires(pub, backend, stack): + """E2E case 11 in miniature: the request wins, no sleep mid-flight.""" + stack.cfg.idle_sleep_min = 1 / 60.0 + stack.cfg.idle_offload_min = 5 / 60.0 + text = stack.manager.services["text"] + text.depth = DEPTH_AWAKE + _age(text, 600) + + async def tick_soon(): + await asyncio.sleep(0) # run after the request started + return await stack.manager.idle_tick() + + tick, response = await asyncio.gather(tick_soon(), pub.post( + "/v1/chat/completions", json={"model": "text"})) + assert response.status_code == 200 + assert backend.count("POST vllm-text/sleep") == 0 + + +async def test_admin_sleep_refuses_when_active(adm, stack, backend): + text = stack.manager.services["text"] + text.depth = DEPTH_AWAKE + stack.manager.begin_request(text) + r = await adm.post("/admin/sleep/text?level=1") + assert r.status_code == 409 + assert r.json()["reason"] == "active_requests" + assert backend.count("/sleep") == 0 + stack.manager.end_request(text) + + +async def test_idle_loop_runs_in_background_when_started(stack, backend): + stack.cfg.idle_enabled = True + stack.cfg.idle_poll_s = 0.01 + stack.cfg.idle_sleep_min = 1 / 60.0 # 60 s + stack.cfg.idle_offload_min = 1000.0 # far away: expect the level-1 tier + text = stack.manager.services["text"] + text.depth = DEPTH_AWAKE + _age(text, 600) + stack.manager.start() + try: + for _ in range(100): + if backend.count("POST vllm-text/sleep"): + break + await asyncio.sleep(0.01) + finally: + await stack.manager.stop() + assert backend.count("POST vllm-text/sleep") == 1 + assert "level=1" in backend.calls[backend.last("POST vllm-text/sleep")] diff --git a/router/tests/test_model_resolution.py b/router/tests/test_model_resolution.py new file mode 100644 index 0000000..f6ceedb --- /dev/null +++ b/router/tests/test_model_resolution.py @@ -0,0 +1,160 @@ +"""Model resolution: JSON / multipart / defaults / embeddings / unknown.""" + +from __future__ import annotations + +import json + +TEXT = "Qwen3.6-35B-A3B-FP8" +OCR = "OvisOCR2" +EMBED = "Qwen3-Embedding-8B" + + +async def test_json_model_exact(pub, backend): + r = await pub.post("/v1/chat/completions", json={"model": OCR, "messages": []}) + assert r.status_code == 200 + assert r.json()["service"] == "vllm-ocr" + + +async def test_json_model_case_insensitive(pub): + r = await pub.post("/v1/chat/completions", json={"model": "qwen3.6-35b-a3b-fp8"}) + assert r.status_code == 200 + assert r.json()["service"] == "vllm-text" + + +async def test_alias_matches(pub): + r = await pub.post("/v1/chat/completions", json={"model": "ocr"}) + assert r.status_code == 200 + assert r.json()["service"] == "vllm-ocr" + r = await pub.post("/v1/chat/completions", json={"model": "Embedding"}) + assert r.json()["service"] == "vllm-embed" + + +async def test_missing_model_defaults_to_text_on_chat(pub): + r = await pub.post("/v1/chat/completions", json={"messages": [{"role": "user", "content": "hi"}]}) + assert r.status_code == 200 + assert r.json()["service"] == "vllm-text" + + +async def test_missing_model_defaults_to_text_on_completions(pub): + r = await pub.post("/v1/completions", json={"prompt": "hi"}) + assert r.status_code == 200 + assert r.json()["service"] == "vllm-text" + + +async def test_missing_model_is_400_elsewhere(pub, backend): + r = await pub.post("/v1/rerank", json={"query": "hi"}) + assert r.status_code == 400 + assert r.json()["error"]["code"] == "missing_model" + assert backend.count("/v1/rerank") == 0 + + +async def test_invalid_json_is_400(pub, backend): + r = await pub.post("/v1/chat/completions", + content=b"{not json", + headers={"content-type": "application/json"}) + assert r.status_code == 400 + assert backend.count("/v1/") == 0 + + +async def test_embeddings_always_routes_to_embed(pub): + r = await pub.post("/v1/embeddings", json={"input": "hello"}) + assert r.status_code == 200 + assert r.json()["service"] == "vllm-embed" + # ... even when the body names a different model + r = await pub.post("/v1/embeddings", json={"input": "hello", "model": TEXT}) + assert r.status_code == 200 + assert r.json()["service"] == "vllm-embed" + + +async def test_unknown_model_404_and_no_wake(pub, backend): + backend.set_sleeping("ocr", True) + r = await pub.post("/v1/chat/completions", json={"model": "gpt-4o"}) + assert r.status_code == 404 + body = r.json()["error"] + assert body["code"] == "model_not_found" + assert body["type"] == "invalid_request_error" + assert body["param"] == "model" + # no probe, no wake, no proxy + assert backend.calls == [] + + +async def test_multipart_with_model_field(pub): + r = await pub.post( + "/v1/chat/completions", + data={"model": OCR}, + files={"image": ("page.png", b"PNGDATA" * 64, "image/png")}, + ) + assert r.status_code == 200 + echo = r.json() + assert echo["service"] == "vllm-ocr" + # raw body forwarded unchanged: the file bytes and the boundary survive + assert "PNGDATA" * 64 in echo["echo_body"] + assert echo["echo_content_type"].startswith("multipart/form-data; boundary=") + + +async def test_multipart_without_model_defaults_to_ocr(pub): + r = await pub.post( + "/v1/chat/completions", + files={"image": ("page.png", b"PNGDATA", "image/png")}, + ) + assert r.status_code == 200 + assert r.json()["service"] == "vllm-ocr" + + +async def test_multipart_unknown_model_404(pub, backend): + r = await pub.post( + "/v1/chat/completions", + data={"model": "nope"}, + files={"image": ("page.png", b"x", "image/png")}, + ) + assert r.status_code == 404 + assert r.json()["error"]["code"] == "model_not_found" + assert backend.calls == [] + + +async def test_multipart_model_extraction_ignores_file_parts(pub): + # a file part literally named "model" must not be read as the model field + r = await pub.post( + "/v1/chat/completions", + data={"prompt": "hi"}, + files={"model": ("fake.json", b"NOT-A-MODEL-NAME", "application/octet-stream")}, + ) + # no usable model field -> multipart default on the chat path is OCR + assert r.status_code == 200 + assert r.json()["service"] == "vllm-ocr" + assert "NOT-A-MODEL-NAME" in r.json()["echo_body"] + + +async def test_model_in_path_for_models_detail(pub): + r = await pub.get(f"/v1/models/{OCR}") + assert r.status_code == 200 + assert r.json()["service"] == "vllm-ocr" + r = await pub.get("/v1/models/does-not-exist") + assert r.status_code == 404 + assert r.json()["error"]["code"] == "model_not_found" + + +async def test_body_and_headers_forwarded(pub): + payload = {"model": TEXT, "messages": [{"role": "user", "content": "hi"}], "stream": False} + r = await pub.post("/v1/chat/completions", json=payload, + headers={"Authorization": "Bearer x", "x-custom": "1"}) + assert r.status_code == 200 + echo = r.json() + assert echo["echo_path"] == "/v1/chat/completions" + assert echo["echo_method"] == "POST" + assert json.loads(echo["echo_body"]) == payload + + +async def test_query_string_forwarded(pub): + r = await pub.post("/v1/chat/completions?foo=bar&baz=1", json={"model": TEXT}) + assert r.status_code == 200 + assert r.json()["echo_query"] == "foo=bar&baz=1" + + +async def test_streaming_passthrough(pub, backend): + backend.stream_chunks = [b"data: {\"a\":1}\n\n", b"data: {\"a\":2}\n\n", b"data: [DONE]\n\n"] + r = await pub.post("/v1/chat/completions", json={"model": TEXT, "stream": True}) + assert r.status_code == 200 + assert r.headers["content-type"].startswith("text/event-stream") + assert r.text == "".join(chunk.decode() for chunk in backend.stream_chunks) + backend.stream_chunks = [] diff --git a/router/tests/test_paths.py b/router/tests/test_paths.py new file mode 100644 index 0000000..7018677 --- /dev/null +++ b/router/tests/test_paths.py @@ -0,0 +1,146 @@ +"""Allowlist: only /v1/*, /health (and /metrics when enabled) are public.""" + +from __future__ import annotations + +import httpx + +from app import build_apps +from conftest import raw_asgi + +DEV_ENDPOINTS = [ + "/sleep", + "/wake_up", + "/collective_rpc", + "/reset_prefix_cache", + "/is_sleeping", +] +ADMIN_PATHS = ["/admin/status", "/admin/wake/text", "/admin/sleep/text?level=1"] + + +async def test_dev_endpoints_404_on_public(pub, backend): + for path in DEV_ENDPOINTS: + r = await pub.post(path) + assert r.status_code == 404, path + assert r.json()["error"]["code"] == "not_found" + for path in DEV_ENDPOINTS: + r = await pub.get(path) + assert r.status_code == 404, path + assert backend.calls == [] + + +async def test_admin_not_on_public(pub, backend): + for path in ADMIN_PATHS: + r = await pub.request("POST", path) + assert r.status_code == 404, path + r = await pub.get(path) + assert r.status_code == 404, path + assert backend.calls == [] + + +async def test_docs_and_openapi_404(pub): + for path in ("/docs", "/redoc", "/openapi.json", "/"): + r = await pub.get(path) + assert r.status_code == 404, path + + +async def test_dotted_paths_404(pub): + for path in ("/v1/../sleep", "/v1/..%2fsleep", "/health/../sleep", "/./sleep"): + r = await pub.post(path) + assert r.status_code == 404, path + + +async def test_raw_traversal_404(stack, backend): + """Percent-encoded traversal that a normal client would normalise.""" + cases = [ + ("POST", "/v1%2f..%2fsleep"), + ("POST", "/v1%2F..%2Fsleep"), + ("GET", "/v1/%2e%2e/wake_up"), + ("POST", "/v1/%2e%2e%2fcollective_rpc"), + ("POST", "/v1//../sleep"), + ("GET", "/v1/models/../../is_sleeping"), + ] + for method, raw in cases: + status, body = await raw_asgi(stack.public, method, raw, + headers=[("content-type", "application/json")], + body=b"{}") + assert status == 404, (method, raw, body) + assert backend.calls == [] + + +async def test_normalised_traversal_still_reaches_v1(stack, backend): + """A traversal that lands inside /v1 must still work (not over-block).""" + status, body = await raw_asgi( + stack.public, "POST", "/v1/../v1/chat/completions", + headers=[("content-type", "application/json")], + body=b'{"model": "text"}', + ) + assert status == 200, body + assert backend.count("POST vllm-text/v1/chat/completions") == 1 + + +async def test_public_health(pub, backend): + r = await pub.get("/health") + assert r.status_code == 200 + body = r.json() + assert body["status"] == "ok" + assert set(body["services"]) == {"text", "ocr", "embed"} + # liveness must not depend on backend state + backend.services["vllm-text"]["reachable"] = False + r = await pub.get("/health") + assert r.status_code == 200 + + +async def test_public_models_lists_all_three(pub, backend): + r = await pub.get("/v1/models") + assert r.status_code == 200 + body = r.json() + assert body["object"] == "list" + assert [m["id"] for m in body["data"]] == [ + "Qwen3.6-35B-A3B-FP8", "OvisOCR2", "Qwen3-Embedding-8B", + ] + # answered by the router itself, no backend contact + assert backend.calls == [] + + +async def test_metrics_disabled_by_default(pub): + assert (await pub.get("/metrics")).status_code == 404 + + +async def test_metrics_enabled_when_configured(stack, pub): + stack.cfg.metrics_enabled = True + r = await pub.get("/metrics") + assert r.status_code == 200 + assert "vllm_router_service_depth" in r.text + + +async def test_unknown_http_method_on_v1(stack): + status, _ = await raw_asgi(stack.public, "PROPFIND", "/v1/chat/completions") + assert status in (404, 405) + + +async def test_admin_listener_is_a_separate_app(stack): + """The admin app answers /admin/*; the public app must not.""" + status, body = await raw_asgi(stack.admin, "GET", "/admin/status") + assert status == 200 + status, _ = await raw_asgi(stack.public, "GET", "/admin/status") + assert status == 404 + + +async def test_request_body_cap(cfg, backend): + """ROUTER_MAX_BODY_BYTES must be wired as middleware (413, not OOM).""" + cfg.max_body_bytes = 64 + public_app, _admin, manager = build_apps(cfg, transport=backend) + ok_body = b'{"model": "text", "messages": []}' + assert len(ok_body) < 64 + try: + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=public_app), + base_url="http://router.test") as client: + r = await client.post("/v1/chat/completions", content=ok_body, + headers={"content-type": "application/json"}) + assert r.status_code == 200 + r = await client.post("/v1/chat/completions", content=ok_body + b" " * 512, + headers={"content-type": "application/json"}) + assert r.status_code == 413 + finally: + await manager.proxy_client.aclose() + await manager.ctrl_client.aclose() diff --git a/router/tests/test_shutdown_and_cli.py b/router/tests/test_shutdown_and_cli.py new file mode 100644 index 0000000..e716d58 --- /dev/null +++ b/router/tests/test_shutdown_and_cli.py @@ -0,0 +1,101 @@ +"""Shutdown behaviour (E2E case 12) and the vllmctl CLI glue.""" + +from __future__ import annotations + +import asyncio +import json +import os +import subprocess + +import pytest + +from conftest import raw_asgi +from routing import ShutdownGuardMiddleware + + +async def test_shutdown_guard_returns_503_when_not_started(stack): + """uvicorn cancels in-flight tasks after timeout_graceful_shutdown; that + CancelledError must become a depth-aware 503, not uvicorn's bare 500.""" + stack.manager.begin_shutdown() + + async def app(scope, receive, send): + raise asyncio.CancelledError() + + guard = ShutdownGuardMiddleware(app, stack.manager) + status, body = await raw_asgi(guard, "POST", "/v1/chat/completions") + assert status == 503 + payload = json.loads(body) + assert payload["error"]["code"] == "router_shutting_down" + assert payload["error"]["retry_after_seconds"] == 5 + + +async def test_shutdown_guard_passes_through_when_not_shutting_down(stack): + class Boom(RuntimeError): + pass + + async def app(scope, receive, send): + raise Boom("nope") + + guard = ShutdownGuardMiddleware(app, stack.manager) + with pytest.raises(Boom): + await raw_asgi(guard, "POST", "/v1/chat/completions") + + +async def test_shutdown_guard_cannot_unsend_a_started_response(stack): + """A response that already started streaming is closed, not replaced.""" + stack.manager.begin_shutdown() + + async def app(scope, receive, send): + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"partial"}) + raise asyncio.CancelledError() + + guard = ShutdownGuardMiddleware(app, stack.manager) + with pytest.raises(asyncio.CancelledError): + await raw_asgi(guard, "POST", "/v1/chat/completions") + + +async def test_shutdown_guard_forwards_normal_requests(stack): + async def app(scope, receive, send): + await send({"type": "http.response.start", "status": 204, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + + guard = ShutdownGuardMiddleware(app, stack.manager) + status, _ = await raw_asgi(guard, "GET", "/health") + assert status == 204 + + +def test_vllmctl_service_name_mapping(): + """`logs router` / `restart router` must map to the compose service + `router`; the model services map to vllm-*.""" + script = ( + "source /data/home/renbaibing/vllm/vllmctl >/dev/null 2>&1; " + 'for n in router vllm-router Router text vllm-text TEXT ocr OvisOCR2 ' + "embed Qwen3-Embedding-8B bogus; do " + 'printf "%s:%s\\n" "$n" "$(docker_service_for "$n" || echo ERR)"; done' + ) + out = subprocess.run(["bash", "-c", script], capture_output=True, text=True, + timeout=30) + assert out.returncode == 0, out.stderr + mapping = dict(line.split(":", 1) for line in out.stdout.strip().splitlines()) + assert mapping["router"] == "router" + assert mapping["vllm-router"] == "router" + assert mapping["Router"] == "router" + assert mapping["text"] == "vllm-text" + assert mapping["vllm-text"] == "vllm-text" + assert mapping["TEXT"] == "vllm-text" + assert mapping["ocr"] == "vllm-ocr" + assert mapping["OvisOCR2"] == "vllm-ocr" + assert mapping["embed"] == "vllm-embed" + assert mapping["Qwen3-Embedding-8B"] == "vllm-embed" + assert mapping["bogus"] == "ERR" + + +def test_vllmctl_has_no_removed_machinery(): + source = open("/data/home/renbaibing/vllm/vllmctl", encoding="utf-8").read() + assert ".runas.py" not in source + assert ".user.env" not in source + assert "idle-watch on" not in source + assert "cmd_idle_watch" not in source + assert "admin/wake/" in source and "admin/sleep/" in source + assert os.access("/data/home/renbaibing/vllm/vllmctl", os.X_OK) diff --git a/router/tests/test_singleflight.py b/router/tests/test_singleflight.py new file mode 100644 index 0000000..ffc0359 --- /dev/null +++ b/router/tests/test_singleflight.py @@ -0,0 +1,95 @@ +"""Single-flight wake: N concurrent requests -> exactly one wake sequence.""" + +from __future__ import annotations + +import asyncio + +from config import DEPTH_AWAKE + + +async def test_ten_concurrent_requests_trigger_one_wake(pub, backend): + backend.set_sleeping("text", True, level=2) + backend.services["vllm-text"]["wake_delay"] = 0.15 + + async def one(i: int): + r = await pub.post("/v1/chat/completions", json={"model": "text", "seed": i}) + assert r.status_code == 200, r.text + return r.json() + + results = await asyncio.gather(*(one(i) for i in range(10))) + assert all(r["service"] == "vllm-text" for r in results) + + svc = backend.services["vllm-text"] + assert svc["wake_up_calls"] == 1, backend.calls + assert svc["reload_calls"] == 1 + assert svc["api_calls"] == 10 + + +async def test_hold_timeout_does_not_cancel_the_wake(pub, backend, stack): + """The first caller times out; the wake keeps going and finishes for the + next caller (no second wake sequence, no half-awake proxy).""" + backend.set_sleeping("text", True) + backend.services["vllm-text"]["wake_delay"] = 0.3 + stack.cfg.hold_offload_s = 0.05 + + r = await pub.post("/v1/chat/completions", json={"model": "text"}) + assert r.status_code == 503 + assert r.json()["error"]["sleep_depth"] == "offloaded" + + await asyncio.sleep(0.5) # let the shielded wake task finish + svc_state = stack.manager.services["text"] + assert svc_state.depth == DEPTH_AWAKE + assert backend.services["vllm-text"]["wake_up_calls"] == 1 + + r = await pub.post("/v1/chat/completions", json={"model": "text"}) + assert r.status_code == 200 + assert backend.services["vllm-text"]["wake_up_calls"] == 1 # cached awake + + +async def test_two_services_wake_concurrently_without_cross_talk(pub, backend): + backend.set_sleeping("text", True) + backend.set_sleeping("embed", True) + backend.services["vllm-text"]["wake_delay"] = 0.2 + backend.services["vllm-embed"]["wake_delay"] = 0.05 + + chat = pub.post("/v1/chat/completions", json={"model": "text"}) + embed = pub.post("/v1/embeddings", json={"input": "hi"}) + r1, r2 = await asyncio.gather(chat, embed) + assert r1.status_code == 200 and r1.json()["service"] == "vllm-text" + assert r2.status_code == 200 and r2.json()["service"] == "vllm-embed" + assert backend.services["vllm-text"]["wake_up_calls"] == 1 + assert backend.services["vllm-embed"]["wake_up_calls"] == 1 + + +async def test_wake_state_is_cached_between_requests(pub, backend, stack): + stack.cfg.state_cache_ttl = 30.0 + backend.set_sleeping("ocr", True) + assert (await pub.post("/v1/chat/completions", json={"model": "ocr"})).status_code == 200 + probes_after_first = backend.count("GET vllm-ocr/is_sleeping") + assert probes_after_first >= 1 + for _ in range(5): + assert (await pub.post("/v1/chat/completions", json={"model": "ocr"})).status_code == 200 + # no extra probe per request while the awake state is cached + assert backend.count("GET vllm-ocr/is_sleeping") == probes_after_first + + +async def test_expired_cache_reprobes(pub, backend, stack): + stack.cfg.state_cache_ttl = 0.0 + for _ in range(3): + assert (await pub.post("/v1/chat/completions", json={"model": "ocr"})).status_code == 200 + assert backend.count("GET vllm-ocr/is_sleeping") >= 3 + + +async def test_admin_and_request_path_share_one_lock(adm, pub, backend): + """A wake driven from the admin port is joined by the public request path.""" + backend.set_sleeping("text", True) + backend.services["vllm-text"]["wake_delay"] = 0.2 + + admin_wake = asyncio.create_task(adm.post("/admin/wake/text")) + await asyncio.sleep(0.05) + r = await pub.post("/v1/chat/completions", json={"model": "text"}) + admin_result = await admin_wake + + assert r.status_code == 200 + assert admin_result.status_code == 200 + assert backend.services["vllm-text"]["wake_up_calls"] == 1 diff --git a/router/tests/test_wake_recovery.py b/router/tests/test_wake_recovery.py new file mode 100644 index 0000000..7624b43 --- /dev/null +++ b/router/tests/test_wake_recovery.py @@ -0,0 +1,296 @@ +"""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 diff --git a/vllmctl b/vllmctl new file mode 100755 index 0000000..89f9e05 --- /dev/null +++ b/vllmctl @@ -0,0 +1,297 @@ +#!/usr/bin/env bash +# vllmctl — control the vLLM serving stack through the router's admin API. +# +# ./vllmctl status per-service state: awake / sleeping / +# offloaded, activity, wake in progress +# ./vllmctl up [MODEL] force-wake MODEL now (no MODEL = show +# status: all services are always running, +# the router wakes them on demand) +# ./vllmctl down [MODEL] offload MODEL (sleep level 2, frees host +# RAM); no MODEL = all services +# ./vllmctl sleep [MODEL] light sleep (level 1, weights stay in RAM) +# ./vllmctl list models on disk + which ones this stack serves +# ./vllmctl logs [-f] [N] [SVC] docker compose logs (SVC: text|ocr|embed|router) +# ./vllmctl restart [SVC] docker compose restart (rare, manual) +# ./vllmctl pull REPO [NAME] download a HF model into MODEL_ROOT +# +# MODEL accepts the service key (text/ocr/embed), the model name +# (Qwen3.6-35B-A3B-FP8, OvisOCR2, Qwen3-Embedding-8B) or an alias, +# case-insensitively; the router resolves them all. +# +# Wake/sleep/idle is owned by the router (127.0.0.1:8010 admin listener); +# `idle-watch` no longer exists. Public API: http://:8000/v1. +# +# docker: this user is in the docker group, so plain `docker` works in fresh +# login shells. In a session started before the group change, invoke this +# script as: sg docker -c "./vllmctl pull " + +set -u + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +C_GREEN=$'\e[32m'; C_RED=$'\e[31m'; C_YELLOW=$'\e[33m'; C_CYAN=$'\e[36m'; C_DIM=$'\e[2m'; C_OFF=$'\e[0m' +say() { printf '%s\n' "$*"; } +ok() { printf '%s%s%s\n' "$C_GREEN" "$*" "$C_OFF"; } +warn() { printf '%s%s%s\n' "$C_YELLOW" "$*" "$C_OFF"; } +err() { printf '%s%s%s\n' "$C_RED" "$*" "$C_OFF" >&2; } + +ADMIN="http://127.0.0.1:${VLLMCTL_ADMIN_PORT:-8010}" +PUBLIC_PORT="${VLLMCTL_PUBLIC_PORT:-8000}" +KEYS="text ocr embed" + +# ---------------------------------------------------------------- config ---- +env_get() { # env_get KEY [default] + local v + v="$(grep -E "^$1=" "$ROOT/.env" 2>/dev/null | tail -n1 | cut -d= -f2-)" + if [ -z "$v" ]; then printf '%s' "${2:-}"; else printf '%s' "$v"; fi +} + +MODEL_ROOT="$(env_get MODEL_ROOT /data/home/renbaibing/huggingface)" + +dcompose() { + docker compose --progress plain --ansi never \ + --project-directory "$ROOT" -f "$ROOT/compose.yml" --env-file "$ROOT/.env" "$@" +} + +# name -> docker compose SERVICE (for logs / restart only). +# NB: the router's compose service is `router` (container_name vllm-router); +# the vllm-* services are named after themselves. +docker_service_for() { + case "$(printf '%s' "$1" | tr '[:upper:]' '[:lower:]')" in + embed|embedding*|qwen3-embedding*) echo vllm-embed ;; + ocr|ovis*) echo vllm-ocr ;; + text|qwen*|chat|default) echo vllm-text ;; + router|vllm-router) echo router ;; + vllm-text|vllm-ocr|vllm-embed) echo "$1" ;; + *) return 1 ;; + esac +} + +# ---------------------------------------------------------------- admin ----- +admin_alive() { curl -fs -m 3 "$ADMIN/health" >/dev/null 2>&1; } + +RESPONSE_FILE="$(mktemp)"; trap 'rm -f "$RESPONSE_FILE"' EXIT + +admin_call() { # admin_call METHOD PATH TIMEOUT -> sets HTTP_CODE and RESPONSE + HTTP_CODE="$(curl -s -m "$3" -o "$RESPONSE_FILE" -w '%{http_code}' \ + -X "$1" "$ADMIN$2" 2>/dev/null)" + [ -n "$HTTP_CODE" ] || HTTP_CODE=000 + RESPONSE="$(cat "$RESPONSE_FILE")" +} + +json_field() { # json_field JSON PYEXPR -> value (python3 json) + printf '%s' "$1" | python3 -c " +import json, sys +try: + d = json.load(sys.stdin) +except Exception: + sys.exit(1) +print(eval(sys.argv[1]))" "$2" 2>/dev/null +} + +die_admin_down() { + err "Router admin API not reachable at $ADMIN" + err "Is the stack up? Try: docker compose -f $ROOT/compose.yml ps" + return 1 +} + +# ---------------------------------------------------------------- status ---- +cmd_status() { + admin_call GET /admin/status 10 + [ "$HTTP_CODE" = "200" ] || { die_admin_down; return 1; } + + # NB: the heredoc *is* stdin for `python3 -`, so the payload is passed as + # a file argument rather than piped. + python3 - "$PUBLIC_PORT" "$RESPONSE_FILE" <<'PY' +import json, sys + +port, payload_path = sys.argv[1], sys.argv[2] +with open(payload_path) as fh: + d = json.load(fh) +r = d["router"] +idle = r["idle"] +up = r.get("uptime_s") +up_txt = f"up {up:.0f}s" if up is not None and up < 3600 else f"up {up/3600:.1f}h" +tiers = f"idle tiers: sleep {idle['sleep_min']:g} min, offload {idle['offload_min']:g} min" +if not idle["enabled"]: + tiers += " (disabled)" + +G, Y, D, R, OFF = "\033[32m", "\033[33m", "\033[2m", "\033[31m", "\033[0m" +print(f"router {G}{up_txt}{OFF} {tiers}") +print(f"Public API : http://:{port}/v1 (OpenAI-compatible, all 3 models)") +print(f"Admin API : 127.0.0.1:{r['admin_port']} (localhost only)") +print() +print(f" {'SERVICE':<7} {'MODEL':<22} {'STATE':<11} {'ACTIVE':>6} {'IDLE':>8} NOTES") +for key in ("text", "ocr", "embed"): + s = d["services"][key] + state = s["depth"] or "unknown" + colour = {"awake": G, "offloaded": D}.get(state, Y) + idle_s = s["last_activity_ago_s"] + if idle_s is None: + idle_txt = "-" + elif idle_s >= 3600: + idle_txt = f"{idle_s/3600:.1f}h" + elif idle_s >= 60: + idle_txt = f"{idle_s/60:.0f}m" + else: + idle_txt = f"{idle_s:.0f}s" + notes = [] + if not s["reachable"]: + notes.append(R + "unreachable" + OFF) + if s.get("wake_recovery_pending"): + notes.append(Y + "reload recovery pending" + OFF) + if s["wake_in_progress"]: + notes.append("waking...") + if s["last_wake_latency_s"] is not None: + notes.append(f"last wake {s['last_wake_latency_s']}s") + if s["last_error"]: + notes.append(s["last_error"]) + print(f" {key:<7} {s['model']:<22} {colour}{state:<11}{OFF} " + f"{s['active_requests']:>6} {idle_txt:>8} {' '.join(notes)}") +PY +} + +# -------------------------------------------------------------------- up ---- +cmd_up() { + local target="${1:-}" body key + if [ -z "$target" ]; then + say "All services are always running -- the router wakes them on demand." + say "Nothing to do. Use '${C_CYAN}./vllmctl up MODEL${C_OFF}' to force-wake one now." + cmd_status + return 0 + fi + admin_alive || { die_admin_down; return 1; } + # Level-2 wakes can take a minute or two (weights come back from NFS). + admin_call POST "/admin/wake/$(uri_escape "$target")" 400 + local body="$RESPONSE" + case "$HTTP_CODE" in + 200) + local lat + lat="$(json_field "$body" "d.get('latency_s')")" + [ -n "$lat" ] && lat=" in ${lat}s" + ok "Awake: $(json_field "$body" "d['model']") ($target)${lat:-}" + return 0 + ;; + 404) err "Unknown model or service: '$target'"; return 1 ;; + 503) + warn "'$target' is still waking ($(json_field "$body" "d['error']['sleep_depth']"))." + say " Retry-After: $(json_field "$body" "d['retry_after_s']")s est. $(json_field "$body" "d['estimated_wake_seconds']")s" + [ -n "$(json_field "$body" "d['error'].get('message','')")" ] && \ + say " $(json_field "$body" "d['error']['message']")" + return 1 + ;; + *) err "Router returned HTTP $HTTP_CODE"; printf '%s\n' "$body"; return 1 ;; + esac +} + +uri_escape() { python3 -c "import urllib.parse,sys;print(urllib.parse.quote(sys.argv[1],safe=''))" "$1"; } + +# ------------------------------------------------------------------ down ---- +cmd_down() { + local level="$1"; shift + local targets="${1:-}" body rc=0 key + admin_alive || { die_admin_down; return 1; } + if [ -z "$targets" ]; then targets="$KEYS"; say "Offloading all services (sleep level 2)..."; fi + for key in $targets; do + admin_call POST "/admin/sleep/$(uri_escape "$key")?level=$level" 120 + local body="$RESPONSE" + case "$HTTP_CODE" in + 200) ok "$(json_field "$body" "d['model']") ($key): $(sleep_state "$level")" ;; + 409) + warn "$key: refused -- $(json_field "$body" "d.get('reason','busy')")" + rc=1 ;; + 404) err "Unknown model or service: '$key'"; rc=1 ;; + *) err "$key: HTTP $HTTP_CODE $(printf '%s' "$body" | head -c 200)"; rc=1 ;; + esac + done + return $rc +} + +sleep_state() { [ "$1" = "2" ] && printf '%s' 'offloaded (weights freed)' || printf '%s' 'sleeping (weights in RAM)'; } + +# ------------------------------------------------------------------ list ---- +cmd_list() { + say "${C_CYAN}Available models in $MODEL_ROOT${C_OFF}" + local served d name found=0 + served="$(curl -fs -m 3 "http://127.0.0.1:$PUBLIC_PORT/v1/models" 2>/dev/null)" + for d in "$MODEL_ROOT"/*/; do + [ -d "$d" ] || continue + name="$(basename "$d")" + if compgen -G "$d/*.safetensors" >/dev/null || [ -f "$d/config.json" ]; then + found=1 + local mark="${C_DIM}on disk${C_OFF}" + [ -n "$served" ] && printf '%s' "$served" | grep -q "\"$name\"" && mark="${C_GREEN}*served${C_OFF}" + local size + size="$(du -sh "$d" 2>/dev/null | cut -f1)" + printf ' %-40s %6s %s\n' "$name" "$size" "$mark" + fi + done + [ "$found" = 1 ] || warn ' (none found — pull one with: ./vllmctl pull )' +} + +# ------------------------------------------------------------------ logs ---- +cmd_logs() { + local follow='' n=100 svc='' + for a in "$@"; do + case "$a" in + -f|--follow) follow='--follow' ;; + *[!0-9]*) svc="$a" ;; + *) n="$a" ;; + esac + done + local args=($follow --tail "$n") + if [ -n "$svc" ]; then + local service + service="$(docker_service_for "$svc")" || { err "Unknown service '$svc' (text|ocr|embed|router)"; return 1; } + args+=("$service") + fi + dcompose logs "${args[@]}" +} + +# --------------------------------------------------------------- restart ---- +cmd_restart() { + local svc="${1:-}" service + [ -n "$svc" ] || { err 'Usage: ./vllmctl restart SERVICE (text|ocr|embed|router)'; return 1; } + service="$(docker_service_for "$svc")" || { err "Unknown service '$svc'"; return 1; } + warn "Restarting $service -- cold start can take 2-10 min (NFS weights)." + dcompose restart "$service" +} + +# ------------------------------------------------------------------ pull ---- +cmd_pull() { + local repo="${1:-}" name="${2:-}" + [ -n "$repo" ] || { err 'Usage: ./vllmctl pull [local-name]'; return 1; } + [ -n "$name" ] || name="$(basename "$repo")" + if [ -e "$MODEL_ROOT/$name" ]; then + err "$MODEL_ROOT/$name already exists" + return 1 + fi + say "Downloading '$repo' → $MODEL_ROOT/$name (Ctrl-C safe to abort)" + local image="vllm/vllm-openai:$(env_get VLLM_VERSION v0.27.1)" + docker run --rm -i \ + --name vllm-pull \ + -v "$MODEL_ROOT:/models" \ + --env-file "$ROOT/.env" \ + --entrypoint python3 \ + "$image" \ + -c "from huggingface_hub import snapshot_download; p=snapshot_download('$repo', local_dir='/models/$name'); print('done:', p)" + local rc=$? + [ $rc -eq 0 ] && ok "Pulled '$name'." || err "Pull failed (rc=$rc)" + return $rc +} + +# ------------------------------------------------------------------ main ---- +case "${1:-help}" in + status) shift; cmd_status "$@" ;; + up) shift; cmd_up "$@" ;; + down|offload) shift; cmd_down 2 "$@" ;; + sleep|nap) shift; cmd_down 1 "$@" ;; + stop) shift; cmd_down 2 "$@" ;; + list) shift; cmd_list "$@" ;; + logs) shift; cmd_logs "$@" ;; + restart) shift; cmd_restart "$@" ;; + pull) shift; cmd_pull "$@" ;; + help|--help|-h|*) + sed -n '2,26p' "$0" | sed 's/^# \{0,1\}//' + ;; +esac