From 1a0f9e3633318ab1ee9428d2129161b583786b18 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 2d9196fe6c4fc44a62bbab518fa873d030110ece Mon Sep 17 00:00:00 2001 From: inf-yasong Date: Mon, 13 Jul 2026 13:36:09 +0000 Subject: [PATCH 2/2] feat(vllm): support per-role DP launch modes --- docs/config-reference.md | 8 ++ src/srtctl/backends/vllm.py | 104 +++++++++++++++++++++-- src/srtctl/cli/mixins/benchmark_stage.py | 11 ++- tests/test_configs.py | 49 +++++++++++ tests/test_health_expectations.py | 16 ++++ 5 files changed, 178 insertions(+), 10 deletions(-) diff --git a/docs/config-reference.md b/docs/config-reference.md index 2e298931..d6b80222 100644 --- a/docs/config-reference.md +++ b/docs/config-reference.md @@ -432,6 +432,7 @@ manage the local DP ranks in a shared CUDA namespace: backend: type: vllm dp_launch_mode: per_node + decode_dp_launch_mode: per_gpu # Optional role override vllm_config: prefill: data-parallel-size: 8 @@ -446,6 +447,13 @@ backend: | `per_gpu` | One process per DP rank/GPU (default) | | `per_node` | One process manages all DP ranks allocated per node | +`prefill_dp_launch_mode`, `decode_dp_launch_mode`, and +`aggregated_dp_launch_mode` can override the global mode for one role. For +example, keeping prefill on `per_node` while setting decode to `per_gpu` +preserves a shared CUDA namespace for prefill collectives while exposing each +decode rank as an independent Dynamo worker. Roles without an override inherit +`dp_launch_mode`. + 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 diff --git a/src/srtctl/backends/vllm.py b/src/srtctl/backends/vllm.py index d586d87c..0484d701 100644 --- a/src/srtctl/backends/vllm.py +++ b/src/srtctl/backends/vllm.py @@ -140,6 +140,7 @@ class VLLMProtocol: 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 + decode_dp_launch_mode: per_gpu # optional per-role override prefill_environment: PYTHONUNBUFFERED: "1" vllm_config: @@ -197,6 +198,11 @@ class VLLMProtocol: # per-node lets vLLM manage local DP ranks in one CUDA namespace. dp_launch_mode: DPLaunchMode = "per_gpu" + # Optional per-role overrides. Unset roles inherit dp_launch_mode. + prefill_dp_launch_mode: DPLaunchMode | None = None + decode_dp_launch_mode: DPLaunchMode | None = None + aggregated_dp_launch_mode: DPLaunchMode | None = None + Schema: ClassVar[builtins.type[Schema]] = Schema # ========================================================================= @@ -420,6 +426,50 @@ 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_dp_launch_mode_for_mode(self, mode: WorkerMode) -> DPLaunchMode: + """Return the role-specific launch mode or the global default.""" + overrides = { + "prefill": self.prefill_dp_launch_mode, + "decode": self.decode_dp_launch_mode, + "agg": self.aggregated_dp_launch_mode, + } + return overrides[mode] or self.dp_launch_mode + + def get_expected_dynamo_worker_counts(self, processes: Sequence[Process]) -> tuple[int, int]: + """Return the generate-instance counts expected from a process layout. + + Per-GPU processes each register one instance. Hybrid per-node processes + register their local rank range, while non-hybrid headless followers do + not register separately. + """ + 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 = 1 + if self._is_dp_mode(mode): + launch_mode = self.get_dp_launch_mode_for_mode(mode) + if launch_mode == "per_gpu": + registrations = len(endpoint_processes) + else: + config = self.get_config_for_mode(mode) + 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"} + ) + registrations = len(endpoint_processes) if hybrid_lb_enabled 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. @@ -440,22 +490,62 @@ def endpoints_to_processes( 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 + from srtctl.core.topology import NodePortAllocator, endpoints_to_processes - # Check if any endpoint uses DP mode - has_dp_mode = any(self._is_dp_mode(ep.mode) for ep in endpoints) - - if not has_dp_mode: + dp_endpoints = [endpoint for endpoint in endpoints if self._is_dp_mode(endpoint.mode)] + if not dp_endpoints: # 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": + launch_modes = {self.get_dp_launch_mode_for_mode(endpoint.mode) for endpoint in dp_endpoints} + if launch_modes == {"per_node"}: return self._dp_per_node_endpoints_to_processes( endpoints, base_sys_port=base_sys_port, port_allocator=port_allocator, ) + if launch_modes == {"per_gpu"}: + return self._dp_per_gpu_endpoints_to_processes( + endpoints, + base_sys_port=base_sys_port, + port_allocator=port_allocator, + ) + + # Mixed role modes share one allocator so every coordination port stays + # unique across the per-node and per-GPU process groups. + processes: list[Process] = [] + current_sys_port = base_sys_port + if port_allocator is None: + port_allocator = NodePortAllocator() + + for endpoint in endpoints: + if self._is_dp_mode(endpoint.mode) and self.get_dp_launch_mode_for_mode(endpoint.mode) == "per_node": + endpoint_processes = self._dp_per_node_endpoints_to_processes( + [endpoint], + base_sys_port=current_sys_port, + port_allocator=port_allocator, + ) + else: + endpoint_processes = self._dp_per_gpu_endpoints_to_processes( + [endpoint], + base_sys_port=current_sys_port, + port_allocator=port_allocator, + ) + processes.extend(endpoint_processes) + current_sys_port += len(endpoint_processes) + + return processes + + def _dp_per_gpu_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 GPU.""" + from srtctl.core.topology import NodePortAllocator, Process + # DP+EP mode: one process per GPU processes: list[Process] = [] current_sys_port = base_sys_port @@ -672,7 +762,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 self.get_dp_launch_mode_for_mode(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 diff --git a/src/srtctl/cli/mixins/benchmark_stage.py b/src/srtctl/cli/mixins/benchmark_stage.py index fe6cbf1c..1c455c0b 100644 --- a/src/srtctl/cli/mixins/benchmark_stage.py +++ b/src/srtctl/cli/mixins/benchmark_stage.py @@ -46,7 +46,9 @@ 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: list["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 @@ -65,7 +67,10 @@ 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: + registration_counter = getattr(config.backend, "get_expected_dynamo_worker_counts", None) + if processes is not None and callable(registration_counter): + n_prefill, n_decode = registration_counter(processes) + elif r.num_agg > 0: n_prefill = 0 n_decode = logical_decode * _vllm_data_parallel_size(config, "aggregated") else: @@ -131,7 +136,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/tests/test_configs.py b/tests/test_configs.py index 687918d2..fa6902a5 100644 --- a/tests/test_configs.py +++ b/tests/test_configs.py @@ -1975,6 +1975,55 @@ 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_launch_mode_supports_per_role_overrides(self): + """Prefill and decode can use different DP process layouts.""" + from srtctl.backends import VLLMProtocol, VLLMServerConfig + from srtctl.core.topology import Endpoint + + backend = VLLMProtocol( + dp_launch_mode="per_node", + decode_dp_launch_mode="per_gpu", + vllm_config=VLLMServerConfig( + prefill={ + "data-parallel-size": 8, + "data-parallel-hybrid-lb": True, + "enable-expert-parallel": True, + }, + decode={"data-parallel-size": 16, "enable-expert-parallel": True}, + ), + ) + endpoints = [ + Endpoint( + mode="prefill", + index=0, + nodes=("pnode0", "pnode1"), + gpu_indices=frozenset(range(4)), + gpus_per_node=4, + ), + Endpoint( + mode="decode", + index=0, + nodes=("dnode0", "dnode1", "dnode2", "dnode3"), + gpu_indices=frozenset(range(4)), + gpus_per_node=4, + ), + ] + + processes = backend.endpoints_to_processes(endpoints) + prefill_processes = [process for process in processes if process.endpoint_mode == "prefill"] + decode_processes = [process for process in processes if process.endpoint_mode == "decode"] + + assert backend.get_dp_launch_mode_for_mode("prefill") == "per_node" + assert backend.get_dp_launch_mode_for_mode("decode") == "per_gpu" + assert len(prefill_processes) == 2 + assert all(process.gpu_indices == frozenset(range(4)) for process in prefill_processes) + assert [process.node_rank for process in prefill_processes] == [0, 4] + assert len(decode_processes) == 16 + assert all(len(process.gpu_indices) == 1 for process in decode_processes) + assert [process.node_rank for process in decode_processes] == list(range(16)) + assert len({process.sys_port for process in processes}) == len(processes) + assert backend.get_expected_dynamo_worker_counts(processes) == (2, 16) + 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 diff --git a/tests/test_health_expectations.py b/tests/test_health_expectations.py index c2c39ce1..42aae81d 100644 --- a/tests/test_health_expectations.py +++ b/tests/test_health_expectations.py @@ -54,6 +54,22 @@ def test_dynamo_vllm_without_dp_config_defaults_to_logical_counts(): assert (n_prefill, n_decode, num_workers) == (6, 1, 7) +def test_dynamo_vllm_uses_process_aware_counts_for_mixed_dp_layout(): + """Mixed launch modes use actual registrations rather than global DP sizes.""" + vllm_config = SimpleNamespace( + prefill={"data-parallel-size": 8}, + decode={"data-parallel-size": 16}, + aggregated=None, + ) + config = _config("dynamo", "vllm", num_prefill=1, num_decode=1, vllm_config=vllm_config) + config.backend.get_expected_dynamo_worker_counts = lambda _processes: (2, 16) + + n_prefill, n_decode, count_desc, num_workers = _get_health_expectations(config, [object()]) + + assert (n_prefill, n_decode, num_workers) == (2, 16, 18) + assert count_desc == "2P + 16D Dynamo generate instances; logical workers: 1P + 1D" + + def test_non_dynamo_frontend_uses_logical_worker_counts(): """Only Dynamo reports per-DP-rank generate instances; others stay logical.""" vllm_config = SimpleNamespace(prefill={"data-parallel-size": 2}, decode={"data-parallel-size": 8}, aggregated=None)