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 6, 09:30 UTC
Fix DP coordinator ZMQ port TOCTOU (#37452)
Signed-off-by: Itay Alroy <ialroy@nvidia.com>
Itay Alroy committed
Mar 20, 2026 at 02:58 UTC
ca1ac1a4b44f46c53747f6792507ec1927ade617
2 files changed
+58
-8
vllm/utils/network_utils.py
+1
-1
@@ -247,7 +247,7 @@ def split_zmq_path(path: str) -> tuple[str, str, str]:
247
248
scheme = parsed.scheme
249
host = parsed.hostname or ""
250
- port = str(parsed.port or "")
250
+ port = "" if parsed.port is None else str(parsed.port)
251
if host.startswith("[") and host.endswith("]"):
252
host = host[1:-1] # Remove brackets for IPv6 address
253
vllm/v1/engine/coordinator.py
+57
-7
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import copy
4
import multiprocessing
5
+import multiprocessing.connection
6
import time
7
import weakref
8
@@ -10,7 +11,7 @@ import zmq
11
12
from vllm.config import ParallelConfig
13
from vllm.logger import init_logger
13
-from vllm.utils.network_utils import make_zmq_socket
14
+from vllm.utils.network_utils import get_tcp_uri, make_zmq_socket
15
from vllm.utils.system_utils import get_mp_context, set_process_title
16
from vllm.v1.engine import EngineCoreOutputs, EngineCoreRequestType
17
from vllm.v1.serial_utils import MsgpackDecoder
@@ -55,6 +56,25 @@ class DPCoordinator:
56
request wave / running state changes.
57
"""
58
59
+ def _wait_for_zmq_addrs(self, zmq_addr_pipe) -> tuple[str, str, str]:
60
+ try:
61
+ ready = multiprocessing.connection.wait(
62
+ [zmq_addr_pipe, self.proc.sentinel], timeout=30
63
+ )
64
+ if not ready:
65
+ raise RuntimeError(
66
+ "DP Coordinator process failed to report ZMQ addresses "
67
+ "during startup."
68
+ )
69
+ try:
70
+ return zmq_addr_pipe.recv()
71
+ except EOFError:
72
+ raise RuntimeError(
73
+ "DP Coordinator process failed during startup."
74
+ ) from None
75
+ finally:
76
+ zmq_addr_pipe.close()
77
+
78
def __init__(
79
self, parallel_config: ParallelConfig, enable_wave_coordination: bool = True
80
):
@@ -66,18 +86,24 @@ class DPCoordinator:
86
# Assume coordinator is colocated with front-end procs when not in
87
# either external or hybrid DP LB mode.
88
local_only = not parallel_config.local_engines_only
69
- front_publish_address = get_engine_client_zmq_addr(
70
- local_only=local_only, host=host
71
- )
72
-
89
local_only_eng = dp_size == parallel_config.data_parallel_size_local
90
# NOTE(yongji): handling scaling from intra-node to inter-node
91
if parallel_config.enable_elastic_ep:
92
local_only_eng = False
77
- back_publish_address = get_engine_client_zmq_addr(local_only_eng, host)
78
- back_output_address = get_engine_client_zmq_addr(local_only_eng, host)
93
+
94
+ def bind_address(local_only: bool) -> str:
95
+ return (
96
+ get_engine_client_zmq_addr(local_only=True, host=host)
97
+ if local_only
98
+ else get_tcp_uri(host, 0)
99
+ )
100
+
101
+ front_publish_address = bind_address(local_only)
102
+ back_publish_address = bind_address(local_only_eng)
103
+ back_output_address = bind_address(local_only_eng)
104
105
context = get_mp_context()
106
+ parent_zmq_addr_pipe, child_zmq_addr_pipe = context.Pipe(duplex=False)
107
self.proc: multiprocessing.Process = context.Process(
108
target=DPCoordinatorProc.run_coordinator,
109
name="VLLM_DP_Coordinator",
@@ -86,11 +112,18 @@ class DPCoordinator:
112
"front_publish_address": front_publish_address,
113
"back_output_address": back_output_address,
114
"back_publish_address": back_publish_address,
115
+ "zmq_addr_pipe": child_zmq_addr_pipe,
116
"enable_wave_coordination": enable_wave_coordination,
117
},
118
daemon=True,
119
)
120
self.proc.start()
121
+ child_zmq_addr_pipe.close()
122
+ (
123
+ front_publish_address,
124
+ back_output_address,
125
+ back_publish_address,
126
+ ) = self._wait_for_zmq_addrs(parent_zmq_addr_pipe)
127
128
self.stats_publish_address = front_publish_address
129
self.coord_in_address = back_publish_address
@@ -136,6 +169,7 @@ class DPCoordinatorProc:
169
front_publish_address: str,
170
back_output_address: str,
171
back_publish_address: str,
172
+ zmq_addr_pipe=None,
173
min_stats_update_interval_ms: int = 100,
174
enable_wave_coordination: bool = True,
175
):
@@ -149,15 +183,20 @@ class DPCoordinatorProc:
183
front_publish_address,
184
back_output_address,
185
back_publish_address,
186
+ zmq_addr_pipe,
187
)
188
except KeyboardInterrupt:
189
logger.info("DP Coordinator process exiting")
190
+ finally:
191
+ if zmq_addr_pipe is not None:
192
+ zmq_addr_pipe.close()
193
194
def process_input_socket(
195
self,
196
front_publish_address: str,
197
back_output_address: str,
198
back_publish_address: str,
199
+ zmq_addr_pipe=None,
200
):
201
decoder = MsgpackDecoder(EngineCoreOutputs)
202
@@ -191,6 +230,17 @@ class DPCoordinatorProc:
230
bind=True,
231
) as publish_back,
232
):
233
+ if zmq_addr_pipe is not None:
234
+ try:
235
+ zmq_addr_pipe.send(
236
+ (
237
+ publish_front.getsockopt(zmq.LAST_ENDPOINT).decode(),
238
+ output_back.getsockopt(zmq.LAST_ENDPOINT).decode(),
239
+ publish_back.getsockopt(zmq.LAST_ENDPOINT).decode(),
240
+ )
241
+ )
242
+ finally:
243
+ zmq_addr_pipe.close()
244
# Wait until all engines subscribe.
245
for _ in self.engines:
246
if publish_back.recv() != b"\x01":