[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>
Xun Sun committed
Sep 28, 2026 at 01:48 UTC
e7900156e130c9880eb03b7c1f2df32820e7a2be
2 files changed
+106
tests/model_executor/model_loader/test_weight_cache.py
+19
@@ -12,6 +12,8 @@ import subprocess
12
import sys
13
import tempfile
14
import threading
15
+import urllib.error
16
+import urllib.request
17
from dataclasses import dataclass, field
18
from typing import Any
19
@@ -20,6 +22,7 @@ import pytest
22
from vllm import SamplingParams
23
from vllm.assets.image import ImageAsset
24
from vllm.platforms import current_platform
25
+from vllm.utils.network_utils import get_open_port
26
27
28
class WeightCacheDaemon:
@@ -33,6 +36,7 @@ class WeightCacheDaemon:
36
):
37
# Short base path: Unix socket paths are limited to ~107 characters.
38
self.socket_dir = tempfile.mkdtemp(prefix="vllm_ipc_")
39
+ self.health_port = get_open_port()
40
self._cmd = [
41
sys.executable,
42
"-m",
@@ -45,6 +49,10 @@ class WeightCacheDaemon:
49
"--weight-cache-socket-dir",
50
self.socket_dir,
51
"--enforce-eager",
52
+ "--weight-cache-health-host",
53
+ "127.0.0.1",
54
+ "--weight-cache-health-port",
55
+ str(self.health_port),
56
*(extra_args or []),
57
]
58
self._proc: subprocess.Popen | None = None
@@ -63,6 +71,16 @@ class WeightCacheDaemon:
71
assert self._proc is not None and self._proc.stderr is not None
72
self._proc.stderr.read()
73
74
+ def check_health(self) -> None:
75
+ url = f"http://127.0.0.1:{self.health_port}/health"
76
+ try:
77
+ with urllib.request.urlopen(url, timeout=5) as response:
78
+ assert response.status == 200
79
+ except urllib.error.HTTPError as error:
80
+ raise AssertionError(
81
+ f"Weight cache daemon is not ready: HTTP {error.code}"
82
+ ) from error
83
+
84
def _stop(self) -> None:
85
assert self._proc is not None
86
self._proc.terminate()
@@ -200,6 +218,7 @@ def test_ipc_cache_cold_start_and_warm_restart(vllm_runner, case: ModelCase):
218
extra_args=case.daemon_args,
219
) as d:
220
warm_outputs = generate(vllm_runner, case, d.socket_dir, fallback=False)
221
+ d.check_health()
222
# Warm restart: a second engine lifetime against the same daemon.
223
restart_outputs = generate(vllm_runner, case, d.socket_dir, fallback=False)
224
vllm/entrypoints/cli/preload.py
+87
@@ -23,8 +23,14 @@ import multiprocessing
23
import queue
24
import signal
25
import sys
26
+import threading
27
import typing
28
+from collections.abc import Sequence
29
from itertools import product
30
+from multiprocessing.process import BaseProcess
31
+
32
+from fastapi import FastAPI
33
+from fastapi.responses import Response
34
35
from vllm.engine.arg_utils import EngineArgs
36
from vllm.entrypoints.cli.types import CLISubcommand
@@ -44,6 +50,55 @@ else:
50
logger = init_logger(__name__)
51
52
53
+class _HealthState:
54
+ def __init__(self, processes: Sequence[BaseProcess]) -> None:
55
+ self.processes = processes
56
+ self._ready = False
57
+ self._lock = threading.Lock()
58
+
59
+ def mark_ready(self) -> None:
60
+ with self._lock:
61
+ self._ready = True
62
+
63
+ @property
64
+ def is_ready(self) -> bool:
65
+ with self._lock:
66
+ ready = self._ready
67
+ return ready and not any(proc.exitcode is not None for proc in self.processes)
68
+
69
+
70
+def _build_health_app(state: _HealthState) -> FastAPI:
71
+ app = FastAPI(openapi_url=None, docs_url=None, redoc_url=None)
72
+
73
+ @app.get("/health", include_in_schema=False)
74
+ async def health() -> Response:
75
+ return Response(status_code=200 if state.is_ready else 503)
76
+
77
+ return app
78
+
79
+
80
+def _start_health_server(
81
+ host: str, port: int, state: _HealthState
82
+) -> tuple[typing.Any, threading.Thread]:
83
+ import uvicorn
84
+
85
+ config = uvicorn.Config(
86
+ _build_health_app(state),
87
+ host=host,
88
+ port=port,
89
+ log_level="error",
90
+ access_log=False,
91
+ )
92
+ server = uvicorn.Server(config)
93
+ thread = threading.Thread(
94
+ target=server.run,
95
+ name="vllm-weight-cache-health",
96
+ daemon=True,
97
+ )
98
+ thread.start()
99
+ return server, thread
100
+
101
+
102
class PreloadSubcommand(CLISubcommand):
103
"""The `preload` subcommand for the vLLM CLI."""
104
@@ -75,6 +130,18 @@ class PreloadSubcommand(CLISubcommand):
130
"--weight-cache-master-port + 1 for multi-node, or a free port "
131
"for single-node.",
132
)
133
+ parser.add_argument(
134
+ "--weight-cache-health-port",
135
+ type=int,
136
+ default=None,
137
+ help="Port for the parent /health endpoint; disabled by default.",
138
+ )
139
+ parser.add_argument(
140
+ "--weight-cache-health-host",
141
+ type=str,
142
+ default="0.0.0.0",
143
+ help="Host for the parent /health endpoint (default: 0.0.0.0).",
144
+ )
145
146
@staticmethod
147
def cmd(args: argparse.Namespace) -> None:
@@ -190,10 +257,21 @@ class PreloadSubcommand(CLISubcommand):
257
tp_rank,
258
) in product(groups, placements)
259
]
260
+ health_state = _HealthState(procs)
261
+ health_server = None
262
+ health_thread = None
263
+ if args.weight_cache_health_port is not None:
264
+ health_server, health_thread = _start_health_server(
265
+ args.weight_cache_health_host,
266
+ args.weight_cache_health_port,
267
+ health_state,
268
+ )
269
for proc in procs:
270
proc.start()
271
272
def _shutdown(signum, frame):
273
+ if health_server is not None:
274
+ health_server.should_exit = True
275
for proc in procs:
276
proc.terminate()
277
@@ -216,7 +294,12 @@ class PreloadSubcommand(CLISubcommand):
294
proc.terminate()
295
for proc in procs:
296
proc.join()
297
+ if health_server is not None:
298
+ health_server.should_exit = True
299
+ if health_thread is not None:
300
+ health_thread.join()
301
sys.exit(max((p.exitcode or 0) for p in procs))
302
+ health_state.mark_ready()
303
socket_dir_msg = args.weight_cache_socket_dir or "the default socket dir"
304
logger.info_once(
305
"===== Weight cache daemon READY: node %d/%d serving %d local "
@@ -230,6 +313,10 @@ class PreloadSubcommand(CLISubcommand):
313
314
for proc in procs:
315
proc.join()
316
+ if health_server is not None:
317
+ health_server.should_exit = True
318
+ if health_thread is not None:
319
+ health_thread.join()
320
sys.exit(max(proc.exitcode or 0 for proc in procs))
321
322
def subparser_init(