[Fast Start] Wait for weight cache daemon readiness (#58370)

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.

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})