[BugFix] Pick the DP world-group port at bind time via the coordination store (#51018)

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: Codex <noreply@openai.com>

tests/v1/engine/test_core_engine_actor_manager.py

@@ -13,8 +13,9 @@ from unittest.mock import Mock
 import pytest
 import ray
 import zmq
+from torch.distributed import TCPStore
 
-from vllm.utils.network_utils import make_zmq_socket, split_zmq_path
+from vllm.utils.network_utils import get_open_port, make_zmq_socket, split_zmq_path
 from vllm.v1.engine.core import EngineCoreActorMixin
 from vllm.v1.engine.core_client import BackgroundResources
 from vllm.v1.engine.utils import (
@@ -134,6 +135,72 @@ def _make_addresses() -> EngineZmqAddresses:
     )
 
 
+@pytest.mark.parametrize(
+    ("overrides", "expect_store"),
+    [
+        ({"data_parallel_backend": "mp"}, True),
+        ({}, True),
+        (
+            {
+                "data_parallel_backend": "mp",
+                "data_parallel_rank": 1,
+                "local_engines_only": True,
+            },
+            False,
+        ),
+        ({"data_parallel_backend": "mp", "data_parallel_rank_local": 0}, False),
+        ({"data_parallel_backend": "mp", "enable_elastic_ep": True}, False),
+        ({"data_parallel_backend": "mp", "data_parallel_size": 1}, False),
+    ],
+    ids=["mp", "ray", "non-master-node", "offline", "elastic-ep", "single-engine"],
+)
+def test_coordination_store_held_only_by_the_online_dp_master(
+    monkeypatch: pytest.MonkeyPatch, overrides: dict[str, Any], expect_store: bool
+) -> None:
+    """The DP master node holds a coordination store for both engine backends,
+    so engines there pick world-group ports at bind time. Every other
+    deployment keeps the pre-allocated ports and must not hold one."""
+
+    class FakeManager:
+        def __init__(self, **kwargs) -> None:
+            self.kwargs = kwargs
+
+    vllm_config = _make_vllm_config_ray_dp_multinode()
+    parallel_config = vllm_config.parallel_config
+    for name, value in overrides.items():
+        setattr(parallel_config, name, value)
+    # A fresh port per case so a lingering handshake socket cannot fail the next.
+    parallel_config.data_parallel_rpc_port = get_open_port()
+    parallel_config._coord_store_port = 0
+    monkeypatch.setattr("vllm.v1.engine.utils.CoreEngineProcManager", FakeManager)
+    monkeypatch.setattr("vllm.v1.engine.utils.CoreEngineActorManager", FakeManager)
+    monkeypatch.setattr(
+        "vllm.v1.engine.utils.wait_for_engine_startup", lambda *args, **kwargs: None
+    )
+
+    with launch_core_engines(
+        vllm_config,
+        executor_class=_DummyExecutor,
+        log_stats=False,
+        addresses=_make_addresses(),
+    ) as engine_launch:
+        assert engine_launch.engine_manager is not None
+        if not expect_store:
+            assert not parallel_config._coord_store_port
+            return
+        assert parallel_config._coord_store_port
+        # Engines look the store up while they start, so it must be reachable
+        # for as long as this frame is alive.
+        client = TCPStore(
+            parallel_config.data_parallel_master_ip,
+            parallel_config._coord_store_port,
+            is_master=False,
+            wait_for_workers=False,
+        )
+        client.set("probe", b"1")
+        assert client.get("probe") == b"1"
+
+
 def _make_cpu_placement_group():
     pg = ray.util.placement_group(
         [{"CPU": 0.001}, {"CPU": 1.0}],

tests/v1/engine/test_dp_world_group_port.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The DP world group picks its port when it binds, not ahead of time."""
+
+import multiprocessing as mp
+import socket
+import traceback
+from multiprocessing.queues import Queue
+
+import pytest
+import torch
+
+from vllm.distributed.utils import create_tcp_store
+from vllm.utils.network_utils import get_open_port
+
+RESULT_TIMEOUT_S = 240
+
+
+def _join_dp_world_group(
+    dp_rank: int, coord_store_port: int, dp_master_port: int, result: Queue
+) -> None:
+    """Bring up one DP rank of the world group and report what it saw."""
+    try:
+        import torch.distributed as dist
+
+        from vllm.config import VllmConfig, set_current_vllm_config
+        from vllm.config.parallel import ParallelConfig
+        from vllm.distributed.parallel_state import init_distributed_environment
+
+        parallel_config = ParallelConfig(
+            data_parallel_size=2, data_parallel_rank=dp_rank
+        )
+        parallel_config.data_parallel_master_ip = "127.0.0.1"
+        parallel_config._data_parallel_master_port_list = [dp_master_port]
+        parallel_config._coord_store_port = coord_store_port
+
+        with set_current_vllm_config(VllmConfig(parallel_config=parallel_config)):
+            init_distributed_environment(
+                world_size=1, rank=0, local_rank=0, backend="gloo"
+            )
+
+        reduced = torch.tensor([dp_rank + 1])
+        dist.all_reduce(reduced)
+        result.put((dp_rank, int(reduced.item()), None))
+    except Exception:
+        result.put((dp_rank, None, traceback.format_exc()))
+
+
+@pytest.mark.skipif(
+    not torch.distributed.is_gloo_available(), reason="gloo backend required"
+)
+@pytest.mark.parametrize("with_coord_store", [True, False])
+def test_dp_world_group_survives_a_taken_preallocated_port(with_coord_store: bool):
+    """With a coordination store, the two ranks agree on a port the group binds
+    itself, so a pre-allocated port taken in the meantime cannot stop startup.
+    Without one, the pre-allocated port is still used as before.
+    """
+    squatter = None
+    coord_store = None
+    try:
+        if with_coord_store:
+            # Something else grabbed the pre-allocated port after it was picked.
+            squatter = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            squatter.bind(("127.0.0.1", 0))
+            squatter.listen(1)
+            dp_master_port = squatter.getsockname()[1]
+            coord_store = create_tcp_store(
+                "127.0.0.1", 0, is_master=True, world_size=-1, wait_for_workers=False
+            )
+            coord_store_port = coord_store.port
+        else:
+            dp_master_port = get_open_port()
+            coord_store_port = 0
+
+        ctx = mp.get_context("spawn")
+        result: Queue = ctx.Queue()
+        procs = [
+            ctx.Process(
+                target=_join_dp_world_group,
+                args=(dp_rank, coord_store_port, dp_master_port, result),
+            )
+            for dp_rank in range(2)
+        ]
+        for proc in procs:
+            proc.start()
+        reports: list[tuple[int, int | None, str | None]] = []
+        try:
+            while len(reports) < len(procs):
+                reports.append(result.get(timeout=RESULT_TIMEOUT_S))
+                if reports[-1][2] is not None:
+                    break  # the peer would wait on this rank; do not wait for it
+        finally:
+            for proc in procs:
+                proc.join(timeout=30)
+                if proc.is_alive():
+                    proc.kill()
+
+        for dp_rank, reduced, error in sorted(reports):
+            assert error is None, f"DP rank {dp_rank} failed:\n{error}"
+            # Both ranks joined the same group: 1 + 2 over the two ranks.
+            assert reduced == 3, f"DP rank {dp_rank} all-reduced to {reduced}"
+
+        if with_coord_store:
+            assert coord_store is not None
+            bound_port = int(coord_store.get("world_pg_port").decode())
+            assert bound_port != dp_master_port
+    finally:
+        if squatter is not None:
+            squatter.close()

vllm/distributed/parallel_state.py

@@ -42,6 +42,7 @@ import torch.distributed
 import torch.distributed._functional_collectives as funcol
 import torch.distributed._symmetric_memory
 from torch.distributed import Backend, ProcessGroup, Store
+from torch.distributed.constants import default_pg_timeout
 
 import vllm.envs as envs
 from vllm.distributed.device_communicators.base_device_communicator import (
@@ -49,6 +50,7 @@ from vllm.distributed.device_communicators.base_device_communicator import (
 )
 from vllm.distributed.utils import (
     StatelessProcessGroup,
+    create_tcp_store,
     get_cached_tcp_store_client,
 )
 from vllm.logger import init_logger
@@ -1829,6 +1831,7 @@ def init_distributed_environment(
 
     config = get_current_vllm_config_or_none()
     enable_elastic_ep = config is not None and config.parallel_config.enable_elastic_ep
+    world_pg_store: Store | None = None
     if (
         config is not None
         and config.parallel_config.distributed_executor_backend != "external_launcher"
@@ -1852,7 +1855,46 @@ def init_distributed_environment(
             distributed_init_method = get_distributed_init_method(ip, port)
         else:
             ip = parallel_config.data_parallel_master_ip
-            port = parallel_config.get_next_dp_init_port()
+            if (
+                parallel_config._coord_store_port
+                and not torch.distributed.is_initialized()
+            ):
+                # Group rank 0 lets the world-group TCPStore bind port 0
+                # and publishes the kernel-assigned port via the
+                # coordination store. A one-shot client of our own: the
+                # cached one is inherited across fork and shared with
+                # sibling workers.
+                coord_store = create_tcp_store(
+                    ip,
+                    parallel_config._coord_store_port,
+                    is_master=False,
+                    wait_for_workers=False,
+                )
+                if rank == 0:
+                    # Workers can only connect once the port is published.
+                    world_pg_store = create_tcp_store(
+                        ip,
+                        0,
+                        world_size=world_size,
+                        is_master=True,
+                        wait_for_workers=False,
+                        timeout=timeout or default_pg_timeout,
+                        multi_tenant=True,
+                    )
+                    port = world_pg_store.port
+                    coord_store.set("world_pg_port", str(port).encode())
+                else:
+                    port = int(coord_store.get("world_pg_port").decode())
+                    world_pg_store = create_tcp_store(
+                        ip,
+                        port,
+                        world_size=world_size,
+                        is_master=False,
+                        timeout=timeout or default_pg_timeout,
+                        multi_tenant=True,
+                    )
+            else:
+                port = parallel_config.get_next_dp_init_port()
             distributed_init_method = get_distributed_init_method(ip, port)
             logger.debug(
                 "Adjusting world_size=%d rank=%d distributed_init_method=%s for DP",
@@ -1882,8 +1924,8 @@ def init_distributed_environment(
                 "Fallback Gloo backend is not available."
             )
             backend = "gloo"
-        store = None
-        if distributed_init_method.startswith("file://"):
+        store = world_pg_store
+        if store is None and distributed_init_method.startswith("file://"):
             store = torch.distributed.FileStore(
                 distributed_init_method.removeprefix("file://"), world_size
             )

vllm/v1/engine/utils.py

@@ -1195,6 +1195,23 @@ def launch_core_engines(
     else:
         coordinator = None
 
+    # Hold a coordination TCPStore (alive for this frame) so engines pick DP
+    # master ports at bind time; pre-allocated ports can be taken before use.
+    coord_store = None
+    if (
+        dp_size > 1
+        and not offline_mode
+        and dp_rank == 0
+        and local_engine_count > 0
+        and not parallel_config.enable_elastic_ep
+    ):
+        from vllm.distributed.utils import create_tcp_store
+
+        coord_store = create_tcp_store(
+            host, 0, is_master=True, world_size=-1, wait_for_workers=False
+        )
+        parallel_config._coord_store_port = coord_store.port
+
     if parallel_config.data_parallel_backend == "ray":
         logger.info("Starting ray-based data parallel backend")
 
@@ -1415,6 +1432,7 @@ def wait_for_engine_startup(
                             "data_parallel_master_ip",
                             "data_parallel_master_port",
                             "_data_parallel_master_port_list",
+                            "_coord_store_port",
                             "data_parallel_size",
                         )
                     }