diff --git a/tests/v1/engine/test_core_engine_actor_manager.py b/tests/v1/engine/test_core_engine_actor_manager.py index f1a67ad9ba..768f51f437 100644 --- a/tests/v1/engine/test_core_engine_actor_manager.py +++ b/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}], diff --git a/tests/v1/engine/test_dp_world_group_port.py b/tests/v1/engine/test_dp_world_group_port.py new file mode 100644 index 0000000000..49f5a59ab5 --- /dev/null +++ b/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() diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index de6680e027..feb8f25d1e 100644 --- a/vllm/distributed/parallel_state.py +++ b/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 ) diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index 203d1b407c..49e6884ac9 100644 --- a/vllm/v1/engine/utils.py +++ b/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", ) }