vllm-frontdoor/router/admin_api.py

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