diff --git a/docs/config-reference.md b/docs/config-reference.md index 9717a58f..81836861 100644 --- a/docs/config-reference.md +++ b/docs/config-reference.md @@ -422,6 +422,41 @@ Each worker leader gets a globally unique port starting at 5550: | decode_0 | 5552 | | decode_1 | 5553 | +### vLLM DP process layout + +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 + vllm_config: + prefill: + data-parallel-size: 8 + decode: + data-parallel-size: 16 +``` + +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. + +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 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..2e248da0 100644 --- a/src/srtctl/backends/vllm.py +++ b/src/srtctl/backends/vllm.py @@ -191,6 +191,10 @@ class VLLMProtocol: # node pools. Defaults off to preserve the original one-node-only policy. allow_prefill_decode_colocation_across_nodes: bool = False + # 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 # ========================================================================= @@ -198,7 +202,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 +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. - In this mode, each GPU runs its own process (rather than TP across GPUs). + 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 @@ -414,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. @@ -431,7 +458,8 @@ 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 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 @@ -443,6 +471,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 not self.legacy_dp_per_gpu: + 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 +555,69 @@ 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) + 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 + + 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 + ), + # 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, + ) + ) + current_sys_port += 1 + dp_start_rank += local_dp_size + + return processes + def build_worker_command( self, process: Process, @@ -598,7 +696,34 @@ 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 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 + 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) + config.pop("data-parallel-hybrid-lb", None) + config.pop("data_parallel_hybrid_lb", 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), + "--data-parallel-hybrid-lb", + ] + ) + 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 @@ -612,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 e662e5f0..8049604a 100644 --- a/src/srtctl/core/topology.py +++ b/src/srtctl/core/topology.py @@ -98,6 +98,21 @@ 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. + + 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") + 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..75488dd1 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, @@ -1471,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} @@ -1488,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) @@ -1839,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 @@ -1884,15 +1888,103 @@ 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_mode_allocates_unique_ports_for_multiple_endpoints_per_node(self): - """Test DP endpoints sharing a node get non-colliding coordination ports.""" + 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( + 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] + 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( 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( + 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_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 = [ @@ -1924,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 @@ -1933,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 @@ -1999,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 @@ -2006,6 +2101,103 @@ 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( + 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": False, + "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_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( + vllm_config=VLLMServerConfig(decode={"data-parallel-size": 8, "data-parallel-hybrid-lb": False}), + ) + 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 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.""" from srtctl.backends import VLLMProtocol, VLLMServerConfig 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)