[Fast Start] Wait for weight cache daemon readiness (#58370)
Xun Sun committed
Sep 25, 2026 at 03:18 UTC
e773247f8aee272d0262928117e3049cd2902b5e
2 files changed
+32
-47
tests/model_executor/model_loader/test_weight_cache.py
+1
-47
@@ -7,14 +7,11 @@ warm restarts (weights mapped from the daemon via CUDA IPC) must both serve
7
identical outputs.
8
"""
9
10
-import glob
11
-import os
10
import shutil
11
import subprocess
12
import sys
13
import tempfile
14
import threading
17
-import time
15
from dataclasses import dataclass, field
16
from typing import Any
17
@@ -24,8 +21,6 @@ from vllm import SamplingParams
21
from vllm.assets.image import ImageAsset
22
from vllm.platforms import current_platform
23
27
-DAEMON_TIMEOUT_S = 600
28
-
24
25
class WeightCacheDaemon:
26
"""Context manager running the real weight cache daemon as a subprocess."""
@@ -35,14 +30,9 @@ class WeightCacheDaemon:
30
model: str,
31
tp_size: int,
32
extra_args: list[str] | None = None,
38
- num_groups: int = 1,
33
):
34
# Short base path: Unix socket paths are limited to ~107 characters.
35
self.socket_dir = tempfile.mkdtemp(prefix="vllm_ipc_")
42
- self.tp_size = tp_size
43
- # Each daemon group (target plus cached drafts) binds one socket per
44
- # local rank.
45
- self.num_sockets = tp_size * num_groups
36
self._cmd = [
37
sys.executable,
38
"-m",
@@ -57,18 +47,12 @@ class WeightCacheDaemon:
47
*(extra_args or []),
48
]
49
self._proc: subprocess.Popen | None = None
60
- self._lines: list[str] = []
50
51
def __enter__(self) -> "WeightCacheDaemon":
52
self._proc = subprocess.Popen(
53
self._cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True
54
)
55
threading.Thread(target=self._drain_stderr, daemon=True).start()
67
- try:
68
- self._wait_ready(DAEMON_TIMEOUT_S)
69
- except Exception:
70
- self._stop()
71
- raise
56
return self
57
58
def __exit__(self, *exc_info) -> None:
@@ -76,32 +60,7 @@ class WeightCacheDaemon:
60
61
def _drain_stderr(self) -> None:
62
assert self._proc is not None and self._proc.stderr is not None
79
- for line in self._proc.stderr:
80
- self._lines.append(line)
81
-
82
- def _logs(self) -> str:
83
- return "".join(self._lines)
84
-
85
- def _wait_ready(self, timeout_s: float) -> None:
86
- assert self._proc is not None
87
- # Each rank binds its socket only once the model is fully cached.
88
- # Poll for the socket files rather than a log line: model loading can
89
- # pull in JIT compilers that swap the process's stderr and swallow
90
- # everything logged afterwards, making log-based readiness flaky.
91
- pattern = os.path.join(self.socket_dir, "vllm_weight_cache_*.sock")
92
- deadline = time.monotonic() + timeout_s
93
- while time.monotonic() < deadline:
94
- if self._proc.poll() is not None:
95
- raise RuntimeError(
96
- f"Weight cache daemon exited with {self._proc.returncode}:\n"
97
- f"{self._logs()}"
98
- )
99
- if len(glob.glob(pattern)) >= self.num_sockets:
100
- return
101
- time.sleep(1.0)
102
- raise TimeoutError(
103
- f"Weight cache daemon not ready after {timeout_s}s:\n{self._logs()}"
104
- )
63
+ self._proc.stderr.read()
64
65
def _stop(self) -> None:
66
assert self._proc is not None
@@ -121,9 +80,6 @@ class ModelCase:
80
images: list | None = None
81
llm_kwargs: dict[str, Any] = field(default_factory=dict)
82
daemon_args: list[str] = field(default_factory=list)
124
- # Daemon groups the launcher starts; 2 when a cached draft group joins the
125
- # target group (MTP/EAGLE speculative decoding).
126
- daemon_groups: int = 1
83
84
85
def generate(
@@ -203,7 +159,6 @@ QWEN_MTP_CASE = ModelCase(
159
"--speculative-config",
160
'{"method": "mtp", "num_speculative_tokens": 1}',
161
],
206
- daemon_groups=2,
162
)
163
164
@@ -242,7 +197,6 @@ def test_ipc_cache_cold_start_and_warm_restart(vllm_runner, case: ModelCase):
197
case.model,
198
tp_size=1,
199
extra_args=case.daemon_args,
245
- num_groups=case.daemon_groups,
200
) as d:
201
warm_outputs = generate(vllm_runner, case, d.socket_dir, fallback=False)
202
# Warm restart: a second engine lifetime against the same daemon.
vllm/model_executor/model_loader/weight_cache/ipc_loader.py
+31
@@ -5,6 +5,7 @@ daemon via CUDA IPC instead of loading from disk."""
5
6
import dataclasses
7
import socket
8
+import time
9
from copy import copy
10
11
import torch
@@ -46,6 +47,7 @@ logger = init_logger(__name__)
47
48
_CONNECT_TIMEOUT_S = 5.0
49
_STATE_TIMEOUT_S = 300.0
50
+_STARTUP_RETRY_INTERVAL_S = 0.5
51
52
53
class IpcModelLoader(BaseModelLoader):
@@ -295,8 +297,37 @@ class IpcModelLoader(BaseModelLoader):
297
dp_rank=dp_group.rank_in_group,
298
is_draft=self.is_draft,
299
)
300
+ if not self.fallback:
301
+ return self._request_state_with_startup_wait(cache_config)
302
return self._request_state(cache_config)
303
304
+ def _request_state_with_startup_wait(
305
+ self, cache_config: WeightCacheKey
306
+ ) -> WeightCacheState:
307
+ deadline = time.monotonic() + self.state_timeout_s
308
+ waiting_logged = False
309
+ while True:
310
+ try:
311
+ return self._request_state(cache_config)
312
+ except WeightCacheUnavailableError as e:
313
+ if time.monotonic() >= deadline:
314
+ raise WeightCacheUnavailableError(
315
+ "Weight cache daemon did not become ready within "
316
+ f"{self.state_timeout_s:.1f}s"
317
+ ) from e
318
+ if not waiting_logged:
319
+ logger.info(
320
+ "Waiting up to %.1fs for the weight cache daemon to start",
321
+ self.state_timeout_s,
322
+ )
323
+ waiting_logged = True
324
+ time.sleep(
325
+ max(
326
+ 0.0,
327
+ min(_STARTUP_RETRY_INTERVAL_S, deadline - time.monotonic()),
328
+ )
329
+ )
330
+
331
def _request_state(self, cache_config: WeightCacheKey) -> WeightCacheState:
332
with self._connect(self.state_timeout_s) as conn:
333
send_msg(conn, {"cmd": "get_state", "cache_config": cache_config})