[BugFix] Bind RayExecutorV2 TCPStore before publishing its port (#50969)
Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com> Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
aoshen02 committed
Aug 29, 2026 at 03:27 UTC
67e86d1e6fbcab31a404a620521a5179a4c5fe67
2 files changed
+57
-105
tests/distributed/test_ray_v2_executor.py
+26
-58
@@ -7,12 +7,17 @@ Validates executor initialization, placement group support, RPC calls,
7
and distributed execution with various TP/PP configurations.
8
"""
9
10
+import errno
11
import gc
12
+import socket
13
import threading
14
+from types import SimpleNamespace
15
from unittest.mock import patch
16
+from urllib.parse import urlsplit
17
18
import pytest
19
import ray
20
+from torch.distributed import TCPStore
21
22
from vllm import LLM
23
from vllm.config import VllmConfig
@@ -102,67 +107,30 @@ def assert_executor(executor, tp_size, pp_size):
107
assert handle.node_id is not None
108
109
105
-def test_select_tcpstore_port_seeds_disjoint_windows(monkeypatch):
106
- """Co-located DP engines scan distinct, adjacent port windows, so two
107
- engines on a node cannot pick the same TCPStore port."""
108
- requested = []
110
+def test_dist_init_store_owns_port(monkeypatch):
111
+ """The rank-0 actor keeps the selected TCPStore port bound."""
112
+ worker = ray_executor_v2.RayWorkerProc.__new__(ray_executor_v2.RayWorkerProc)
113
+ worker._parallel_config = SimpleNamespace(world_size=1)
114
+ monkeypatch.setattr(ray.util, "get_node_ip_address", lambda: "127.0.0.1")
115
110
- def fake_get_open_port(start_port, max_attempts):
111
- requested.append((start_port, max_attempts))
112
- return start_port
116
+ init_method = worker.create_dist_init_method()
117
+ port = urlsplit(init_method).port
118
+ assert port is not None and port != 0
119
114
- monkeypatch.setattr(ray_executor_v2, "_get_open_port", fake_get_open_port)
120
+ with socket.socket() as contender, pytest.raises(OSError) as exc_info:
121
+ contender.bind(("127.0.0.1", port))
122
+ assert exc_info.value.errno == errno.EADDRINUSE
123
116
- ports = [
117
- RayExecutorV2._select_tcpstore_port(rank, master_port=29500)
118
- for rank in range(4)
119
- ]
120
-
121
- assert requested == [(29600, 32), (29632, 32), (29664, 32), (29696, 32)]
122
- assert len(set(ports)) == 4
123
-
124
-
125
-def test_select_tcpstore_port_non_dp_uses_random(monkeypatch):
126
- """A non-DP engine has no local rank and uses a random port."""
127
- monkeypatch.setattr(ray_executor_v2, "get_open_port", lambda: 54321)
128
- assert RayExecutorV2._select_tcpstore_port(None, master_port=29500) == 54321
129
-
130
-
131
-def test_colocated_tcpstore_ports_differ():
132
- """Two or more replicas on one host can collide on the same TCPStore port.
133
- Their ports must differ and never be a low system port.
134
- """
135
- ports = []
136
- real_select = RayExecutorV2._select_tcpstore_port
137
-
138
- def record(local_dp_rank, master_port):
139
- port = real_select(local_dp_rank, master_port)
140
- ports.append(port)
141
- return port
142
-
143
- with patch.object(RayExecutorV2, "_select_tcpstore_port", staticmethod(record)):
144
- first = RayExecutorV2(vllm_config=create_vllm_config())
145
- try:
146
- second = RayExecutorV2(vllm_config=create_vllm_config())
147
- try:
148
- assert len(ports) == 2
149
- assert min(ports) >= 1024
150
- assert ports[0] != ports[1]
151
- finally:
152
- second.shutdown()
153
- finally:
154
- first.shutdown()
155
-
156
-
157
-def test_select_tcpstore_port_full_window_uses_random(monkeypatch):
158
- """A fully occupied window falls back to a random port."""
159
-
160
- def raise_full(start_port, max_attempts):
161
- raise RuntimeError("no open port")
162
-
163
- monkeypatch.setattr(ray_executor_v2, "_get_open_port", raise_full)
164
- monkeypatch.setattr(ray_executor_v2, "get_open_port", lambda: 54321)
165
- assert RayExecutorV2._select_tcpstore_port(0, master_port=29500) == 54321
124
+ reused_store = TCPStore(
125
+ "127.0.0.1",
126
+ port,
127
+ world_size=1,
128
+ is_master=True,
129
+ wait_for_workers=False,
130
+ multi_tenant=True,
131
+ )
132
+ worker._dist_init_store.set("key", "value")
133
+ assert reused_store.get("key") == b"value"
134
135
136
@pytest.mark.parametrize("tp_size, pp_size", [(1, 1), (2, 1), (4, 1), (2, 2)])
vllm/v1/executor/ray_executor_v2.py
+31
-47
@@ -8,6 +8,8 @@ from collections import defaultdict, deque
8
from dataclasses import dataclass
9
from typing import Any
10
11
+from torch.distributed import TCPStore
12
+
13
import vllm.envs as envs
14
from vllm.config import VllmConfig
15
from vllm.distributed.device_communicators.shm_broadcast import (
@@ -17,9 +19,7 @@ from vllm.distributed.device_communicators.shm_broadcast import (
19
from vllm.logger import init_logger
20
from vllm.platforms import current_platform
21
from vllm.utils.network_utils import (
20
- _get_open_port,
22
get_distributed_init_method,
22
- get_open_port,
23
)
24
from vllm.v1.executor.multiproc_executor import (
25
FutureWrapper,
@@ -106,22 +106,35 @@ class RayWorkerProc(WorkerProc):
106
self,
107
vllm_config: VllmConfig,
108
rank: int,
109
- distributed_init_method: str,
109
input_shm_handle: Handle,
110
is_driver_worker: bool,
111
is_driver_node: bool = False,
112
):
113
# Defer WorkerProc.__init__ until GPU IDs are known.
114
self._is_driver_node = is_driver_node
115
+ self._parallel_config = vllm_config.parallel_config
116
self._init_kwargs = dict(
117
vllm_config=vllm_config,
118
rank=rank,
119
- distributed_init_method=distributed_init_method,
119
input_shm_handle=input_shm_handle,
120
shared_worker_lock=None,
121
is_driver_worker=is_driver_worker,
122
)
123
124
+ def create_dist_init_method(self) -> str:
125
+ host = ray.util.get_node_ip_address()
126
+ store = TCPStore(
127
+ host_name=host,
128
+ port=0,
129
+ world_size=self._parallel_config.world_size,
130
+ is_master=True,
131
+ wait_for_workers=False,
132
+ multi_tenant=True,
133
+ )
134
+ # Keep the bound server alive until init_process_group reuses it.
135
+ self._dist_init_store = store
136
+ return get_distributed_init_method(host, store.port)
137
+
138
def get_node_and_physical_gpu_ids(self) -> tuple[str, list[int]]:
139
"""Return (node_id, physical_gpu_ids) assigned to this actor by Ray."""
140
node_id = ray.get_runtime_context().get_node_id()
@@ -140,6 +153,7 @@ class RayWorkerProc(WorkerProc):
153
self,
154
local_rank: int,
155
env_vars: dict[str, str],
156
+ distributed_init_method: str,
157
driver_env_vars: dict[str, str] | None = None,
158
assigned_physical_gpu_ids: list[int] | None = None,
159
) -> None:
@@ -163,6 +177,7 @@ class RayWorkerProc(WorkerProc):
177
assigned_physical_gpu_ids
178
)
179
180
+ self._init_kwargs["distributed_init_method"] = distributed_init_method
181
self.local_rank = local_rank
182
super().__init__(
183
local_rank=local_rank,
@@ -260,25 +275,6 @@ class RayExecutorV2(MultiprocExecutor):
275
return {"num_gpus": num_devices}
276
return {"num_gpus": 0, "resources": {device_key: num_devices}}
277
263
- @staticmethod
264
- def _select_tcpstore_port(local_dp_rank: int | None, master_port: int) -> int:
265
- """Pick the torch.distributed TCPStore port for this engine.
266
-
267
- Co-located DP engines choosing this port with a shared random search
268
- collide intermittently. Seeding by node-local DP rank gives each a
269
- disjoint window. Non-DP engines and full windows fall back to a
270
- random port.
271
- """
272
- if local_dp_rank is None:
273
- return get_open_port()
274
- # Offset past the DP master port reserved range, one window per rank.
275
- window = 32
276
- start_port = master_port + 100 + local_dp_rank * window
277
- try:
278
- return _get_open_port(start_port=start_port, max_attempts=window)
279
- except RuntimeError:
280
- return get_open_port()
281
-
278
def _init_executor(self) -> None:
279
"""Initialize the RayExecutorV2 executor."""
280
self._finalizer = weakref.finalize(self, self.shutdown)
@@ -324,22 +320,7 @@ class RayExecutorV2(MultiprocExecutor):
320
}
321
)
322
327
- # Step 3: Resolve the IP for torch.distributed TCPStore.
328
- # The TCPStore server runs on rank 0's node, so all workers
329
- # must be able to reach this address.
330
- dist_ip = bundle_assignments[0]["node_ip"]
331
- parallel_config = self.vllm_config.parallel_config
332
- local_dp_rank = (
333
- parallel_config.data_parallel_rank_local
334
- if parallel_config.data_parallel_size > 1
335
- else None
336
- )
337
- port = self._select_tcpstore_port(
338
- local_dp_rank, parallel_config.data_parallel_master_port
339
- )
340
- distributed_init_method = get_distributed_init_method(dist_ip, port)
341
-
342
- # Step 4: Create broadcast MessageQueue.
323
+ # Step 3: Create broadcast MessageQueue.
324
# Workers on the driver node use shared memory; the rest use TCP.
325
max_chunk_bytes = envs.VLLM_MQ_MAX_CHUNK_BYTES_MB * 1024 * 1024
326
n_local = sum(1 for a in bundle_assignments if a["node_id"] == driver_node)
@@ -351,9 +332,9 @@ class RayExecutorV2(MultiprocExecutor):
332
)
333
scheduler_output_handle = self.rpc_broadcast_mq.export_handle()
334
354
- # Step 5: Spawn RayWorkerProc actors into PG bundles (deferred init).
335
+ # Step 4: Spawn RayWorkerProc actors into PG bundles (deferred init).
336
# Workers are created lightweight here; full initialization happens
356
- # in Step 7 after GPU IDs are discovered.
337
+ # in Step 6 after GPU IDs are discovered.
338
self.ray_worker_handles: list[RayWorkerHandle] = []
339
instance_id = self.vllm_config.instance_id
340
@@ -405,7 +386,6 @@ class RayExecutorV2(MultiprocExecutor):
386
.remote(
387
vllm_config=self.vllm_config,
388
rank=bundle["rank"],
408
- distributed_init_method=distributed_init_method,
389
input_shm_handle=scheduler_output_handle,
390
is_driver_worker=is_driver_worker,
391
is_driver_node=is_driver_node,
@@ -421,7 +401,7 @@ class RayExecutorV2(MultiprocExecutor):
401
)
402
self.ray_worker_handles.append(handle)
403
424
- # Step 6: Discover physical GPU IDs assigned to each worker via Ray
404
+ # Step 5: Discover physical GPU IDs assigned to each worker via Ray
405
# runtime context.
406
worker_node_and_physical_gpu_ids = ray.get(
407
[
@@ -440,8 +420,11 @@ class RayExecutorV2(MultiprocExecutor):
420
for node_id, physical_gpu_ids in node_physical_gpu_ids.items():
421
node_physical_gpu_ids[node_id] = sorted(physical_gpu_ids)
422
443
- # Step 7: Initialize workers with local logical ranks and the
423
+ # Step 6: Initialize workers with local logical ranks and the
424
# logical-to-physical GPU mapping discovered from Ray placement.
425
+ distributed_init_method = ray.get(
426
+ self.ray_worker_handles[0].actor.create_dist_init_method.remote()
427
+ )
428
init_worker_refs = []
429
for i, (node_id, _) in enumerate(worker_node_and_physical_gpu_ids):
430
local_rank = node_workers[node_id].index(i)
@@ -452,6 +435,7 @@ class RayExecutorV2(MultiprocExecutor):
435
self.ray_worker_handles[i].actor.initialize_worker.remote(
436
local_rank,
437
worker_env_vars,
438
+ distributed_init_method,
439
self.driver_env_vars,
440
assigned_physical_gpu_ids=assigned_physical_gpu_ids,
441
)
@@ -465,7 +449,7 @@ class RayExecutorV2(MultiprocExecutor):
449
)
450
ray.get(init_worker_refs)
451
468
- # Step 8: Collect response MQ handles
452
+ # Step 7: Collect response MQ handles
453
init_results = ray.get(
454
[h.actor.wait_for_init.remote() for h in self.ray_worker_handles]
455
)
@@ -478,12 +462,12 @@ class RayExecutorV2(MultiprocExecutor):
462
MessageQueue.create_from_handle(result["handle"], 0)
463
)
464
481
- # Step 9: Start run() before wait_until_ready() to avoid
465
+ # Step 8: Start run() before wait_until_ready() to avoid
466
# deadlock — workers send subscriptions inside run().
467
for handle in self.ray_worker_handles:
468
handle.run()
469
486
- # Step 10: wait_until_ready() barrier
470
+ # Step 9: wait_until_ready() barrier
471
self.rpc_broadcast_mq.wait_until_ready()
472
for response_mq in self.response_mqs:
473
response_mq.wait_until_ready()