A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 03:45 UTC
Make engine core client handshake timeout configurable (#27444)
Signed-off-by: Seiji Eicher <seiji@anyscale.com>
Seiji Eicher committed
Dec 19, 2025 at 12:38 UTC
1ab521353155e170adf9d9bae34e2a96b994aef4
3 files changed
+105
-3
tests/v1/engine/test_engine_core_client.py
+92
-1
@@ -2,12 +2,14 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import asyncio
5
+import importlib
6
import os
7
import signal
8
import time
9
import uuid
10
from dataclasses import dataclass
11
from threading import Thread
12
+from types import SimpleNamespace
13
from typing import Any
14
from unittest.mock import MagicMock
15
@@ -24,7 +26,11 @@ from vllm.usage.usage_lib import UsageContext
26
from vllm.utils.torch_utils import set_default_torch_num_threads
27
from vllm.v1.engine import EngineCoreRequest
28
from vllm.v1.engine.core import EngineCore
27
-from vllm.v1.engine.core_client import AsyncMPClient, EngineCoreClient, SyncMPClient
29
+from vllm.v1.engine.core_client import (
30
+ AsyncMPClient,
31
+ EngineCoreClient,
32
+ SyncMPClient,
33
+)
34
from vllm.v1.engine.utils import CoreEngineProcManager
35
from vllm.v1.executor.abstract import Executor
36
@@ -60,6 +66,91 @@ def make_request(
66
)
67
68
69
+def _reload_envs_module():
70
+ import vllm.envs as envs_mod
71
+
72
+ cache_clear = getattr(getattr(envs_mod, "__getattr__", None), "cache_clear", None)
73
+ if cache_clear is not None:
74
+ cache_clear()
75
+ return importlib.reload(envs_mod)
76
+
77
+
78
+def _reload_core_client_module():
79
+ module = importlib.import_module("vllm.v1.engine.core_client")
80
+ return importlib.reload(module)
81
+
82
+
83
+def test_mp_client_uses_env_timeout(monkeypatch: pytest.MonkeyPatch):
84
+ timeout_value = 654
85
+ monkeypatch.setenv("VLLM_ENGINE_READY_TIMEOUT_S", str(timeout_value))
86
+
87
+ # Ensure that the environment variable is loaded if caching is enabled
88
+ _reload_envs_module()
89
+ core_client_mod = _reload_core_client_module()
90
+
91
+ poll_timeouts: list[int] = []
92
+
93
+ class ShadowSocket:
94
+ def poll(self, timeout: int) -> int:
95
+ # Capture the timeout value for each poll call
96
+ poll_timeouts.append(timeout)
97
+ return 1
98
+
99
+ def recv_multipart(self):
100
+ return (b"\x00\x00", b"ready")
101
+
102
+ class DummySocket:
103
+ def send_multipart(self, _msg, *, copy: bool = False, track: bool = False):
104
+ if track:
105
+ return SimpleNamespace(done=True)
106
+
107
+ def recv_multipart(self, *, copy: bool = False):
108
+ return (b"", b"")
109
+
110
+ def close(self, *, linger: int = 0):
111
+ pass
112
+
113
+ def bind(self, _address):
114
+ pass
115
+
116
+ def connect(self, _address):
117
+ pass
118
+
119
+ def setsockopt(self, *_args, **_kwargs):
120
+ pass
121
+
122
+ monkeypatch.setattr(core_client_mod.zmq.Socket, "shadow", lambda *_: ShadowSocket())
123
+ monkeypatch.setattr(
124
+ core_client_mod, "make_zmq_socket", lambda *_, **__: DummySocket()
125
+ )
126
+
127
+ parallel_config = SimpleNamespace(
128
+ data_parallel_size=1,
129
+ data_parallel_rank=0,
130
+ data_parallel_size_local=1,
131
+ data_parallel_rank_local=None,
132
+ data_parallel_hybrid_lb=False,
133
+ data_parallel_external_lb=False,
134
+ )
135
+ vllm_config = SimpleNamespace(parallel_config=parallel_config)
136
+
137
+ client = core_client_mod.MPClient(
138
+ asyncio_mode=False,
139
+ vllm_config=vllm_config,
140
+ executor_class=object,
141
+ log_stats=False,
142
+ client_addresses={
143
+ "input_address": "inproc://input",
144
+ "output_address": "inproc://output",
145
+ },
146
+ )
147
+ try:
148
+ # timeout_value is in seconds, but poll receives milliseconds
149
+ assert poll_timeouts == [timeout_value * 1000]
150
+ finally:
151
+ client.shutdown()
152
+
153
+
154
def loop_until_done(client: EngineCoreClient, outputs: dict):
155
while True:
156
engine_core_outputs = client.get_output().outputs
vllm/envs.py
+6
@@ -24,6 +24,7 @@ if TYPE_CHECKING:
24
LOCAL_RANK: int = 0
25
CUDA_VISIBLE_DEVICES: str | None = None
26
VLLM_ENGINE_ITERATION_TIMEOUT_S: int = 60
27
+ VLLM_ENGINE_READY_TIMEOUT_S: int = 600
28
VLLM_API_KEY: str | None = None
29
VLLM_DEBUG_LOG_API_SERVER_RESPONSE: bool = False
30
S3_ACCESS_KEY_ID: str | None = None
@@ -604,6 +605,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
605
"VLLM_ENGINE_ITERATION_TIMEOUT_S": lambda: int(
606
os.environ.get("VLLM_ENGINE_ITERATION_TIMEOUT_S", "60")
607
),
608
+ # Timeout in seconds for waiting for engine cores to become ready
609
+ # during startup. Default is 600 seconds (10 minutes).
610
+ "VLLM_ENGINE_READY_TIMEOUT_S": lambda: int(
611
+ os.environ.get("VLLM_ENGINE_READY_TIMEOUT_S", "600")
612
+ ),
613
# API key for vLLM API server
614
"VLLM_API_KEY": lambda: os.environ.get("VLLM_API_KEY", None),
615
# Whether to log responses from API Server for debugging
vllm/v1/engine/core_client.py
+7
-2
@@ -20,6 +20,7 @@ import zmq
20
import zmq.asyncio
21
22
from vllm.config import VllmConfig
23
+from vllm.envs import VLLM_ENGINE_READY_TIMEOUT_S
24
from vllm.logger import init_logger
25
from vllm.lora.request import LoRARequest
26
from vllm.tasks import SupportedTask
@@ -528,7 +529,9 @@ class MPClient(EngineCoreClient):
529
identities = set(self.core_engines)
530
sync_input_socket = zmq.Socket.shadow(self.input_socket)
531
while identities:
531
- if not sync_input_socket.poll(timeout=600_000):
532
+ if not sync_input_socket.poll(
533
+ timeout=VLLM_ENGINE_READY_TIMEOUT_S * 1000 # convert to ms
534
+ ):
535
raise TimeoutError(
536
"Timed out waiting for engines to send"
537
"initial message on input socket."
@@ -1340,7 +1343,9 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
1343
# Wait for ready messages from new engines on the input socket
1344
sync_input_socket = zmq.Socket.shadow(self.input_socket)
1345
while new_engine_identities:
1343
- if not sync_input_socket.poll(timeout=600_000):
1346
+ if not sync_input_socket.poll(
1347
+ timeout=VLLM_ENGINE_READY_TIMEOUT_S * 1000 # convert to ms
1348
+ ):
1349
raise TimeoutError(
1350
"Timed out waiting for new engines to send initial "
1351
"message on input socket."