[Fast Start] Add /health endpoint for the weight cache daemon (#58552)

Signed-off-by: Xun Sun <UNIDY2002@outlook.com>
Co-authored-by: Isotr0py <Isotr0py@outlook.com>

tests/model_executor/model_loader/test_weight_cache.py

@@ -12,6 +12,8 @@ import subprocess
 import sys
 import tempfile
 import threading
+import urllib.error
+import urllib.request
 from dataclasses import dataclass, field
 from typing import Any
 
@@ -20,6 +22,7 @@ import pytest
 from vllm import SamplingParams
 from vllm.assets.image import ImageAsset
 from vllm.platforms import current_platform
+from vllm.utils.network_utils import get_open_port
 
 
 class WeightCacheDaemon:
@@ -33,6 +36,7 @@ class WeightCacheDaemon:
     ):
         # Short base path: Unix socket paths are limited to ~107 characters.
         self.socket_dir = tempfile.mkdtemp(prefix="vllm_ipc_")
+        self.health_port = get_open_port()
         self._cmd = [
             sys.executable,
             "-m",
@@ -45,6 +49,10 @@ class WeightCacheDaemon:
             "--weight-cache-socket-dir",
             self.socket_dir,
             "--enforce-eager",
+            "--weight-cache-health-host",
+            "127.0.0.1",
+            "--weight-cache-health-port",
+            str(self.health_port),
             *(extra_args or []),
         ]
         self._proc: subprocess.Popen | None = None
@@ -63,6 +71,16 @@ class WeightCacheDaemon:
         assert self._proc is not None and self._proc.stderr is not None
         self._proc.stderr.read()
 
+    def check_health(self) -> None:
+        url = f"http://127.0.0.1:{self.health_port}/health"
+        try:
+            with urllib.request.urlopen(url, timeout=5) as response:
+                assert response.status == 200
+        except urllib.error.HTTPError as error:
+            raise AssertionError(
+                f"Weight cache daemon is not ready: HTTP {error.code}"
+            ) from error
+
     def _stop(self) -> None:
         assert self._proc is not None
         self._proc.terminate()
@@ -200,6 +218,7 @@ def test_ipc_cache_cold_start_and_warm_restart(vllm_runner, case: ModelCase):
         extra_args=case.daemon_args,
     ) as d:
         warm_outputs = generate(vllm_runner, case, d.socket_dir, fallback=False)
+        d.check_health()
         # Warm restart: a second engine lifetime against the same daemon.
         restart_outputs = generate(vllm_runner, case, d.socket_dir, fallback=False)
 

vllm/entrypoints/cli/preload.py

@@ -23,8 +23,14 @@ import multiprocessing
 import queue
 import signal
 import sys
+import threading
 import typing
+from collections.abc import Sequence
 from itertools import product
+from multiprocessing.process import BaseProcess
+
+from fastapi import FastAPI
+from fastapi.responses import Response
 
 from vllm.engine.arg_utils import EngineArgs
 from vllm.entrypoints.cli.types import CLISubcommand
@@ -44,6 +50,55 @@ else:
 logger = init_logger(__name__)
 
 
+class _HealthState:
+    def __init__(self, processes: Sequence[BaseProcess]) -> None:
+        self.processes = processes
+        self._ready = False
+        self._lock = threading.Lock()
+
+    def mark_ready(self) -> None:
+        with self._lock:
+            self._ready = True
+
+    @property
+    def is_ready(self) -> bool:
+        with self._lock:
+            ready = self._ready
+        return ready and not any(proc.exitcode is not None for proc in self.processes)
+
+
+def _build_health_app(state: _HealthState) -> FastAPI:
+    app = FastAPI(openapi_url=None, docs_url=None, redoc_url=None)
+
+    @app.get("/health", include_in_schema=False)
+    async def health() -> Response:
+        return Response(status_code=200 if state.is_ready else 503)
+
+    return app
+
+
+def _start_health_server(
+    host: str, port: int, state: _HealthState
+) -> tuple[typing.Any, threading.Thread]:
+    import uvicorn
+
+    config = uvicorn.Config(
+        _build_health_app(state),
+        host=host,
+        port=port,
+        log_level="error",
+        access_log=False,
+    )
+    server = uvicorn.Server(config)
+    thread = threading.Thread(
+        target=server.run,
+        name="vllm-weight-cache-health",
+        daemon=True,
+    )
+    thread.start()
+    return server, thread
+
+
 class PreloadSubcommand(CLISubcommand):
     """The `preload` subcommand for the vLLM CLI."""
 
@@ -75,6 +130,18 @@ class PreloadSubcommand(CLISubcommand):
             "--weight-cache-master-port + 1 for multi-node, or a free port "
             "for single-node.",
         )
+        parser.add_argument(
+            "--weight-cache-health-port",
+            type=int,
+            default=None,
+            help="Port for the parent /health endpoint; disabled by default.",
+        )
+        parser.add_argument(
+            "--weight-cache-health-host",
+            type=str,
+            default="0.0.0.0",
+            help="Host for the parent /health endpoint (default: 0.0.0.0).",
+        )
 
     @staticmethod
     def cmd(args: argparse.Namespace) -> None:
@@ -190,10 +257,21 @@ class PreloadSubcommand(CLISubcommand):
                 tp_rank,
             ) in product(groups, placements)
         ]
+        health_state = _HealthState(procs)
+        health_server = None
+        health_thread = None
+        if args.weight_cache_health_port is not None:
+            health_server, health_thread = _start_health_server(
+                args.weight_cache_health_host,
+                args.weight_cache_health_port,
+                health_state,
+            )
         for proc in procs:
             proc.start()
 
         def _shutdown(signum, frame):
+            if health_server is not None:
+                health_server.should_exit = True
             for proc in procs:
                 proc.terminate()
 
@@ -216,7 +294,12 @@ class PreloadSubcommand(CLISubcommand):
                         proc.terminate()
                     for proc in procs:
                         proc.join()
+                    if health_server is not None:
+                        health_server.should_exit = True
+                    if health_thread is not None:
+                        health_thread.join()
                     sys.exit(max((p.exitcode or 0) for p in procs))
+        health_state.mark_ready()
         socket_dir_msg = args.weight_cache_socket_dir or "the default socket dir"
         logger.info_once(
             "===== Weight cache daemon READY: node %d/%d serving %d local "
@@ -230,6 +313,10 @@ class PreloadSubcommand(CLISubcommand):
 
         for proc in procs:
             proc.join()
+        if health_server is not None:
+            health_server.should_exit = True
+        if health_thread is not None:
+            health_thread.join()
         sys.exit(max(proc.exitcode or 0 for proc in procs))
 
     def subparser_init(