From 7483114a5de09b7823e5303cca03d98817e35760 Mon Sep 17 00:00:00 2001 From: inf-yasong Date: Sat, 11 Jul 2026 01:42:34 +0000 Subject: [PATCH 1/2] feat(vllm): add optional per-node DP launch --- docs/config-reference.md | 30 ++++++ src/srtctl/backends/vllm.py | 113 ++++++++++++++++++++- src/srtctl/core/topology.py | 14 +++ tests/test_configs.py | 190 ++++++++++++++++++++++++++++++++++++ 4 files changed, 343 insertions(+), 4 deletions(-) diff --git a/docs/config-reference.md b/docs/config-reference.md index 9717a58f..2e298931 100644 --- a/docs/config-reference.md +++ b/docs/config-reference.md @@ -422,6 +422,36 @@ Each worker leader gets a globally unique port starting at 5550: | decode_0 | 5552 | | decode_1 | 5553 | +### vLLM DP launch mode + +vLLM data-parallel endpoints use one process per GPU by default. Set +`dp_launch_mode: per_node` to launch one process per node and let vLLM +manage the local DP ranks in a shared CUDA namespace: + +```yaml +backend: + type: vllm + dp_launch_mode: per_node + vllm_config: + prefill: + data-parallel-size: 8 + data-parallel-hybrid-lb: true + decode: + data-parallel-size: 16 + data-parallel-hybrid-lb: true +``` + +| Value | Process layout | +| ---------- | --------------------------------------------------- | +| `per_gpu` | One process per DP rank/GPU (default) | +| `per_node` | One process manages all DP ranks allocated per node | + +In `per_node` mode, srtslurm derives `--data-parallel-size-local` and +`--data-parallel-start-rank` from the allocated topology. Do not set those +two flags manually. Set `data-parallel-hybrid-lb: true` when the frontend +must route to an exact DP rank; without hybrid LB, non-leader node processes +are launched with `--headless` for vLLM's internal DP load balancer. + ### TRTLLM Backend When using `type: trtllm`, the backend uses TRTLLM with MPI-style launching: diff --git a/src/srtctl/backends/vllm.py b/src/srtctl/backends/vllm.py index 8b324950..d586d87c 100644 --- a/src/srtctl/backends/vllm.py +++ b/src/srtctl/backends/vllm.py @@ -41,6 +41,7 @@ # Type alias for worker modes WorkerMode = Literal["prefill", "decode", "agg"] +DPLaunchMode = Literal["per_gpu", "per_node"] # Filename for the mooncake-store JSON config srtslurm writes to log_dir at job # start. log_dir is mounted into every worker at /logs, so workers read the JSON @@ -138,6 +139,7 @@ class VLLMProtocol: connector: nixl # translated to --kv-transfer-config JSON allow_prefill_decode_colocation: true # pack P/D on one node when all workers fit allow_prefill_decode_colocation_across_nodes: true # continue packing on later nodes + dp_launch_mode: per_node # one process manages all local DP ranks prefill_environment: PYTHONUNBUFFERED: "1" vllm_config: @@ -191,6 +193,10 @@ class VLLMProtocol: # node pools. Defaults off to preserve the original one-node-only policy. allow_prefill_decode_colocation_across_nodes: bool = False + # DP process layout. Keep the existing per-GPU behavior by default; + # per-node lets vLLM manage local DP ranks in one CUDA namespace. + dp_launch_mode: DPLaunchMode = "per_gpu" + Schema: ClassVar[builtins.type[Schema]] = Schema # ========================================================================= @@ -198,7 +204,7 @@ class VLLMProtocol: # ========================================================================= def get_srun_config(self) -> SrunConfig: - """vLLM uses per-process launching (one srun per node).""" + """vLLM launches one srun step for each generated process.""" from srtctl.backends.base import SrunConfig return SrunConfig(mpi=None, oversubscribe=False, launch_per_endpoint=False) @@ -404,7 +410,7 @@ def _is_dp_mode(self, mode: WorkerMode) -> bool: """Check if this mode uses Data Parallel + Expert Parallel pattern. DP+EP mode is detected when data-parallel-size is set in the mode's config. - In this mode, each GPU runs its own process (rather than TP across GPUs). + ``dp_launch_mode`` controls whether a process owns one rank or all local ranks. """ config = self.get_config_for_mode(mode) return config.get("data-parallel-size") is not None or config.get("data_parallel_size") is not None @@ -431,7 +437,7 @@ def endpoints_to_processes( ) -> list[Process]: """Convert endpoints to processes. - For DP+EP mode (data-parallel-size set), creates one process per GPU. + DP+EP mode uses the configured per-GPU or per-node process layout. For standard TP mode, creates one process per node. """ from srtctl.core.topology import NodePortAllocator, Process, endpoints_to_processes @@ -443,6 +449,13 @@ def endpoints_to_processes( # Standard TP mode: one process per node return endpoints_to_processes(endpoints, base_sys_port=base_sys_port, port_allocator=port_allocator) + if self.dp_launch_mode == "per_node": + return self._dp_per_node_endpoints_to_processes( + endpoints, + base_sys_port=base_sys_port, + port_allocator=port_allocator, + ) + # DP+EP mode: one process per GPU processes: list[Process] = [] current_sys_port = base_sys_port @@ -520,6 +533,67 @@ def endpoints_to_processes( return processes + def _dp_per_node_endpoints_to_processes( + self, + endpoints: list[Endpoint], + base_sys_port: int = DYN_SYSTEM_PORT_BASE, + port_allocator: NodePortAllocator | None = None, + ) -> list[Process]: + """Convert DP endpoints to one process per node.""" + from srtctl.core.topology import NodePortAllocator, Process, endpoints_to_processes + + processes: list[Process] = [] + current_sys_port = base_sys_port + if port_allocator is None: + port_allocator = NodePortAllocator() + + for endpoint in endpoints: + if not self._is_dp_mode(endpoint.mode): + non_dp = endpoints_to_processes( + [endpoint], + base_sys_port=current_sys_port, + port_allocator=port_allocator, + ) + processes.extend(non_dp) + current_sys_port += len(non_dp) + continue + + dp_size = self._get_dp_size(endpoint.mode) or endpoint.total_gpus + if dp_size != endpoint.total_gpus: + raise ValueError( + f"{endpoint.mode} data-parallel-size={dp_size} does not match " + f"the endpoint's {endpoint.total_gpus} allocated GPUs" + ) + + local_dp_size = len(endpoint.gpu_indices) + dp_rpc_port = port_allocator.next_dp_rpc_port(endpoint.leader_node) + nixl_base_port = port_allocator.next_nixl_port_block(dp_size) + dp_start_rank = 0 + + for node in endpoint.nodes: + processes.append( + Process( + node=node, + gpu_indices=endpoint.gpu_indices, + sys_port=current_sys_port, + http_port=port_allocator.next_http_port(node), + endpoint_mode=endpoint.mode, + endpoint_index=endpoint.index, + node_rank=dp_start_rank, + bootstrap_port=( + port_allocator.next_bootstrap_port(node) if endpoint.mode == "prefill" else None + ), + kv_events_port=port_allocator.next_kv_events_port_block(local_dp_size), + nixl_port=nixl_base_port, + dp_rpc_port=dp_rpc_port, + het_group=endpoint.het_group, + ) + ) + current_sys_port += 1 + dp_start_rank += local_dp_size + + return processes + def build_worker_command( self, process: Process, @@ -598,7 +672,38 @@ def build_worker_command( # Check if this is DP+EP mode (data-parallel-size set) is_dp_mode = self._is_dp_mode(mode) - if is_dp_mode: + if is_dp_mode and self.dp_launch_mode == "per_node": + rpc_port_kebab = config.pop("data-parallel-rpc-port", None) + rpc_port_snake = config.pop("data_parallel_rpc_port", None) + config_dp_rpc_port = rpc_port_kebab or rpc_port_snake + dp_rpc_port = process.dp_rpc_port or config_dp_rpc_port or VLLM_DATA_PARALLEL_RPC_PORT + + # These values are derived from the allocated process topology. + config.pop("data-parallel-size-local", None) + config.pop("data_parallel_size_local", None) + config.pop("data-parallel-start-rank", None) + config.pop("data_parallel_start_rank", None) + + cmd.extend( + [ + "--data-parallel-size-local", + str(len(process.gpu_indices)), + "--data-parallel-start-rank", + str(process.node_rank), + "--data-parallel-address", + leader_ip, + "--data-parallel-rpc-port", + str(dp_rpc_port), + ] + ) + + hybrid_lb = config.get("data-parallel-hybrid-lb", config.get("data_parallel_hybrid_lb", False)) + hybrid_lb_enabled = hybrid_lb is True or ( + isinstance(hybrid_lb, str) and hybrid_lb.strip().lower() in {"1", "true", "yes", "on"} + ) + if process.node_rank > 0 and not hybrid_lb_enabled: + cmd.append("--headless") + elif is_dp_mode: # DP+EP mode: each GPU runs its own process # process.node_rank is the dp_rank (set in endpoints_to_processes) dp_rank = process.node_rank diff --git a/src/srtctl/core/topology.py b/src/srtctl/core/topology.py index e662e5f0..b34cfaca 100644 --- a/src/srtctl/core/topology.py +++ b/src/srtctl/core/topology.py @@ -98,6 +98,20 @@ def next_kv_events_port(self) -> int: self._next_kv_events_port += 1 return port + def next_kv_events_port_block(self, size: int) -> int: + """Reserve consecutive KV-event ports and return the base port. + + A vLLM hybrid-DP process opens one publisher per local DP rank, so + adjacent worker processes must receive non-overlapping port ranges. + """ + if size < 1: + raise ValueError("KV-event port block size must be at least 1") + if self._next_kv_events_port == 0: + self._next_kv_events_port = self.base_kv_events_port + port = self._next_kv_events_port + self._next_kv_events_port += size + return port + def next_nixl_port(self) -> int: """Get next available NIXL side channel port (globally unique across all nodes).""" if self._next_nixl_port == 0: diff --git a/tests/test_configs.py b/tests/test_configs.py index a5075a0b..687918d2 100644 --- a/tests/test_configs.py +++ b/tests/test_configs.py @@ -12,6 +12,7 @@ from srtctl.backends import SGLangProtocol, SGLangServerConfig from srtctl.core.schema import SrtConfig from srtctl.ports import ( + KV_EVENTS_PORT_BASE, SGLANG_BOOTSTRAP_PORT_BASE, SGLANG_HTTP_PORT_BASE, SGLANG_HTTP_PORT_STRIDE, @@ -1884,6 +1885,96 @@ def test_dp_mode_creates_per_gpu_processes(self): dp_ranks = [p.node_rank for p in processes] assert dp_ranks == list(range(16)) + def test_dp_per_node_mode_creates_per_node_processes(self): + """Per-node DP owns all local GPUs and reserves rank-sized port blocks.""" + from srtctl.backends import VLLMProtocol, VLLMServerConfig + from srtctl.core.topology import Endpoint + + backend = VLLMProtocol( + dp_launch_mode="per_node", + vllm_config=VLLMServerConfig( + prefill={"data-parallel-size": 8, "enable-expert-parallel": True}, + ), + ) + endpoint = Endpoint( + mode="prefill", + index=0, + nodes=("node0", "node1"), + gpu_indices=frozenset(range(4)), + gpus_per_node=4, + het_group=1, + ) + + processes = backend.endpoints_to_processes([endpoint]) + + assert len(processes) == 2 + assert [p.node for p in processes] == ["node0", "node1"] + assert all(p.gpu_indices == frozenset(range(4)) for p in processes) + assert [p.node_rank for p in processes] == [0, 4] + assert [p.kv_events_port for p in processes] == [KV_EVENTS_PORT_BASE, KV_EVENTS_PORT_BASE + 4] + assert {p.nixl_port for p in processes} == {VLLM_NIXL_PORT_BASE} + assert {p.dp_rpc_port for p in processes} == {VLLM_DATA_PARALLEL_RPC_PORT} + assert {p.het_group for p in processes} == {1} + assert all(p.http_port > 0 for p in processes) + assert all(p.bootstrap_port is not None for p in processes) + + def test_dp_per_node_mode_allocates_non_overlapping_endpoint_ports(self): + """Co-located per-node DP endpoints get disjoint coordination ranges.""" + from srtctl.backends import VLLMProtocol, VLLMServerConfig + from srtctl.core.topology import Endpoint + + backend = VLLMProtocol( + dp_launch_mode="per_node", + vllm_config=VLLMServerConfig( + decode={"data-parallel-size": 4, "enable-expert-parallel": True}, + ), + ) + endpoints = [ + Endpoint( + mode="decode", + index=0, + nodes=("node0",), + gpu_indices=frozenset(range(4)), + gpus_per_node=8, + ), + Endpoint( + mode="decode", + index=1, + nodes=("node0",), + gpu_indices=frozenset(range(4, 8)), + gpus_per_node=8, + ), + ] + + processes = backend.endpoints_to_processes(endpoints) + + assert [p.kv_events_port for p in processes] == [KV_EVENTS_PORT_BASE, KV_EVENTS_PORT_BASE + 4] + assert [p.nixl_port for p in processes] == [VLLM_NIXL_PORT_BASE, VLLM_NIXL_PORT_BASE + 4] + assert [p.dp_rpc_port for p in processes] == [ + VLLM_DATA_PARALLEL_RPC_PORT, + VLLM_DATA_PARALLEL_RPC_PORT + 1, + ] + + def test_dp_per_node_mode_rejects_dp_size_mismatch(self): + """The configured global DP size must match the allocated GPUs.""" + from srtctl.backends import VLLMProtocol, VLLMServerConfig + from srtctl.core.topology import Endpoint + + backend = VLLMProtocol( + dp_launch_mode="per_node", + vllm_config=VLLMServerConfig(prefill={"data-parallel-size": 7}), + ) + endpoint = Endpoint( + mode="prefill", + index=0, + nodes=("node0", "node1"), + gpu_indices=frozenset(range(4)), + gpus_per_node=4, + ) + + with pytest.raises(ValueError, match="data-parallel-size=7"): + backend.endpoints_to_processes([endpoint]) + def test_dp_mode_allocates_unique_ports_for_multiple_endpoints_per_node(self): """Test DP endpoints sharing a node get non-colliding coordination ports.""" from srtctl.backends import VLLMProtocol, VLLMServerConfig @@ -2006,6 +2097,105 @@ def test_dp_mode_command_includes_dp_flags(self): assert "--node-rank" not in cmd assert "--headless" not in cmd + def test_dp_per_node_hybrid_command_targets_local_rank_range(self): + """Hybrid per-node DP exposes the local rank range without headless.""" + from unittest.mock import MagicMock, patch + + from srtctl.backends import VLLMProtocol, VLLMServerConfig + from srtctl.core.topology import Process + + backend = VLLMProtocol( + dp_launch_mode="per_node", + vllm_config=VLLMServerConfig( + decode={ + "data-parallel-size": 8, + "data-parallel-size-local": 99, + "data-parallel-start-rank": 99, + "data-parallel-rpc-port": 13345, + "data-parallel-hybrid-lb": True, + "enable-expert-parallel": True, + }, + ), + ) + leader = Process( + node="node0", + gpu_indices=frozenset(range(4)), + sys_port=8081, + http_port=6100, + endpoint_mode="decode", + endpoint_index=0, + node_rank=0, + dp_rpc_port=VLLM_DATA_PARALLEL_RPC_PORT, + ) + process = Process( + node="node1", + gpu_indices=frozenset(range(4)), + sys_port=8082, + http_port=6100, + endpoint_mode="decode", + endpoint_index=0, + node_rank=4, + dp_rpc_port=VLLM_DATA_PARALLEL_RPC_PORT, + ) + runtime = MagicMock() + runtime.model_path = Path("/model") + runtime.is_hf_model = False + runtime.request_plane = "tcp" + + with patch("srtctl.core.slurm.get_hostname_ip", return_value="10.0.0.1"): + cmd = backend.build_worker_command(process, [leader, process], runtime) + + assert cmd.count("--data-parallel-hybrid-lb") == 1 + assert cmd.count("--data-parallel-size-local") == 1 + assert cmd.count("--data-parallel-start-rank") == 1 + assert cmd[cmd.index("--data-parallel-size-local") + 1] == "4" + assert cmd[cmd.index("--data-parallel-start-rank") + 1] == "4" + assert cmd[cmd.index("--data-parallel-rpc-port") + 1] == str(VLLM_DATA_PARALLEL_RPC_PORT) + assert "--data-parallel-rank" not in cmd + assert "--headless" not in cmd + + def test_dp_per_node_internal_lb_follower_is_headless(self): + """Non-hybrid per-node DP uses a headless follower process.""" + from unittest.mock import MagicMock, patch + + from srtctl.backends import VLLMProtocol, VLLMServerConfig + from srtctl.core.topology import Process + + backend = VLLMProtocol( + dp_launch_mode="per_node", + vllm_config=VLLMServerConfig(decode={"data-parallel-size": 8}), + ) + leader = Process( + node="node0", + gpu_indices=frozenset(range(4)), + sys_port=8081, + http_port=6100, + endpoint_mode="decode", + endpoint_index=0, + node_rank=0, + ) + process = Process( + node="node1", + gpu_indices=frozenset(range(4)), + sys_port=8082, + http_port=6100, + endpoint_mode="decode", + endpoint_index=0, + node_rank=4, + ) + runtime = MagicMock() + runtime.model_path = Path("/model") + runtime.is_hf_model = False + runtime.request_plane = "tcp" + + with patch("srtctl.core.slurm.get_hostname_ip", return_value="10.0.0.1"): + cmd = backend.build_worker_command(process, [leader, process], runtime) + + assert "--data-parallel-size-local" in cmd + assert "--data-parallel-start-rank" in cmd + assert "--data-parallel-hybrid-lb" not in cmd + assert "--headless" in cmd + def test_standard_tp_mode_still_works(self): """Test that standard TP mode (no DP) still creates per-node processes.""" from srtctl.backends import VLLMProtocol, VLLMServerConfig From 3df7d6c7ef414f0bc54d2862768dcbad474e7c3b Mon Sep 17 00:00:00 2001 From: Alec Flowers Date: Sat, 18 Jul 2026 07:27:03 -0700 Subject: [PATCH 2/2] feat(vllm): default DP launch to per-node hybrid LB --- docs/config-reference.md | 37 +++++++----- src/srtctl/backends/vllm.py | 56 ++++++++++++------ src/srtctl/cli/mixins/benchmark_stage.py | 27 +++++---- src/srtctl/core/topology.py | 5 +- tests/test_configs.py | 56 +++++++++--------- tests/test_health_expectations.py | 75 ++++++++++++++++++++++++ 6 files changed, 184 insertions(+), 72 deletions(-) diff --git a/docs/config-reference.md b/docs/config-reference.md index 2e298931..81836861 100644 --- a/docs/config-reference.md +++ b/docs/config-reference.md @@ -422,35 +422,40 @@ Each worker leader gets a globally unique port starting at 5550: | decode_0 | 5552 | | decode_1 | 5553 | -### vLLM DP launch mode +### vLLM DP process layout -vLLM data-parallel endpoints use one process per GPU by default. Set -`dp_launch_mode: per_node` to launch one process per node and let vLLM -manage the local DP ranks in a shared CUDA namespace: +vLLM data-parallel endpoints use hybrid load balancing with one top-level +process per node by default. Each process manages the DP ranks on its local +GPUs and registers that rank range with Dynamo: ```yaml backend: type: vllm - dp_launch_mode: per_node vllm_config: prefill: data-parallel-size: 8 - data-parallel-hybrid-lb: true decode: data-parallel-size: 16 - data-parallel-hybrid-lb: true ``` -| Value | Process layout | -| ---------- | --------------------------------------------------- | -| `per_gpu` | One process per DP rank/GPU (default) | -| `per_node` | One process manages all DP ranks allocated per node | +srtslurm derives `--data-parallel-size-local` and +`--data-parallel-start-rank` from the allocated topology and enables +`--data-parallel-hybrid-lb`. Do not set those flags manually. -In `per_node` mode, srtslurm derives `--data-parallel-size-local` and -`--data-parallel-start-rank` from the allocated topology. Do not set those -two flags manually. Set `data-parallel-hybrid-lb: true` when the frontend -must route to an exact DP rank; without hybrid LB, non-leader node processes -are launched with `--headless` for vLLM's internal DP load balancer. +Set `legacy_dp_per_gpu: true` to restore the previous external-load-balancing +layout, which launches one top-level process and Dynamo engine per DP rank: + +```yaml +backend: + type: vllm + legacy_dp_per_gpu: true + vllm_config: + decode: + data-parallel-size: 16 +``` + +The compatibility flag applies to every vLLM DP role in the deployment. +Standard non-DP tensor/pipeline-parallel endpoints are unaffected. ### TRTLLM Backend diff --git a/src/srtctl/backends/vllm.py b/src/srtctl/backends/vllm.py index d586d87c..2e248da0 100644 --- a/src/srtctl/backends/vllm.py +++ b/src/srtctl/backends/vllm.py @@ -41,7 +41,6 @@ # Type alias for worker modes WorkerMode = Literal["prefill", "decode", "agg"] -DPLaunchMode = Literal["per_gpu", "per_node"] # Filename for the mooncake-store JSON config srtslurm writes to log_dir at job # start. log_dir is mounted into every worker at /logs, so workers read the JSON @@ -139,7 +138,6 @@ class VLLMProtocol: connector: nixl # translated to --kv-transfer-config JSON allow_prefill_decode_colocation: true # pack P/D on one node when all workers fit allow_prefill_decode_colocation_across_nodes: true # continue packing on later nodes - dp_launch_mode: per_node # one process manages all local DP ranks prefill_environment: PYTHONUNBUFFERED: "1" vllm_config: @@ -193,9 +191,9 @@ class VLLMProtocol: # node pools. Defaults off to preserve the original one-node-only policy. allow_prefill_decode_colocation_across_nodes: bool = False - # DP process layout. Keep the existing per-GPU behavior by default; - # per-node lets vLLM manage local DP ranks in one CUDA namespace. - dp_launch_mode: DPLaunchMode = "per_gpu" + # Compatibility escape hatch for the previous DP layout. By default, + # one hybrid-LB process manages all local DP ranks on each node. + legacy_dp_per_gpu: bool = False Schema: ClassVar[builtins.type[Schema]] = Schema @@ -410,7 +408,7 @@ def _is_dp_mode(self, mode: WorkerMode) -> bool: """Check if this mode uses Data Parallel + Expert Parallel pattern. DP+EP mode is detected when data-parallel-size is set in the mode's config. - ``dp_launch_mode`` controls whether a process owns one rank or all local ranks. + The default layout owns all local ranks in one process per node. """ config = self.get_config_for_mode(mode) return config.get("data-parallel-size") is not None or config.get("data_parallel_size") is not None @@ -420,6 +418,29 @@ def _get_dp_size(self, mode: WorkerMode) -> int | None: config = self.get_config_for_mode(mode) return config.get("data-parallel-size") or config.get("data_parallel_size") + def get_expected_dynamo_worker_counts(self, processes: Sequence[Process]) -> tuple[int, int]: + """Return generate-instance counts for the generated process topology. + + External DP registers one instance per GPU process. Hybrid DP registers + one instance per node process. Standard multi-node TP registers only the + endpoint leader, so non-DP processes are grouped by logical endpoint. + """ + grouped: dict[tuple[WorkerMode, int], list[Process]] = {} + for process in processes: + key = (process.endpoint_mode, process.endpoint_index) + grouped.setdefault(key, []).append(process) + + prefill_count = 0 + decode_count = 0 + for (mode, _endpoint_index), endpoint_processes in grouped.items(): + registrations = len(endpoint_processes) if self._is_dp_mode(mode) else 1 + if mode == "prefill": + prefill_count += registrations + else: + decode_count += registrations + + return prefill_count, decode_count + def should_set_cuda_visible_devices(self, process: Process) -> bool: """Whether worker_stage should set CUDA_VISIBLE_DEVICES. @@ -437,7 +458,8 @@ def endpoints_to_processes( ) -> list[Process]: """Convert endpoints to processes. - DP+EP mode uses the configured per-GPU or per-node process layout. + DP+EP mode uses one hybrid-LB process per node by default. The legacy + compatibility flag restores one external-LB process per GPU. For standard TP mode, creates one process per node. """ from srtctl.core.topology import NodePortAllocator, Process, endpoints_to_processes @@ -449,7 +471,7 @@ def endpoints_to_processes( # Standard TP mode: one process per node return endpoints_to_processes(endpoints, base_sys_port=base_sys_port, port_allocator=port_allocator) - if self.dp_launch_mode == "per_node": + if not self.legacy_dp_per_gpu: return self._dp_per_node_endpoints_to_processes( endpoints, base_sys_port=base_sys_port, @@ -567,6 +589,7 @@ def _dp_per_node_endpoints_to_processes( local_dp_size = len(endpoint.gpu_indices) dp_rpc_port = port_allocator.next_dp_rpc_port(endpoint.leader_node) + kv_events_base_port = port_allocator.next_kv_events_port_block(dp_size) nixl_base_port = port_allocator.next_nixl_port_block(dp_size) dp_start_rank = 0 @@ -583,7 +606,8 @@ def _dp_per_node_endpoints_to_processes( bootstrap_port=( port_allocator.next_bootstrap_port(node) if endpoint.mode == "prefill" else None ), - kv_events_port=port_allocator.next_kv_events_port_block(local_dp_size), + # vLLM adds the global DP rank to this endpoint-wide base. + kv_events_port=kv_events_base_port, nixl_port=nixl_base_port, dp_rpc_port=dp_rpc_port, het_group=endpoint.het_group, @@ -672,7 +696,7 @@ def build_worker_command( # Check if this is DP+EP mode (data-parallel-size set) is_dp_mode = self._is_dp_mode(mode) - if is_dp_mode and self.dp_launch_mode == "per_node": + if is_dp_mode and not self.legacy_dp_per_gpu: rpc_port_kebab = config.pop("data-parallel-rpc-port", None) rpc_port_snake = config.pop("data_parallel_rpc_port", None) config_dp_rpc_port = rpc_port_kebab or rpc_port_snake @@ -683,6 +707,8 @@ def build_worker_command( config.pop("data_parallel_size_local", None) config.pop("data-parallel-start-rank", None) config.pop("data_parallel_start_rank", None) + config.pop("data-parallel-hybrid-lb", None) + config.pop("data_parallel_hybrid_lb", None) cmd.extend( [ @@ -694,15 +720,9 @@ def build_worker_command( leader_ip, "--data-parallel-rpc-port", str(dp_rpc_port), + "--data-parallel-hybrid-lb", ] ) - - hybrid_lb = config.get("data-parallel-hybrid-lb", config.get("data_parallel_hybrid_lb", False)) - hybrid_lb_enabled = hybrid_lb is True or ( - isinstance(hybrid_lb, str) and hybrid_lb.strip().lower() in {"1", "true", "yes", "on"} - ) - if process.node_rank > 0 and not hybrid_lb_enabled: - cmd.append("--headless") elif is_dp_mode: # DP+EP mode: each GPU runs its own process # process.node_rank is the dp_rank (set in endpoints_to_processes) @@ -717,6 +737,8 @@ def build_worker_command( # Pop from config so it doesn't get added again by _config_to_cli_args config.pop("data-parallel-rpc-port", None) config.pop("data_parallel_rpc_port", None) + config.pop("data-parallel-hybrid-lb", None) + config.pop("data_parallel_hybrid_lb", None) cmd.extend( [ diff --git a/src/srtctl/cli/mixins/benchmark_stage.py b/src/srtctl/cli/mixins/benchmark_stage.py index fe6cbf1c..9707b1e6 100644 --- a/src/srtctl/cli/mixins/benchmark_stage.py +++ b/src/srtctl/cli/mixins/benchmark_stage.py @@ -11,6 +11,7 @@ import shlex import threading import time +from collections.abc import Sequence from pathlib import Path from typing import TYPE_CHECKING @@ -46,12 +47,14 @@ def _vllm_data_parallel_size(config: "SrtConfig", mode: str) -> int: return int(mode_config.get("data-parallel-size") or mode_config.get("data_parallel_size") or 1) -def _get_health_expectations(config: "SrtConfig") -> tuple[int, int, str, int]: +def _get_health_expectations( + config: "SrtConfig", processes: Sequence["Process"] | None = None +) -> tuple[int, int, str, int]: """Compute expected health counts in the units reported by the frontend. - Dynamo's /health endpoint reports registered generate instances. For vLLM - DP workers, that means one entry per DP rank, not one entry per logical - srt-slurm worker. Other frontends keep using logical worker counts. + Dynamo's /health endpoint reports registered generate instances. External + vLLM DP registers one per rank, while hybrid DP registers one per node. + Other frontends keep using logical worker counts. """ r = config.resources @@ -65,12 +68,16 @@ def _get_health_expectations(config: "SrtConfig") -> tuple[int, int, str, int]: worker_desc = f"{r.num_prefill}P + {r.num_decode}D" if config.frontend.type == "dynamo" and getattr(config.backend, "type", None) == "vllm": - if r.num_agg > 0: - n_prefill = 0 - n_decode = logical_decode * _vllm_data_parallel_size(config, "aggregated") + topology_counter = getattr(config.backend, "get_expected_dynamo_worker_counts", None) + if processes is not None and callable(topology_counter): + n_prefill, n_decode = topology_counter(processes) else: - n_prefill = logical_prefill * _vllm_data_parallel_size(config, "prefill") - n_decode = logical_decode * _vllm_data_parallel_size(config, "decode") + if r.num_agg > 0: + n_prefill = 0 + n_decode = logical_decode * _vllm_data_parallel_size(config, "aggregated") + else: + n_prefill = logical_prefill * _vllm_data_parallel_size(config, "prefill") + n_decode = logical_decode * _vllm_data_parallel_size(config, "decode") count_desc = f"{n_prefill}P + {n_decode}D Dynamo generate instances; logical workers: {worker_desc}" return n_prefill, n_decode, count_desc, n_prefill + n_decode @@ -131,7 +138,7 @@ def run_benchmark( """Run the benchmark.""" logger.info("Waiting for workers to be ready...") - n_prefill, n_decode, count_desc, num_workers = _get_health_expectations(self.config) + n_prefill, n_decode, count_desc, num_workers = _get_health_expectations(self.config, self.backend_processes) logger.info("Waiting for server health (expecting %d health entries: %s)...", num_workers, count_desc) hc = self.config.health_check diff --git a/src/srtctl/core/topology.py b/src/srtctl/core/topology.py index b34cfaca..8049604a 100644 --- a/src/srtctl/core/topology.py +++ b/src/srtctl/core/topology.py @@ -101,8 +101,9 @@ def next_kv_events_port(self) -> int: def next_kv_events_port_block(self, size: int) -> int: """Reserve consecutive KV-event ports and return the base port. - A vLLM hybrid-DP process opens one publisher per local DP rank, so - adjacent worker processes must receive non-overlapping port ranges. + vLLM computes each publisher port as ``base + global_dp_rank``. All + node processes in one endpoint therefore share this base while the + allocator reserves the endpoint's full global DP range. """ if size < 1: raise ValueError("KV-event port block size must be at least 1") diff --git a/tests/test_configs.py b/tests/test_configs.py index 687918d2..75488dd1 100644 --- a/tests/test_configs.py +++ b/tests/test_configs.py @@ -1472,11 +1472,13 @@ def test_same_node_dp_prefill_decode_ports_do_not_collide(self): prefill = [p for p in processes if p.endpoint_mode == "prefill"] decode = [p for p in processes if p.endpoint_mode == "decode"] - assert len(prefill) == 4 - assert len(decode) == 4 + assert len(prefill) == 1 + assert len(decode) == 1 assert {p.node for p in prefill + decode} == {"node0"} assert {p.dp_rpc_port for p in prefill} == {VLLM_DATA_PARALLEL_RPC_PORT} assert {p.dp_rpc_port for p in decode} == {VLLM_DATA_PARALLEL_RPC_PORT + 1} + assert {p.kv_events_port for p in prefill} == {KV_EVENTS_PORT_BASE} + assert {p.kv_events_port for p in decode} == {KV_EVENTS_PORT_BASE + 4} assert {p.nixl_port for p in prefill} == {VLLM_NIXL_PORT_BASE} assert {p.nixl_port for p in decode} == {VLLM_NIXL_PORT_BASE + 4} @@ -1489,8 +1491,8 @@ def test_same_node_dp_prefill_decode_ports_do_not_collide(self): SGLANG_BOOTSTRAP_PORT_BASE, ] - prefill_actual_nixl_ports = {next(iter(p.nixl_port for p in prefill)) + p.node_rank for p in prefill} - decode_actual_nixl_ports = {next(iter(p.nixl_port for p in decode)) + p.node_rank for p in decode} + prefill_actual_nixl_ports = {prefill[0].nixl_port + rank for rank in range(4)} + decode_actual_nixl_ports = {decode[0].nixl_port + rank for rank in range(4)} assert prefill_actual_nixl_ports == {VLLM_NIXL_PORT_BASE + i for i in range(4)} assert decode_actual_nixl_ports == {VLLM_NIXL_PORT_BASE + 4 + i for i in range(4)} assert prefill_actual_nixl_ports.isdisjoint(decode_actual_nixl_ports) @@ -1840,15 +1842,16 @@ def test_dp_mode_detection(self): assert backend_dp._is_dp_mode("decode") is True assert backend_dp._get_dp_size("prefill") == 16 - def test_dp_mode_creates_per_gpu_processes(self): - """Test that DP mode creates one process per GPU instead of per node.""" + def test_legacy_dp_per_gpu_creates_per_gpu_processes(self): + """The compatibility flag restores one process per DP rank.""" from srtctl.backends import VLLMProtocol, VLLMServerConfig from srtctl.core.topology import Endpoint backend = VLLMProtocol( + legacy_dp_per_gpu=True, vllm_config=VLLMServerConfig( prefill={"data-parallel-size": 16, "enable-expert-parallel": True}, - ) + ), ) # Create an endpoint spanning 2 nodes with 8 GPUs each = 16 GPUs total @@ -1885,13 +1888,12 @@ def test_dp_mode_creates_per_gpu_processes(self): dp_ranks = [p.node_rank for p in processes] assert dp_ranks == list(range(16)) - def test_dp_per_node_mode_creates_per_node_processes(self): - """Per-node DP owns all local GPUs and reserves rank-sized port blocks.""" + def test_dp_default_creates_per_node_processes(self): + """Default hybrid DP owns all local GPUs and reserves endpoint port blocks.""" from srtctl.backends import VLLMProtocol, VLLMServerConfig from srtctl.core.topology import Endpoint backend = VLLMProtocol( - dp_launch_mode="per_node", vllm_config=VLLMServerConfig( prefill={"data-parallel-size": 8, "enable-expert-parallel": True}, ), @@ -1911,7 +1913,7 @@ def test_dp_per_node_mode_creates_per_node_processes(self): assert [p.node for p in processes] == ["node0", "node1"] assert all(p.gpu_indices == frozenset(range(4)) for p in processes) assert [p.node_rank for p in processes] == [0, 4] - assert [p.kv_events_port for p in processes] == [KV_EVENTS_PORT_BASE, KV_EVENTS_PORT_BASE + 4] + assert [p.kv_events_port for p in processes] == [KV_EVENTS_PORT_BASE, KV_EVENTS_PORT_BASE] assert {p.nixl_port for p in processes} == {VLLM_NIXL_PORT_BASE} assert {p.dp_rpc_port for p in processes} == {VLLM_DATA_PARALLEL_RPC_PORT} assert {p.het_group for p in processes} == {1} @@ -1924,7 +1926,6 @@ def test_dp_per_node_mode_allocates_non_overlapping_endpoint_ports(self): from srtctl.core.topology import Endpoint backend = VLLMProtocol( - dp_launch_mode="per_node", vllm_config=VLLMServerConfig( decode={"data-parallel-size": 4, "enable-expert-parallel": True}, ), @@ -1961,7 +1962,6 @@ def test_dp_per_node_mode_rejects_dp_size_mismatch(self): from srtctl.core.topology import Endpoint backend = VLLMProtocol( - dp_launch_mode="per_node", vllm_config=VLLMServerConfig(prefill={"data-parallel-size": 7}), ) endpoint = Endpoint( @@ -1975,15 +1975,16 @@ def test_dp_per_node_mode_rejects_dp_size_mismatch(self): with pytest.raises(ValueError, match="data-parallel-size=7"): backend.endpoints_to_processes([endpoint]) - def test_dp_mode_allocates_unique_ports_for_multiple_endpoints_per_node(self): - """Test DP endpoints sharing a node get non-colliding coordination ports.""" + def test_legacy_dp_per_gpu_allocates_unique_ports_for_multiple_endpoints_per_node(self): + """Legacy DP endpoints sharing a node get non-colliding coordination ports.""" from srtctl.backends import VLLMProtocol, VLLMServerConfig from srtctl.core.topology import Endpoint backend = VLLMProtocol( + legacy_dp_per_gpu=True, vllm_config=VLLMServerConfig( decode={"data-parallel-size": 4, "enable-expert-parallel": True}, - ) + ), ) endpoints = [ @@ -2015,8 +2016,8 @@ def test_dp_mode_allocates_unique_ports_for_multiple_endpoints_per_node(self): assert [p.node_rank for p in first_endpoint] == list(range(4)) assert [p.node_rank for p in second_endpoint] == list(range(4)) - def test_dp_mode_command_includes_dp_flags(self): - """Test that DP mode command includes correct DP flags instead of TP flags.""" + def test_legacy_dp_per_gpu_command_includes_external_dp_flags(self): + """The compatibility path emits external-DP flags instead of hybrid flags.""" from pathlib import Path from unittest.mock import MagicMock, patch @@ -2024,13 +2025,15 @@ def test_dp_mode_command_includes_dp_flags(self): from srtctl.core.topology import Process backend = VLLMProtocol( + legacy_dp_per_gpu=True, vllm_config=VLLMServerConfig( prefill={ "data-parallel-size": 16, "data-parallel-rpc-port": 13345, + "data-parallel-hybrid-lb": True, "enable-expert-parallel": True, }, - ) + ), ) # Create a process representing GPU 5 with dp_rank=5 @@ -2090,6 +2093,7 @@ def test_dp_mode_command_includes_dp_flags(self): assert "13345" in cmd assert "--data-parallel-size" in cmd assert "16" in cmd + assert "--data-parallel-hybrid-lb" not in cmd # Should NOT include TP multi-node flags assert "--master-addr" not in cmd @@ -2105,14 +2109,13 @@ def test_dp_per_node_hybrid_command_targets_local_rank_range(self): from srtctl.core.topology import Process backend = VLLMProtocol( - dp_launch_mode="per_node", vllm_config=VLLMServerConfig( decode={ "data-parallel-size": 8, "data-parallel-size-local": 99, "data-parallel-start-rank": 99, "data-parallel-rpc-port": 13345, - "data-parallel-hybrid-lb": True, + "data-parallel-hybrid-lb": False, "enable-expert-parallel": True, }, ), @@ -2154,16 +2157,15 @@ def test_dp_per_node_hybrid_command_targets_local_rank_range(self): assert "--data-parallel-rank" not in cmd assert "--headless" not in cmd - def test_dp_per_node_internal_lb_follower_is_headless(self): - """Non-hybrid per-node DP uses a headless follower process.""" + def test_dp_default_forces_hybrid_for_follower(self): + """The default never falls back to a headless internal-LB follower.""" from unittest.mock import MagicMock, patch from srtctl.backends import VLLMProtocol, VLLMServerConfig from srtctl.core.topology import Process backend = VLLMProtocol( - dp_launch_mode="per_node", - vllm_config=VLLMServerConfig(decode={"data-parallel-size": 8}), + vllm_config=VLLMServerConfig(decode={"data-parallel-size": 8, "data-parallel-hybrid-lb": False}), ) leader = Process( node="node0", @@ -2193,8 +2195,8 @@ def test_dp_per_node_internal_lb_follower_is_headless(self): assert "--data-parallel-size-local" in cmd assert "--data-parallel-start-rank" in cmd - assert "--data-parallel-hybrid-lb" not in cmd - assert "--headless" in cmd + assert cmd.count("--data-parallel-hybrid-lb") == 1 + assert "--headless" not in cmd def test_standard_tp_mode_still_works(self): """Test that standard TP mode (no DP) still creates per-node processes.""" diff --git a/tests/test_health_expectations.py b/tests/test_health_expectations.py index c2c39ce1..9bb77a6b 100644 --- a/tests/test_health_expectations.py +++ b/tests/test_health_expectations.py @@ -5,7 +5,9 @@ from types import SimpleNamespace +from srtctl.backends import VLLMProtocol, VLLMServerConfig from srtctl.cli.mixins.benchmark_stage import _get_health_expectations, _vllm_data_parallel_size +from srtctl.core.topology import Endpoint def _config(frontend_type, backend_type, *, num_prefill=0, num_decode=0, num_agg=0, vllm_config=None): @@ -87,3 +89,76 @@ def test_vllm_data_parallel_size_reads_both_key_styles_and_defaults(): non_vllm = _config("dynamo", "sglang") assert _vllm_data_parallel_size(non_vllm, "prefill") == 1 + + +def test_dynamo_vllm_default_uses_per_node_hybrid_registration_count(): + backend = VLLMProtocol( + vllm_config=VLLMServerConfig( + prefill={"data-parallel-size": 8}, + decode={"data-parallel-size": 4}, + ) + ) + config = SimpleNamespace( + frontend=SimpleNamespace(type="dynamo"), + backend=backend, + resources=SimpleNamespace(num_prefill=1, num_decode=1, num_agg=0), + ) + processes = backend.endpoints_to_processes( + [ + Endpoint( + mode="prefill", + index=0, + nodes=("node0", "node1"), + gpu_indices=frozenset(range(4)), + gpus_per_node=4, + ), + Endpoint( + mode="decode", + index=0, + nodes=("node2",), + gpu_indices=frozenset(range(4)), + gpus_per_node=4, + ), + ] + ) + + n_prefill, n_decode, _count_desc, num_workers = _get_health_expectations(config, processes) + + assert (n_prefill, n_decode, num_workers) == (2, 1, 3) + + +def test_dynamo_vllm_legacy_uses_per_gpu_registration_count(): + backend = VLLMProtocol( + legacy_dp_per_gpu=True, + vllm_config=VLLMServerConfig( + prefill={"data-parallel-size": 8}, + decode={"data-parallel-size": 4}, + ), + ) + config = SimpleNamespace( + frontend=SimpleNamespace(type="dynamo"), + backend=backend, + resources=SimpleNamespace(num_prefill=1, num_decode=1, num_agg=0), + ) + processes = backend.endpoints_to_processes( + [ + Endpoint( + mode="prefill", + index=0, + nodes=("node0", "node1"), + gpu_indices=frozenset(range(4)), + gpus_per_node=4, + ), + Endpoint( + mode="decode", + index=0, + nodes=("node2",), + gpu_indices=frozenset(range(4)), + gpus_per_node=4, + ), + ] + ) + + n_prefill, n_decode, _count_desc, num_workers = _get_health_expectations(config, processes) + + assert (n_prefill, n_decode, num_workers) == (8, 4, 12)