diff --git a/tests/model_executor/model_loader/test_weight_cache.py b/tests/model_executor/model_loader/test_weight_cache.py index 7b60aaf04b..2ea77e79c8 100644 --- a/tests/model_executor/model_loader/test_weight_cache.py +++ b/tests/model_executor/model_loader/test_weight_cache.py @@ -7,14 +7,11 @@ warm restarts (weights mapped from the daemon via CUDA IPC) must both serve identical outputs. """ -import glob -import os import shutil import subprocess import sys import tempfile import threading -import time from dataclasses import dataclass, field from typing import Any @@ -24,8 +21,6 @@ from vllm import SamplingParams from vllm.assets.image import ImageAsset from vllm.platforms import current_platform -DAEMON_TIMEOUT_S = 600 - class WeightCacheDaemon: """Context manager running the real weight cache daemon as a subprocess.""" @@ -35,14 +30,9 @@ class WeightCacheDaemon: model: str, tp_size: int, extra_args: list[str] | None = None, - num_groups: int = 1, ): # Short base path: Unix socket paths are limited to ~107 characters. self.socket_dir = tempfile.mkdtemp(prefix="vllm_ipc_") - self.tp_size = tp_size - # Each daemon group (target plus cached drafts) binds one socket per - # local rank. - self.num_sockets = tp_size * num_groups self._cmd = [ sys.executable, "-m", @@ -57,18 +47,12 @@ class WeightCacheDaemon: *(extra_args or []), ] self._proc: subprocess.Popen | None = None - self._lines: list[str] = [] def __enter__(self) -> "WeightCacheDaemon": self._proc = subprocess.Popen( self._cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True ) threading.Thread(target=self._drain_stderr, daemon=True).start() - try: - self._wait_ready(DAEMON_TIMEOUT_S) - except Exception: - self._stop() - raise return self def __exit__(self, *exc_info) -> None: @@ -76,32 +60,7 @@ class WeightCacheDaemon: def _drain_stderr(self) -> None: assert self._proc is not None and self._proc.stderr is not None - for line in self._proc.stderr: - self._lines.append(line) - - def _logs(self) -> str: - return "".join(self._lines) - - def _wait_ready(self, timeout_s: float) -> None: - assert self._proc is not None - # Each rank binds its socket only once the model is fully cached. - # Poll for the socket files rather than a log line: model loading can - # pull in JIT compilers that swap the process's stderr and swallow - # everything logged afterwards, making log-based readiness flaky. - pattern = os.path.join(self.socket_dir, "vllm_weight_cache_*.sock") - deadline = time.monotonic() + timeout_s - while time.monotonic() < deadline: - if self._proc.poll() is not None: - raise RuntimeError( - f"Weight cache daemon exited with {self._proc.returncode}:\n" - f"{self._logs()}" - ) - if len(glob.glob(pattern)) >= self.num_sockets: - return - time.sleep(1.0) - raise TimeoutError( - f"Weight cache daemon not ready after {timeout_s}s:\n{self._logs()}" - ) + self._proc.stderr.read() def _stop(self) -> None: assert self._proc is not None @@ -121,9 +80,6 @@ class ModelCase: images: list | None = None llm_kwargs: dict[str, Any] = field(default_factory=dict) daemon_args: list[str] = field(default_factory=list) - # Daemon groups the launcher starts; 2 when a cached draft group joins the - # target group (MTP/EAGLE speculative decoding). - daemon_groups: int = 1 def generate( @@ -203,7 +159,6 @@ QWEN_MTP_CASE = ModelCase( "--speculative-config", '{"method": "mtp", "num_speculative_tokens": 1}', ], - daemon_groups=2, ) @@ -242,7 +197,6 @@ def test_ipc_cache_cold_start_and_warm_restart(vllm_runner, case: ModelCase): case.model, tp_size=1, extra_args=case.daemon_args, - num_groups=case.daemon_groups, ) as d: warm_outputs = generate(vllm_runner, case, d.socket_dir, fallback=False) # Warm restart: a second engine lifetime against the same daemon. diff --git a/vllm/model_executor/model_loader/weight_cache/ipc_loader.py b/vllm/model_executor/model_loader/weight_cache/ipc_loader.py index 8e37ce3ad6..508e95c83d 100644 --- a/vllm/model_executor/model_loader/weight_cache/ipc_loader.py +++ b/vllm/model_executor/model_loader/weight_cache/ipc_loader.py @@ -5,6 +5,7 @@ daemon via CUDA IPC instead of loading from disk.""" import dataclasses import socket +import time from copy import copy import torch @@ -46,6 +47,7 @@ logger = init_logger(__name__) _CONNECT_TIMEOUT_S = 5.0 _STATE_TIMEOUT_S = 300.0 +_STARTUP_RETRY_INTERVAL_S = 0.5 class IpcModelLoader(BaseModelLoader): @@ -295,8 +297,37 @@ class IpcModelLoader(BaseModelLoader): dp_rank=dp_group.rank_in_group, is_draft=self.is_draft, ) + if not self.fallback: + return self._request_state_with_startup_wait(cache_config) return self._request_state(cache_config) + def _request_state_with_startup_wait( + self, cache_config: WeightCacheKey + ) -> WeightCacheState: + deadline = time.monotonic() + self.state_timeout_s + waiting_logged = False + while True: + try: + return self._request_state(cache_config) + except WeightCacheUnavailableError as e: + if time.monotonic() >= deadline: + raise WeightCacheUnavailableError( + "Weight cache daemon did not become ready within " + f"{self.state_timeout_s:.1f}s" + ) from e + if not waiting_logged: + logger.info( + "Waiting up to %.1fs for the weight cache daemon to start", + self.state_timeout_s, + ) + waiting_logged = True + time.sleep( + max( + 0.0, + min(_STARTUP_RETRY_INTERVAL_S, deadline - time.monotonic()), + ) + ) + def _request_state(self, cache_config: WeightCacheKey) -> WeightCacheState: with self._connect(self.state_timeout_s) as conn: send_msg(conn, {"cmd": "get_state", "cache_config": cache_config})