diff --git a/tests/model_executor/model_loader/test_weight_cache.py b/tests/model_executor/model_loader/test_weight_cache.py index 2280fff8b7..e4ac86e0eb 100644 --- a/tests/model_executor/model_loader/test_weight_cache.py +++ b/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) diff --git a/vllm/entrypoints/cli/preload.py b/vllm/entrypoints/cli/preload.py index d312b6c36a..50734a052b 100644 --- a/vllm/entrypoints/cli/preload.py +++ b/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(