diff --git a/CLAUDE.md b/CLAUDE.md index 44b5e0c90..fd7dd11a9 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -84,7 +84,7 @@ Two patterns for checking worker readiness: check_dynamo_health(response_json, expected_prefill=2, expected_decode=4) # SGLang router -check_sglang_router_health(response_json, expected_prefill=2, expected_decode=4) +check_static_router_health(response_json, expected_prefill=2, expected_decode=4) ``` For aggregated mode, pass `expected_prefill=0, expected_decode=num_agg`. @@ -279,4 +279,3 @@ tail -f outputs//logs/sweep_.log | grep "srun command" ``` Per-worker env vars and commands are also logged individually (search for `Env:` and `Command:` lines). - diff --git a/docs/config-reference.md b/docs/config-reference.md index 2d751cf4c..6d142ac67 100644 --- a/docs/config-reference.md +++ b/docs/config-reference.md @@ -270,7 +270,8 @@ Frontend/router configuration. ```yaml frontend: - # Frontend type: "dynamo" (default), "sglang", or "trtllm_serve" + # Frontend type: "dynamo" (default), "sglang", "vllm-router", + # "trtllm_serve", or direct "vllm". type: dynamo # Scaling @@ -290,20 +291,43 @@ frontend: # Environment variables for frontend processes env: MY_VAR: "value" + + # Optional router-specific image; defaults to model.container + container_image: "router-image" ``` | Field | Type | Default | Description | | --------------------------- | ---- | ------------- | ----------------------------------- | -| `type` | str | dynamo | Frontend type: "dynamo", "sglang", or "trtllm_serve" | +| `type` | str | dynamo | Frontend type: `dynamo`, `sglang`, `vllm-router`, `trtllm_serve`, or direct `vllm` | | `enable_multiple_frontends` | bool | true | Scale with nginx + multiple routers | | `num_additional_frontends` | int | 9 | Additional routers beyond master | | `nginx_container` | str | nginx:1.27.4 | Custom nginx container image | | `nginx_raise_ulimit` | bool | false | When true with nginx in use, run `ulimit -n 1048576` before nginx and emit `worker_rlimit_nofile 1048576` in generated `nginx.conf`. Off by default so restrictive clusters do not fail. Cluster `srtslurm.yaml` may set `nginx_raise_ulimit` for jobs that omit this field. | | `args` | dict | null | CLI args for the frontend | | `env` | dict | null | Env vars for frontend processes | +| `container_image` | str | null | Router process image; defaults to `model.container` | + +For `vllm-router`, srtctl sets Router's `--worker-startup-timeout-secs` to the +total `health_check` window so large-model compilation cannot outlive the router. +Set `frontend.args.worker-startup-timeout-secs` to override it explicitly. See [SGLang Router](sglang-router.md) for detailed architecture. +### vllm-router frontend + +`type: vllm-router` pairs with `backend.type: vllm` and launches the official +`vllm-router` process against direct private `vllm serve` endpoints. Aggregate +layouts use `--worker-urls`; disaggregated layouts use +`--vllm-pd-disaggregation` with the allocated prefill and decode URLs. For +data-parallel endpoints, srtctl derives Router's +`--intra-node-data-parallel-size`. Router expands each node-local backend URL +into DP-aware targets and injects `X-Data-Parallel-Rank`; vLLM continues to own +the engine processes behind that HTTP server. Multi-node DP endpoints use one +hybrid-LB `vllm serve` process per node and require +`backend.dp_launch_mode: per_node`. Direct `frontend.type: vllm` retains its +existing single-server behavior. No NATS or etcd infrastructure is started for +this frontend. + ### trtllm_serve frontend `type: trtllm_serve` runs the `trtllm-serve disaggregated` orchestrator as the @@ -424,9 +448,16 @@ Each worker leader gets a globally unique port starting at 5550: ### 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: +Direct `frontend.type: vllm` and single-node `frontend.type: vllm-router` +deployments preserve vLLM's native behavior: one `vllm serve` command owns the +configured TP, PP, PCP, and DP ranks and performs its normal internal DP load +balancing. `dp_launch_mode` controls only the cases where srtctl must split a +DP deployment into separately launched processes (Dynamo compatibility or a +multi-node vLLM Router endpoint). + +Set `dp_launch_mode: per_node` for a multi-node DP endpoint. srtctl launches one +hybrid-LB server per node and derives the local DP rank range from the model +parallelism and Slurm allocation: ```yaml backend: @@ -439,26 +470,52 @@ backend: data-parallel-size: 16 ``` -| Value | Process layout | -| ---------- | --------------------------------------------------- | -| `per_gpu` | One process per DP rank/GPU (default) | -| `per_node` | One process manages all DP ranks allocated per node | - -`per_gpu` remains the compatibility default for now, but srtslurm will switch -the default to `per_node` in a future release. Existing vLLM DP configurations -should set `backend.dp_launch_mode: per_node` now; srtslurm emits a -configuration-time migration warning while they still use `per_gpu`. - -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. srtslurm also always enables `--data-parallel-hybrid-lb` -so every node-local process registers with the Dynamo frontend. This is the -recommended vLLM topology for Dynamo and ensures the frontend can route to -each node-local DP engine. Do not set `data-parallel-hybrid-lb` manually; -srtslurm enables it automatically, warns when it is configured, and ignores -the configured value. `headless` is incompatible with `per_node` DP because a -headless process does not register with Dynamo, so srtslurm rejects that -combination during configuration loading. +| Value | Process layout | +| ---------- | --------------------------------------------------------------------- | +| `per_rank` | Legacy Dynamo layout: one process per DP rank; requires TP=PP=PCP=1 | +| `per_node` | One process manages all complete node-local DP replicas | + +`per_gpu` remains accepted as a deprecated compatibility alias for `per_rank`. +The process boundary is the DP rank; it corresponds to one GPU only under the +required TP=PP=PCP=1 topology. + +For `frontend.type: vllm-router`, Router-native DP expansion keeps one backend +URL per node and sends `X-Data-Parallel-Rank` to select a node-local engine. A +single-node endpoint needs no special launch mode. A multi-node DP endpoint must +use `per_node`; srtctl then derives the global/local DP topology and advertises +each node-local HTTP server to Router. All routed backends must have the same +node-local DP size because Router exposes one global +`--intra-node-data-parallel-size` setting. + +For every vLLM Router worker, srtctl validates that the allocation exactly +matches vLLM's process world: `GPUs = DP * TP * PP * PCP`. In `per_node` mode, +the derived local DP size is `GPUs on the node / (TP * PP * PCP)`. Non-integral +or under/over-allocated layouts fail during recipe loading, before Slurm starts. + +In `per_node` mode, srtctl owns the scheduling fields +`data-parallel-size-local`, `data-parallel-start-rank`, +`data-parallel-address`, `data-parallel-rpc-port`, +`data-parallel-hybrid-lb`, and `headless`. Recipes must not set them manually; +srtctl rejects those combinations rather than silently overriding potentially +contradictory values. + +### vLLM device binding compatibility + +Direct `vllm` and `vllm-router` frontends use vLLM's `--device-ids` option by +default. For stable vLLM releases from before +[vllm-project/vllm#45026](https://github.com/vllm-project/vllm/pull/45026), +select the existing CUDA namespace binding instead: + +```yaml +backend: + type: vllm + set_cuda_visible_devices: true +``` + +srtslurm then omits `--device-ids` and scopes sub-node workers with +`CUDA_VISIBLE_DEVICES`. A worker that owns every GPU on its node needs no +explicit CUDA mask. This setting changes only device binding; worker topology, +private router ports, and P/D KV-transfer arguments remain unchanged. ### TRTLLM Backend diff --git a/docs/sglang-router.md b/docs/sglang-router.md index b4a48b74b..68f3de15f 100644 --- a/docs/sglang-router.md +++ b/docs/sglang-router.md @@ -1,6 +1,7 @@ # SGLang Router Mode -This page explains the sglang router mode for prefill-decode (PD) disaggregation, an alternative to the default Dynamo frontend architecture. +This page explains the first-class SGLang Model Gateway router mode for aggregate +and prefill-decode (PD) topologies, an alternative to the default Dynamo frontend. ## Table of Contents @@ -42,7 +43,9 @@ frontend: type: sglang ``` -That's it. The workers will launch with `sglang.launch_server` instead of `dynamo.sglang`, and the router will handle request distribution. +Workers launch with `sglang.launch_server` instead of `dynamo.sglang`, and the +router receives only logical worker-leader URLs from srtctl's allocated +topology. ### Router Arguments diff --git a/src/srtctl/backends/sglang.py b/src/srtctl/backends/sglang.py index 22510ae98..833442451 100644 --- a/src/srtctl/backends/sglang.py +++ b/src/srtctl/backends/sglang.py @@ -296,7 +296,8 @@ def build_worker_command( process: The process to start endpoint_processes: All processes for this endpoint (for multi-node) runtime: Runtime context with paths and settings - frontend_type: Frontend type - "sglang" uses sglang.launch_server, "dynamo" uses dynamo.sglang + frontend_type: Frontend type - "sglang" uses sglang.launch_server, + while "dynamo" uses dynamo.sglang nsys_prefix: Optional nsys profiling command prefix dump_config_path: Path to dump config JSON """ @@ -354,6 +355,14 @@ def build_worker_command( # Add disaggregation mode for prefill/decode workers (both dynamo and sglang frontend) if mode != "agg": cmd.extend(["--disaggregation-mode", mode]) + if use_sglang: + # Direct P/D workers are started together and validated through the + # static router. SGLang's built-in disaggregation warmup instead + # posts a synthetic request to the local worker with a fake + # bootstrap host, which cannot exercise this topology and blocks + # server readiness until its 30-minute timeout. The benchmark + # warmup exercises the real router-mediated P/D transfer path. + cmd.append("--skip-server-warmup") # Always pass bootstrap port for prefill workers regardless of frontend type. # Dynamo does NOT handle this internally — SGLang's CommonKVBootstrapServer # still runs on every prefill node for KV transfer coordination, and workers diff --git a/src/srtctl/backends/vllm.py b/src/srtctl/backends/vllm.py index 0588df212..13fb9277a 100644 --- a/src/srtctl/backends/vllm.py +++ b/src/srtctl/backends/vllm.py @@ -12,7 +12,6 @@ import builtins import json -import logging from collections.abc import Sequence from dataclasses import field from pathlib import Path @@ -43,9 +42,10 @@ # Type alias for worker modes WorkerMode = Literal["prefill", "decode", "agg"] -DPLaunchMode = Literal["per_gpu", "per_node"] - -logger = logging.getLogger(__name__) +# ``per_gpu`` is retained as a compatibility alias for existing recipes. A +# process owns one DP rank, which is only equivalent to one GPU while the +# per-rank layout is restricted to TP=PP=PCP=1. +DPLaunchMode = Literal["per_rank", "per_node", "per_gpu"] # 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 @@ -198,10 +198,12 @@ 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. - # TODO: Change the default to per_node after the per_gpu migration window. - dp_launch_mode: DPLaunchMode = "per_gpu" + # DP process layout for orchestrated multi-process deployments. Direct vLLM + # and single-node vLLM Router preserve native one-command vLLM behavior. + # Keep the per-rank layout as the Dynamo compatibility default; per_node + # lets vLLM manage complete node-local DP replicas in one CUDA namespace. + # ``per_gpu`` is a deprecated compatibility spelling for ``per_rank``. + dp_launch_mode: DPLaunchMode = "per_rank" Schema: ClassVar[builtins.type[Schema]] = Schema @@ -210,54 +212,65 @@ def __post_init__(self) -> None: if self.vllm_config is None: return + configured_dp_mode_configs: list[tuple[str, dict[str, Any]]] = [] dp_mode_configs: list[tuple[str, dict[str, Any]]] = [] for mode_name, mode_config in ( ("prefill", self.vllm_config.prefill), ("decode", self.vllm_config.decode), ("aggregated", self.vllm_config.aggregated), ): - if mode_config and any( - str(key).replace("_", "-") == "data-parallel-size" and value is not None - for key, value in mode_config.items() - ): + configured_dp_size = next( + ( + value + for key, value in (mode_config or {}).items() + if str(key).replace("_", "-") == "data-parallel-size" + ), + None, + ) + if configured_dp_size is None: + continue + try: + normalized_dp_size = int(configured_dp_size) + except (TypeError, ValueError) as exc: + raise ValidationError( + f"vllm_config.{mode_name}.data-parallel-size must be a positive integer; got {configured_dp_size!r}" + ) from exc + if normalized_dp_size < 1: + raise ValidationError( + f"vllm_config.{mode_name}.data-parallel-size must be a positive integer; got {configured_dp_size!r}" + ) + configured_dp_mode_configs.append((mode_name, mode_config)) + if normalized_dp_size > 1: dp_mode_configs.append((mode_name, mode_config)) + if self.dp_launch_mode == "per_node": + managed_topology_fields: list[str] = [] + for mode_name, mode_config in configured_dp_mode_configs: + normalized_keys = {str(key).replace("_", "-") for key in mode_config} + for field_name in ( + "data-parallel-size-local", + "data-parallel-start-rank", + "data-parallel-address", + "data-parallel-rpc-port", + "data-parallel-hybrid-lb", + "headless", + ): + if field_name in normalized_keys: + managed_topology_fields.append(f"vllm_config.{mode_name}.{field_name}") + + if managed_topology_fields: + fields = ", ".join(managed_topology_fields) + raise ValidationError( + f"{fields} cannot be set when dp_launch_mode=per_node; " + "srt-slurm derives the node-local DP topology from the Slurm allocation" + ) + if not dp_mode_configs: return - if self.dp_launch_mode == "per_gpu": - modes = ", ".join(mode_name for mode_name, _ in dp_mode_configs) - logger.warning( - "vLLM DP mode(s) %s use dp_launch_mode=per_gpu. per_node is the recommended topology " - "and will become the default in a future release; set backend.dp_launch_mode: per_node now", - modes, - ) + if self.dp_launch_mode in {"per_rank", "per_gpu"}: return - hybrid_lb_modes: list[str] = [] - headless_modes: list[str] = [] - for mode_name, mode_config in dp_mode_configs: - normalized_keys = {str(key).replace("_", "-") for key in mode_config} - if "headless" in normalized_keys: - headless_modes.append(mode_name) - if "data-parallel-hybrid-lb" in normalized_keys: - hybrid_lb_modes.append(mode_name) - - if headless_modes: - fields = ", ".join(f"vllm_config.{mode}.headless" for mode in headless_modes) - raise ValidationError( - f"{fields} cannot be set when dp_launch_mode=per_node. " - "Every node-local process must register with the Dynamo frontend; remove headless." - ) - - if hybrid_lb_modes: - fields = ", ".join(f"vllm_config.{mode}.data-parallel-hybrid-lb" for mode in hybrid_lb_modes) - logger.warning( - "%s is unnecessary when dp_launch_mode=per_node; " - "srtslurm always enables --data-parallel-hybrid-lb and ignores the configured value", - fields, - ) - # ========================================================================= # BackendProtocol Implementation # ========================================================================= @@ -471,13 +484,67 @@ def _is_dp_mode(self, mode: WorkerMode) -> bool: 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. """ - 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 + dp_size = self._get_dp_size(mode) + return dp_size is not None and int(dp_size) > 1 def _get_dp_size(self, mode: WorkerMode) -> int | None: """Get the data-parallel-size for a mode, or None if not in DP mode.""" config = self.get_config_for_mode(mode) - return config.get("data-parallel-size") or config.get("data_parallel_size") + if "data-parallel-size" in config: + return config["data-parallel-size"] + return config.get("data_parallel_size") + + def _get_parallel_size(self, mode: WorkerMode, name: str) -> int: + """Return a normalized vLLM parallel dimension, defaulting to one.""" + config = self.get_config_for_mode(mode) + value = config.get(name, config.get(name.replace("-", "_"), 1)) + try: + size = int(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"vLLM {mode} {name} must be a positive integer; got {value!r}") from exc + if size < 1: + raise ValueError(f"vLLM {mode} {name} must be a positive integer; got {value!r}") + return size + + def _get_model_parallel_size(self, mode: WorkerMode) -> int: + """Return GPUs consumed by one DP replica using vLLM's world-size dimensions.""" + return ( + self._get_parallel_size(mode, "tensor-parallel-size") + * self._get_parallel_size(mode, "pipeline-parallel-size") + * self._get_parallel_size(mode, "prefill-context-parallel-size") + ) + + def _get_local_dp_size(self, mode: WorkerMode, local_gpu_count: int) -> int: + """Derive how many complete DP replicas fit in a node-local GPU allocation.""" + model_parallel_size = self._get_model_parallel_size(mode) + if local_gpu_count % model_parallel_size != 0: + raise ValueError( + f"vLLM {mode} node-local allocation has {local_gpu_count} GPUs, which is not divisible by " + f"TP*PP*PCP={model_parallel_size}" + ) + return local_gpu_count // model_parallel_size + + def _validate_endpoint_parallelism(self, endpoint: Endpoint) -> tuple[int, int]: + """Validate vLLM's DP x TP x PP x PCP world against an endpoint allocation.""" + dp_size = self._get_dp_size(endpoint.mode) or 1 + try: + dp_size = int(dp_size) + except (TypeError, ValueError) as exc: + raise ValueError( + f"vLLM {endpoint.mode} data-parallel-size must be a positive integer; got {dp_size!r}" + ) from exc + if dp_size < 1: + raise ValueError(f"vLLM {endpoint.mode} data-parallel-size must be a positive integer; got {dp_size!r}") + + model_parallel_size = self._get_model_parallel_size(endpoint.mode) + required_gpus = dp_size * model_parallel_size + if required_gpus != endpoint.total_gpus: + raise ValueError( + f"vLLM {endpoint.mode} parallelism requires DP*TP*PP*PCP=" + f"{dp_size}*{model_parallel_size}={required_gpus} GPUs, but the endpoint allocates " + f"{endpoint.total_gpus} GPUs" + ) + return dp_size, model_parallel_size def should_set_cuda_visible_devices(self, process: Process) -> bool: """Whether worker_stage should set CUDA_VISIBLE_DEVICES. @@ -497,14 +564,17 @@ def endpoints_to_processes( ) -> list[Process]: """Convert endpoints to processes. - Dynamo DP+EP mode uses the configured per-GPU or per-node process layout. - For direct vLLM aggregate jobs, `vllm serve` manages local DP ranks from - one process, so keep the standard one-process-per-node topology. + Dynamo DP+EP mode uses the configured per-rank or per-node process layout. + For direct vLLM and single-node vLLM Router jobs, `vllm serve` manages + local DP ranks from one process. Multi-node vLLM Router DP jobs use one + hybrid-LB process per node so Router can address each node-local DP pool. For standard TP mode, creates one process per node. """ from srtctl.core.topology import NodePortAllocator, Process, endpoints_to_processes if frontend_type == "vllm": + for endpoint in endpoints: + self._validate_endpoint_parallelism(endpoint) return endpoints_to_processes(endpoints, base_sys_port=base_sys_port, port_allocator=port_allocator) # Check if any endpoint uses DP mode @@ -514,6 +584,17 @@ 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 frontend_type == "vllm-router" and all(not endpoint.is_multi_node for endpoint in endpoints): + # Preserve vLLM's native single-command topology. The server owns TP/PP/DP + # and its internal load balancer; Router expands the advertised URL into + # the configured node-local DP ranks. + for endpoint in endpoints: + self._validate_endpoint_parallelism(endpoint) + return endpoints_to_processes(endpoints, base_sys_port=base_sys_port, port_allocator=port_allocator) + + if frontend_type == "vllm-router" and self.dp_launch_mode != "per_node": + raise ValueError("multi-node vLLM Router DP endpoints require backend.dp_launch_mode: per_node") + if self.dp_launch_mode == "per_node": return self._dp_per_node_endpoints_to_processes( endpoints, @@ -521,7 +602,18 @@ def endpoints_to_processes( port_allocator=port_allocator, ) - # DP+EP mode: one process per GPU + for endpoint in endpoints: + if not self._is_dp_mode(endpoint.mode): + continue + _dp_size, model_parallel_size = self._validate_endpoint_parallelism(endpoint) + if model_parallel_size != 1: + raise ValueError( + f"vLLM {endpoint.mode} dp_launch_mode=per_rank supports only TP=PP=PCP=1; " + "use dp_launch_mode=per_node for combined DP and model parallelism" + ) + + # DP+EP mode: one process per DP rank. The supported TP=PP=PCP=1 + # topology means each rank currently owns one GPU. processes: list[Process] = [] current_sys_port = base_sys_port if port_allocator is None: @@ -556,8 +648,8 @@ def endpoints_to_processes( ) current_sys_port += 1 else: - # DP+EP mode: one process per GPU - # Each process gets a single GPU and a unique dp_rank + # DP+EP mode: one process per DP rank. Each process gets a + # single GPU under the required TP=PP=PCP=1 topology. dp_rank = 0 # Allocate a unique DP RPC port for this endpoint's leader node dp_rpc_port = port_allocator.next_dp_rpc_port(endpoint.leader_node) @@ -623,14 +715,8 @@ def _dp_per_node_endpoints_to_processes( 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_size, _model_parallel_size = self._validate_endpoint_parallelism(endpoint) + local_dp_size = self._get_local_dp_size(endpoint.mode, 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 @@ -675,7 +761,7 @@ def build_worker_command( process: The process to start endpoint_processes: All processes for this endpoint (for multi-node) runtime: Runtime context with paths and settings - frontend_type: Frontend type ("dynamo" or direct "vllm") + frontend_type: Frontend type ("dynamo", direct "vllm", or "vllm-router") nsys_prefix: Optional nsys profiling command prefix dump_config_path: Path to dump config JSON profiling: Profiling config; drives --profiler-config for iteration-based nsys @@ -714,17 +800,50 @@ def build_worker_command( } ) - if frontend_type == "vllm": - if mode != "agg": + if frontend_type in {"vllm", "vllm-router"}: + if frontend_type == "vllm" and mode != "agg": raise ValueError("frontend.type: vllm supports aggregate vLLM jobs only") - if is_multi_node: - raise ValueError("frontend.type: vllm currently supports single-node aggregate jobs only") + is_router_hybrid_dp = ( + frontend_type == "vllm-router" + and is_multi_node + and self._is_dp_mode(mode) + and self.dp_launch_mode == "per_node" + ) + if is_multi_node and not is_router_hybrid_dp: + raise ValueError(f"frontend.type: {frontend_type} requires each vLLM endpoint to fit on one node") config.pop("host", None) config.pop("port", None) - config.pop("connector", None) config.setdefault("served-model-name", served_model_name) + if frontend_type == "vllm": + config.pop("connector", None) + worker_port = runtime.frontend_port + else: + worker_port = process.http_port + mode_connector = config.pop("connector", None) + connector = mode_connector if mode_connector is not None else self.connector + if connector and connector not in ("null", "none", None): + config.setdefault("kv-transfer-config", _connector_to_kv_transfer_config(connector)) + + if is_router_hybrid_dp: + rpc_port_kebab = config.pop("data-parallel-rpc-port", None) + rpc_port_snake = config.pop("data_parallel_rpc_port", None) + dp_rpc_port = process.dp_rpc_port or rpc_port_kebab or rpc_port_snake or VLLM_DATA_PARALLEL_RPC_PORT + + 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) + config.pop("headless", None) + config["data-parallel-size-local"] = self._get_local_dp_size(mode, len(process.gpu_indices)) + config["data-parallel-start-rank"] = process.node_rank + config["data-parallel-address"] = leader_ip + config["data-parallel-rpc-port"] = dp_rpc_port + config["data-parallel-hybrid-lb"] = True + cmd.extend( [ "vllm", @@ -733,7 +852,7 @@ def build_worker_command( "--host", "0.0.0.0", "--port", - str(runtime.frontend_port), + str(worker_port), ] ) if not self.set_cuda_visible_devices: @@ -796,7 +915,7 @@ def build_worker_command( cmd.extend( [ "--data-parallel-size-local", - str(len(process.gpu_indices)), + str(self._get_local_dp_size(mode, len(process.gpu_indices))), "--data-parallel-start-rank", str(process.node_rank), "--data-parallel-address", diff --git a/src/srtctl/benchmarks/router.py b/src/srtctl/benchmarks/router.py index 5d2cfd930..8e15b0a16 100644 --- a/src/srtctl/benchmarks/router.py +++ b/src/srtctl/benchmarks/router.py @@ -43,7 +43,7 @@ def local_script_dir(self) -> str: def validate_config(self, config: SrtConfig) -> list[str]: errors = [] - # Router benchmark requires sglang frontend + # Router benchmark exercises the SGLang router's prefix-aware policies. if config.frontend.type != "sglang": errors.append("router benchmark requires frontend.type: sglang") diff --git a/src/srtctl/cli/do_sweep.py b/src/srtctl/cli/do_sweep.py index 1f846e481..8f385b6f6 100644 --- a/src/srtctl/cli/do_sweep.py +++ b/src/srtctl/cli/do_sweep.py @@ -677,7 +677,7 @@ def run(self) -> int: try: # Stage 1: Head infrastructure (NATS, etcd). Only the dynamo request # plane uses it; static/direct frontends skip it. - if self.config.frontend.type in {"trtllm_serve", "vllm"}: + if self.config.frontend.type in {"sglang", "trtllm_serve", "vllm", "vllm-router"}: logger.info("Skipping NATS/etcd infrastructure (frontend.type=%s)", self.config.frontend.type) else: reporter.report(JobStatus.STARTING, JobStage.HEAD_INFRASTRUCTURE, "Starting head infrastructure") diff --git a/src/srtctl/cli/mixins/benchmark_stage.py b/src/srtctl/cli/mixins/benchmark_stage.py index 24710765c..72085b11d 100644 --- a/src/srtctl/cli/mixins/benchmark_stage.py +++ b/src/srtctl/cli/mixins/benchmark_stage.py @@ -15,7 +15,7 @@ from typing import TYPE_CHECKING from srtctl.core.fingerprint import format_identity_verification, verify_identity -from srtctl.core.health import wait_for_model +from srtctl.core.health import wait_for_http_endpoints, wait_for_model from srtctl.core.lockfile import collect_worker_fingerprints from srtctl.core.slurm import get_hostname_ip, start_srun_process from srtctl.core.status import JobStage, JobStatus, StatusReporter @@ -54,7 +54,7 @@ def _vllm_health_entries( ) -> int: """Return expected Dynamo generate registrations for a vLLM worker mode.""" dp_size = _vllm_data_parallel_size(config, mode) - if dp_size > 1 and getattr(config.backend, "dp_launch_mode", "per_gpu") == "per_node": + if dp_size > 1 and getattr(config.backend, "dp_launch_mode", "per_rank") == "per_node": if backend_processes is None: raise ValueError("backend_processes are required for per-node DP health expectations") endpoint_mode = "agg" if mode == "aggregated" else mode @@ -70,8 +70,9 @@ def _get_health_expectations( Dynamo's /health endpoint reports registered generate instances. For vLLM DP workers, per-GPU launch registers one entry per DP rank, while per-node - launch registers one entry per node-local process. Other frontends keep - using logical worker counts. + launch registers one entry per node-local process. vLLM Router expands each + routed backend URL into its node-local DP ranks. Other frontends keep using + logical worker counts. """ r = config.resources @@ -95,6 +96,24 @@ def _get_health_expectations( 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 + if config.frontend.type == "vllm-router" and backend_processes is not None: + from srtctl.frontends.vllm_router import node_local_data_parallel_size + + local_dp_size = node_local_data_parallel_size(config.backend, backend_processes) + + n_prefill = sum( + local_dp_size + for process in backend_processes + if process.endpoint_mode == "prefill" and process.http_port > 0 + ) + n_decode = sum( + local_dp_size + for process in backend_processes + if process.endpoint_mode in {"decode", "agg"} and process.http_port > 0 + ) + count_desc = f"{n_prefill}P + {n_decode}D Router DP workers; logical workers: {worker_desc}" + return n_prefill, n_decode, count_desc, n_prefill + n_decode + count_desc = worker_desc return logical_prefill, logical_decode, count_desc, logical_prefill + logical_decode @@ -146,11 +165,10 @@ def _benchmark_node(self) -> str: ) def _logical_worker_endpoints(self) -> list[tuple[str, str, int]]: - """Return ``(mode, IP, port)`` for every logical worker leader. + """Return ``(mode, IP, port)`` for every routable worker endpoint. - ``backend_processes`` contains one process per physical node for - multi-node workers. Only rank zero owns the logical worker endpoint, - so follower ranks must not be advertised to benchmark clients. + ``backend_processes`` may contain non-routable TP followers (HTTP port + zero) or multiple node-local vLLM DP pools (one positive port per pool). Dynamo exposes worker metrics on each leader's system port. Other frontends expose them on the worker HTTP port, matching the endpoint @@ -159,7 +177,7 @@ def _logical_worker_endpoints(self) -> list[tuple[str, str, int]]: use_sys_port = self.config.frontend.type == "dynamo" endpoints: list[tuple[str, str, int]] = [] for process in self.backend_processes: - if not process.is_leader: + if use_sys_port and not process.is_leader: continue port = process.sys_port if use_sys_port else process.http_port if port <= 0: @@ -209,6 +227,27 @@ def run_benchmark( reporter.report(JobStatus.FAILED, JobStage.BENCHMARK, "Workers failed health check") return 1 + from srtctl.frontends import get_frontend + + frontend = get_frontend(self.config.frontend.type) + backend_health_urls = frontend.get_backend_health_urls(self.config.backend, self.backend_processes) + if backend_health_urls: + logger.info( + "Frontend requires direct readiness from %d advertised backend URLs", + len(backend_health_urls), + ) + if not wait_for_http_endpoints( + backend_health_urls, + poll_interval=float(hc.interval_seconds), + timeout=float(hc.max_attempts * hc.interval_seconds), + report_every=60.0, + stop_event=stop_event, + ): + logger.error("Advertised backend URLs did not become healthy") + if reporter: + reporter.report(JobStatus.FAILED, JobStage.BENCHMARK, "Backends failed direct health check") + return 1 + logger.info("Server is healthy - starting benchmark") # Identity verification: compare recipe identity against runtime fingerprints @@ -472,11 +511,14 @@ def _get_aiperf_server_metrics_env( logical_endpoints = self._logical_worker_endpoints() urls = [f"http://{host}:{port}/metrics" for _, host, port in logical_endpoints] else: - if self.config.frontend.type == "vllm": + if self.config.frontend.type in {"vllm", "vllm-router"}: for process in self.backend_processes: - if process.endpoint_mode == "agg" and process.is_leader: + if self.config.frontend.type == "vllm" and process.is_leader: host = get_hostname_ip(process.node, self.runtime.network_interface) urls.append(f"http://{host}:{FRONTEND_PUBLIC_PORT}/metrics") + elif self.config.frontend.type == "vllm-router" and process.http_port > 0: + host = get_hostname_ip(process.node, self.runtime.network_interface) + urls.append(f"http://{host}:{process.http_port}/metrics") if urls: return {"AIPERF_SERVER_METRICS_URLS": ",".join(sorted(set(urls)))} diff --git a/src/srtctl/core/__init__.py b/src/srtctl/core/__init__.py index b1276e694..74f4b889b 100644 --- a/src/srtctl/core/__init__.py +++ b/src/srtctl/core/__init__.py @@ -34,9 +34,10 @@ from .health import ( WorkerHealthResult, check_dynamo_health, - check_sglang_router_health, + check_static_router_health, wait_for_etcd, wait_for_health, + wait_for_http_endpoints, wait_for_model, wait_for_port, ) @@ -119,7 +120,7 @@ "WorkerHealthResult", "allocate_endpoints", "check_dynamo_health", - "check_sglang_router_health", + "check_static_router_health", "endpoints_to_processes", "find_cluster_config_path", "get_container_mounts_str", @@ -140,6 +141,7 @@ "start_srun_process", "wait_for_etcd", "wait_for_health", + "wait_for_http_endpoints", "wait_for_model", # Health checks "wait_for_port", diff --git a/src/srtctl/core/config.py b/src/srtctl/core/config.py index 70db894e6..d6e404bea 100755 --- a/src/srtctl/core/config.py +++ b/src/srtctl/core/config.py @@ -171,6 +171,13 @@ def resolve_config_with_defaults(user_config: dict[str, Any], cluster_config: di config["frontend"] = frontend logger.debug(f"Resolved nginx_container alias '{nginx_container}' -> '{resolved_nginx}'") + router_container = frontend.get("container_image", "") + if containers and router_container in containers: + resolved_router = containers[router_container] + frontend["container_image"] = resolved_router + config["frontend"] = frontend + logger.debug(f"Resolved frontend.container_image alias '{router_container}' -> '{resolved_router}'") + # Cluster-level default for nginx nofile ulimit (job yaml wins if present). if "nginx_raise_ulimit" not in frontend and cluster_config.get("nginx_raise_ulimit") is not None: frontend["nginx_raise_ulimit"] = cluster_config["nginx_raise_ulimit"] diff --git a/src/srtctl/core/health.py b/src/srtctl/core/health.py index 35ab7498e..b6bf68f1f 100644 --- a/src/srtctl/core/health.py +++ b/src/srtctl/core/health.py @@ -9,8 +9,9 @@ - wait_for_health(): HTTP health check with worker count validation - wait_for_etcd(): Wait for etcd to be ready - wait_for_model(): Wait for model with worker count validation (replaces bash version) +- wait_for_http_endpoints(): Wait until every adapter-provided HTTP endpoint is ready - check_dynamo_health(): Parse dynamo /health response for worker counts -- check_sglang_router_health(): Parse sglang /workers response for worker counts +- check_static_router_health(): Parse static-router /workers response for worker counts """ import logging @@ -46,12 +47,14 @@ class WorkerHealthResult: # ============================================================================ -def check_sglang_router_health( +# "Static router" means SGLang Router or vLLM Router; both expose worker +# counts through the same /workers stats. +def check_static_router_health( response_json: dict, expected_prefill: int, expected_decode: int, ) -> WorkerHealthResult: - """Check health using sglang router /workers endpoint response. + """Check health using the shared static-router /workers response. Expected response format: { @@ -398,6 +401,60 @@ def wait_for_etcd( # ============================================================================ +def wait_for_http_endpoints( + urls: list[str], + poll_interval: float = 1.0, + timeout: float = 600.0, + report_every: float = 60.0, + stop_event: threading.Event | None = None, +) -> bool: + """Wait until every URL returns HTTP 200 in the same polling pass. + + Frontend adapters use this for direct backend readiness requirements that + are additional to the frontend's own health response. + """ + targets = list(dict.fromkeys(urls)) + if not targets: + return True + + logger.info("Polling %d backend health endpoints every %.1fs", len(targets), poll_interval) + start_time = time.time() + last_report_time = start_time + + while True: + if stop_event and stop_event.is_set(): + logger.warning("Wait for backend health endpoints aborted by stop event") + return False + + if time.time() - start_time >= timeout: + logger.error("Backend health endpoints did not all become ready in %.0f seconds", timeout) + return False + + pending: list[str] = [] + for url in targets: + try: + response = requests.get(url, timeout=5.0) + if response.status_code != 200: + pending.append(url) + except requests.exceptions.RequestException: + pending.append(url) + + if not pending: + logger.info("All %d backend health endpoints are ready", len(targets)) + return True + + if time.time() - last_report_time >= report_every: + logger.info( + "Waiting for %d/%d backend health endpoints: %s", + len(pending), + len(targets), + ", ".join(pending), + ) + last_report_time = time.time() + + time.sleep(poll_interval) + + def wait_for_model( host: str, port: int, @@ -422,20 +479,25 @@ def wait_for_model( poll_interval: Seconds between health checks timeout: Maximum wait time in seconds report_every: Log progress every N seconds - frontend_type: Frontend type - "sglang" uses /workers, "dynamo" uses /health + frontend_type: Registered frontend type; its adapter selects and parses + the appropriate health endpoint. stop_event: Optional threading.Event to abort waiting Returns: True if model is ready with expected workers, False if timeout/aborted """ - if frontend_type == "sglang": - health_url = f"http://{host}:{port}/workers" + from srtctl.frontends import get_frontend + + frontend = get_frontend(frontend_type) + health_url = f"http://{host}:{port}{frontend.health_endpoint}" + if frontend.health_endpoint == "/workers": logger.info( - "Polling %s every %.1fs for %d prefills and %d decodes (sglang frontend)", + "Polling %s every %.1fs for %d prefills and %d decodes (%s frontend)", health_url, poll_interval, n_prefill, n_decode, + frontend_type, ) else: health_url = f"http://{host}:{port}/health" @@ -484,11 +546,7 @@ def wait_for_model( response_json = response.json() - # Check worker counts based on frontend type - if frontend_type == "sglang": - result = check_sglang_router_health(response_json, n_prefill, n_decode) - else: - result = check_dynamo_health(response_json, n_prefill, n_decode) + result = frontend.parse_health(response_json, n_prefill, n_decode) if result.ready: logger.info(result.message) diff --git a/src/srtctl/core/schema.py b/src/srtctl/core/schema.py index f966ba3f6..500b39f0e 100755 --- a/src/srtctl/core/schema.py +++ b/src/srtctl/core/schema.py @@ -1437,7 +1437,8 @@ class FrontendConfig: """Frontend/router configuration. Attributes: - type: Frontend type - "dynamo" (default), "sglang", "trtllm_serve", or "vllm" + type: Frontend type - "dynamo" (default), "sglang", "vllm-router", + "trtllm_serve", or direct "vllm". enable_multiple_frontends: Scale with nginx + multiple routers. When ``True`` (default), srtctl stands up nginx and fans out to ``num_additional_frontends + 1`` router replicas. When @@ -1463,6 +1464,8 @@ class FrontendConfig: carry the session id in that header instead. args: CLI arguments passed to the frontend/router process env: Environment variables for frontend processes + container_image: Optional router-specific container image. Defaults to + the model/backend container when omitted. """ type: str = "dynamo" @@ -1475,6 +1478,7 @@ class FrontendConfig: nginx_keepalive_timeout: str = "600s" args: dict[str, Any] | None = None env: dict[str, str] | None = None + container_image: str | None = None # trtllm_serve orchestrator (ser.yaml) options; ignored by other frontends. ctx_router: dict[str, Any] | None = None # context_servers.router, e.g. {type: conversation} gen_router: dict[str, Any] | None = None # generation_servers.router @@ -1588,6 +1592,8 @@ def __post_init__(self): self._validate_het_jobs() self._validate_trtllm_serve() self._validate_vllm_frontend() + self._validate_static_router_frontend() + self._validate_sglang_data_parallelism() def _validate_trtllm_serve(self): """Catch trtllm_serve misconfigurations at load time (dry-run) instead of @@ -1634,6 +1640,112 @@ def _validate_vllm_frontend(self): raise ValidationError("frontend.type: vllm requires resources.agg_workers >= 1") if (self.resources.agg_nodes or 1) != 1: raise ValidationError("frontend.type: vllm currently supports single-node aggregate jobs only") + try: + dp_size = int(self.backend._get_dp_size("agg") or 1) + if dp_size < 1: + raise ValueError(f"vLLM agg data-parallel-size must be a positive integer; got {dp_size!r}") + model_parallel_size = self.backend._get_model_parallel_size("agg") + except (TypeError, ValueError) as exc: + raise ValidationError(str(exc)) from exc + required_gpus = dp_size * model_parallel_size + if required_gpus != self.resources.gpus_per_agg: + raise ValidationError( + f"direct vLLM parallelism requires DP*TP*PP*PCP={dp_size}*{model_parallel_size}=" + f"{required_gpus} GPUs, but resources allocate {self.resources.gpus_per_agg} GPUs per worker" + ) + + def _validate_static_router_frontend(self): + """Validate native static-router/backend pairings and endpoint shape.""" + required_backend = { + "sglang": "sglang", + "vllm-router": "vllm", + }.get(self.frontend.type) + if required_backend is None: + return + if self.backend_type != required_backend: + raise ValidationError( + f"frontend.type: {self.frontend.type} requires backend.type: {required_backend}; " + f"got {self.backend_type!r}" + ) + + if self.frontend.type == "vllm-router": + endpoint_gpu_counts = { + "prefill": self.resources.gpus_per_prefill if self.resources.num_prefill else 0, + "decode": self.resources.gpus_per_decode if self.resources.num_decode else 0, + "agg": self.resources.gpus_per_agg if self.resources.num_agg else 0, + } + multi_node_modes = [ + mode for mode, count in endpoint_gpu_counts.items() if count > self.resources.gpus_per_node + ] + if multi_node_modes and self.backend.dp_launch_mode != "per_node": + raise ValidationError("multi-node vLLM Router DP endpoints require backend.dp_launch_mode: per_node") + for mode in multi_node_modes: + if not self.backend._is_dp_mode(mode): + raise ValidationError( + f"multi-node vLLM Router {mode} endpoints require data-parallel-size; " + "multi-node TP-only direct serving is not supported" + ) + + local_dp_sizes: dict[str, int] = {} + for mode, gpu_count in endpoint_gpu_counts.items(): + if gpu_count == 0: + continue + try: + dp_size = int(self.backend._get_dp_size(mode) or 1) + if dp_size < 1: + raise ValueError(f"vLLM {mode} data-parallel-size must be a positive integer; got {dp_size!r}") + model_parallel_size = self.backend._get_model_parallel_size(mode) + except (TypeError, ValueError) as exc: + raise ValidationError(str(exc)) from exc + required_gpus = int(dp_size) * model_parallel_size + if required_gpus != gpu_count: + raise ValidationError( + f"vLLM Router {mode} parallelism requires DP*TP*PP*PCP=" + f"{int(dp_size)}*{model_parallel_size}={required_gpus} GPUs, " + f"but resources allocate {gpu_count} GPUs per worker" + ) + local_gpu_count = min(gpu_count, self.resources.gpus_per_node) + try: + local_dp_sizes[mode] = self.backend._get_local_dp_size(mode, local_gpu_count) + except ValueError as exc: + raise ValidationError(str(exc)) from exc + + if len(set(local_dp_sizes.values())) > 1: + sizes = ", ".join(f"{mode}={size}" for mode, size in local_dp_sizes.items()) + raise ValidationError( + "vLLM Router requires the same node-local data-parallel size for every routed backend; " + f"derived {sizes}" + ) + + def _validate_sglang_data_parallelism(self): + """Reject SGLang TP/DP combinations that the server cannot initialize. + + SGLang partitions each tensor-parallel group across its data-parallel + attention ranks, so ``tp_size`` must be divisible by ``dp_size``. Its + CLI otherwise accepts the flags and fails later in ``ServerArgs`` after + the Slurm allocation and container have already started. + """ + if self.backend_type != "sglang": + return + + sglang_cfg = getattr(self.backend, "sglang_config", None) + if sglang_cfg is None: + return + + for mode, mode_cfg in ( + ("prefill", sglang_cfg.prefill), + ("decode", sglang_cfg.decode), + ("aggregated", sglang_cfg.aggregated), + ): + if not mode_cfg: + continue + tp_size = int(mode_cfg.get("tp-size", mode_cfg.get("tp_size", 1))) + dp_size = int(mode_cfg.get("dp-size", mode_cfg.get("dp_size", 1))) + if tp_size % dp_size != 0: + raise ValidationError( + f"sglang_config.{mode}: tp-size={tp_size} must be divisible by " + f"dp-size={dp_size}; SGLang rejects this data-parallel layout" + ) def _validate_het_jobs(self): """When ``resources.het_jobs`` is set to True, enforce supported shape. diff --git a/src/srtctl/core/telemetry.py b/src/srtctl/core/telemetry.py index 12d75391f..79a8974d1 100644 --- a/src/srtctl/core/telemetry.py +++ b/src/srtctl/core/telemetry.py @@ -86,7 +86,12 @@ def generate_telemetry_config( for process in sorted(processes, key=lambda p: (p.endpoint_mode, p.endpoint_index, p.node_rank, p.node)): node_ip = get_hostname_ip(process.node, runtime.network_interface) - port = FRONTEND_PUBLIC_PORT if frontend_type == "vllm" and process.endpoint_mode == "agg" else process.sys_port + if frontend_type == "vllm" and process.endpoint_mode == "agg": + port = FRONTEND_PUBLIC_PORT + elif frontend_type == "vllm-router": + port = process.http_port + else: + port = process.sys_port node_metadata = { "hostname": process.node, "worker_index": str(process.endpoint_index), diff --git a/src/srtctl/frontends/__init__.py b/src/srtctl/frontends/__init__.py index f840b3249..c5b352b26 100644 --- a/src/srtctl/frontends/__init__.py +++ b/src/srtctl/frontends/__init__.py @@ -6,8 +6,9 @@ Supported frontend types: - dynamo: Dynamo frontend with NATS/etcd communication -- sglang: SGLang native router with direct worker connections +- sglang: SGLang Model Gateway with direct worker connections - vllm: Direct vLLM OpenAI server for aggregate jobs +- vllm-router: vLLM Router with direct worker connections """ from srtctl.frontends.base import ( @@ -19,6 +20,7 @@ from srtctl.frontends.sglang import SGLangFrontend from srtctl.frontends.trtllm_serve import TRTLLMServeFrontend from srtctl.frontends.vllm import VLLMFrontend +from srtctl.frontends.vllm_router import VLLMRouterFrontend __all__ = [ "DynamoFrontend", @@ -27,5 +29,6 @@ "SGLangFrontend", "TRTLLMServeFrontend", "VLLMFrontend", + "VLLMRouterFrontend", "get_frontend", ] diff --git a/src/srtctl/frontends/base.py b/src/srtctl/frontends/base.py index b7540c940..f2cb6b405 100644 --- a/src/srtctl/frontends/base.py +++ b/src/srtctl/frontends/base.py @@ -10,8 +10,10 @@ - Building CLI arguments from config """ +import shlex import threading -from typing import TYPE_CHECKING, Any, Literal, Protocol +from collections.abc import Callable +from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeVar if TYPE_CHECKING: from srtctl.core.health import WorkerHealthResult @@ -20,7 +22,45 @@ from srtctl.core.topology import Process # Supported frontend types - extensible by adding new literals -FrontendType = Literal["dynamo", "sglang", "trtllm_serve", "vllm"] +FrontendType = Literal["dynamo", "sglang", "trtllm_serve", "vllm", "vllm-router"] + +FrontendFactory = Callable[[], "FrontendProtocol"] +_FRONTEND_REGISTRY: dict[str, FrontendFactory] = {} +_FrontendClass = TypeVar("_FrontendClass", bound=type) + + +def register_frontend(*names: str) -> Callable[[_FrontendClass], _FrontendClass]: + """Register a frontend implementation under one or more config names.""" + + def decorator(frontend_class: _FrontendClass) -> _FrontendClass: + for name in names: + if name in _FRONTEND_REGISTRY: + raise ValueError(f"Frontend type {name!r} is already registered") + _FRONTEND_REGISTRY[name] = frontend_class + return frontend_class + + return decorator + + +def _load_builtin_frontends() -> None: + """Import built-ins once so their registration decorators run.""" + from srtctl.frontends import dynamo, sglang, trtllm_serve, vllm, vllm_router # noqa: F401 + + +def build_setup_script_preamble(setup_script: str | None) -> str | None: + """Build the standard in-container recipe setup-script invocation.""" + if not setup_script: + return None + script_name = shlex.quote(setup_script) + return ( + f"setup_script={script_name} && " + 'script_path="/configs/${setup_script}" && ' + 'patch_script_path="/configs/patches/${setup_script}" && ' + 'echo "Running setup script: ${script_path} (fallback ${patch_script_path})" && ' + 'if [ -f "${script_path}" ]; then bash "${script_path}"; ' + 'elif [ -f "${patch_script_path}" ]; then bash "${patch_script_path}"; ' + 'else echo "WARNING: ${script_path} or ${patch_script_path} not found"; fi' + ) class FrontendProtocol(Protocol): @@ -51,6 +91,19 @@ def parse_health( """Parse health check response and return worker status.""" ... + def get_backend_health_urls( + self, + backend: Any, + backend_processes: list["Process"], + ) -> list[str]: + """Return backend URLs that must be directly healthy before benchmarking. + + Frontends that discover or gate their own backends return an empty list. + Static adapters may use this hook to require readiness at the exact URLs + they advertise to their router. + """ + ... + def start_frontends( self, topology: Any, # FrontendTopology @@ -93,19 +146,10 @@ def get_frontend(frontend_type: str) -> FrontendProtocol: Raises: ValueError: If frontend type is unknown """ - # Import here to avoid circular imports - from srtctl.frontends.dynamo import DynamoFrontend - from srtctl.frontends.sglang import SGLangFrontend - from srtctl.frontends.trtllm_serve import TRTLLMServeFrontend - from srtctl.frontends.vllm import VLLMFrontend - - if frontend_type == "dynamo": - return DynamoFrontend() - elif frontend_type == "sglang": - return SGLangFrontend() - elif frontend_type == "trtllm_serve": - return TRTLLMServeFrontend() - elif frontend_type == "vllm": - return VLLMFrontend() - else: - raise ValueError(f"Unknown frontend type: {frontend_type!r}. Supported: dynamo, sglang, trtllm_serve, vllm") + _load_builtin_frontends() + try: + factory = _FRONTEND_REGISTRY[frontend_type] + except KeyError as exc: + supported = ", ".join(sorted(_FRONTEND_REGISTRY)) + raise ValueError(f"Unknown frontend type: {frontend_type!r}. Supported: {supported}") from exc + return factory() diff --git a/src/srtctl/frontends/dynamo.py b/src/srtctl/frontends/dynamo.py index 20821a7fe..e1f37c120 100644 --- a/src/srtctl/frontends/dynamo.py +++ b/src/srtctl/frontends/dynamo.py @@ -8,13 +8,13 @@ """ import logging -import shlex import threading from typing import TYPE_CHECKING, Any from srtctl.core.health import WorkerHealthResult, check_dynamo_health from srtctl.core.schema import build_otel_env from srtctl.core.slurm import CONTAINER_REMAP_ROOT_EXPORT, start_srun_process +from srtctl.frontends.base import build_setup_script_preamble, register_frontend from srtctl.ports import ETCD_CLIENT_PORT, NATS_PORT if TYPE_CHECKING: @@ -25,6 +25,7 @@ logger = logging.getLogger(__name__) +@register_frontend("dynamo") class DynamoFrontend: """Dynamo frontend implementation. @@ -49,6 +50,11 @@ def parse_health( """Parse dynamo /health endpoint response.""" return check_dynamo_health(response_json, expected_prefill, expected_decode) + def get_backend_health_urls(self, backend: Any, backend_processes: list["Process"]) -> list[str]: + """Dynamo owns backend discovery and exposes readiness through its frontend.""" + del backend, backend_processes + return [] + def get_frontend_args_list(self, args: dict[str, Any] | None) -> list[str]: """Convert frontend args dict to CLI arguments.""" if not args: @@ -139,18 +145,9 @@ def _build_preamble(self, config: Any) -> str | None: parts = [] # Custom setup script - setup_script = getattr(config, "setup_script", None) - if isinstance(setup_script, str) and setup_script: - script_name = shlex.quote(setup_script) - parts.append( - f"setup_script={script_name} && " - 'script_path="/configs/${setup_script}" && ' - 'patch_script_path="/configs/patches/${setup_script}" && ' - 'echo "Running setup script: ${script_path} (fallback ${patch_script_path})" && ' - 'if [ -f "${script_path}" ]; then bash "${script_path}"; ' - 'elif [ -f "${patch_script_path}" ]; then bash "${patch_script_path}"; ' - 'else echo "WARNING: ${script_path} or ${patch_script_path} not found"; fi' - ) + setup_preamble = build_setup_script_preamble(getattr(config, "setup_script", None)) + if setup_preamble: + parts.append(setup_preamble) # Dynamo installation (required for dynamo frontend) # Skip if dynamo.install is False (container already has dynamo installed) diff --git a/src/srtctl/frontends/sglang.py b/src/srtctl/frontends/sglang.py index a6ec3d4ca..3a6d779cf 100644 --- a/src/srtctl/frontends/sglang.py +++ b/src/srtctl/frontends/sglang.py @@ -1,165 +1,30 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -""" -SGLang router frontend implementation. +"""SGLang Model Gateway router frontend.""" -Uses sglang_router for direct communication with backend workers. -""" +from typing import Any, ClassVar -import logging -import shlex -import threading -from typing import TYPE_CHECKING, Any - -from srtctl.core.health import WorkerHealthResult, check_sglang_router_health from srtctl.core.slurm import get_hostname_ip, start_srun_process +from srtctl.frontends.base import register_frontend +from srtctl.frontends.static_router import StaticRouterFrontend -if TYPE_CHECKING: - from srtctl.core.processes import ManagedProcess - from srtctl.core.runtime import RuntimeContext - from srtctl.core.topology import Process - -logger = logging.getLogger(__name__) - - -class SGLangFrontend: - """SGLang router frontend implementation. - - Uses sglang_router.launch_router for direct worker connections. - Health checks via /workers endpoint. - """ - - @property - def type(self) -> str: - return "sglang" - - @property - def health_endpoint(self) -> str: - return "/workers" - - def parse_health( - self, - response_json: dict, - expected_prefill: int, - expected_decode: int, - ) -> WorkerHealthResult: - """Parse sglang /workers endpoint response.""" - return check_sglang_router_health(response_json, expected_prefill, expected_decode) - - def get_frontend_args_list(self, args: dict[str, Any] | None) -> list[str]: - """Convert frontend args dict to CLI arguments.""" - if not args: - return [] - result = [] - for key, value in args.items(): - if value is True: - result.append(f"--{key}") - elif value is not False and value is not None: - result.extend([f"--{key}", str(value)]) - return result - - def start_frontends( - self, - topology: Any, # FrontendTopology - runtime: "RuntimeContext", - config: Any, # SrtConfig - backend: Any, # BackendProtocol - backend_processes: list["Process"], - stop_event: "threading.Event | None" = None, # unused: returns immediately - ) -> list["ManagedProcess"]: - """Start sglang routers on designated nodes. - - Supports two modes: - - Aggregated: --worker-urls http://w1:port1 http://w2:port2 ... - - Disaggregated: --pd-disaggregation --prefill url bootstrap_port --decode url - """ - from srtctl.backends.sglang import SGLangProtocol - from srtctl.core.processes import ManagedProcess - - r = config.resources - is_disaggregated = r.num_prefill > 0 or r.num_decode > 0 - - # Collect worker info by mode - agg_workers: list[tuple[str, int]] = [] # (ip, http_port) - prefill_leaders: list[tuple[str, int, int | None]] = [] # (ip, http_port, bootstrap_port) - decode_leaders: list[tuple[str, int]] = [] # (ip, http_port) - - # Determine URL schemes based on gRPC mode - prefill_scheme = "http://" - decode_scheme = "http://" - agg_scheme = "http://" - if isinstance(backend, SGLangProtocol): - if backend.is_grpc_mode("prefill"): - prefill_scheme = "grpc://" - if backend.is_grpc_mode("decode"): - decode_scheme = "grpc://" - if backend.is_grpc_mode("agg"): - agg_scheme = "grpc://" - - for process in backend_processes: - if not process.is_leader: - continue - leader_ip = get_hostname_ip(process.node) - if process.endpoint_mode == "agg": - agg_workers.append((leader_ip, process.http_port)) - elif process.endpoint_mode == "prefill": - prefill_leaders.append((leader_ip, process.http_port, process.bootstrap_port)) - elif process.endpoint_mode == "decode": - decode_leaders.append((leader_ip, process.http_port)) - - processes: list[ManagedProcess] = [] - - for idx, node in enumerate(topology.frontend_nodes): - logger.info("Starting sglang-router %d on %s", idx, node) - - router_log = runtime.log_dir / f"{node}_router_{idx}.out" - - cmd = ["python", "-m", "sglang_router.launch_router"] - - if is_disaggregated: - # Disaggregated mode: --pd-disaggregation with --prefill and --decode - cmd.append("--pd-disaggregation") - for ip, http_port, bootstrap_port in prefill_leaders: - cmd.extend(["--prefill", f"{prefill_scheme}{ip}:{http_port}"]) - # Add bootstrap port if available - if bootstrap_port is not None: - cmd.append(str(bootstrap_port)) - for ip, http_port in decode_leaders: - cmd.extend(["--decode", f"{decode_scheme}{ip}:{http_port}"]) - else: - # Aggregated mode: --worker-urls with space-separated URLs - worker_urls = [f"{agg_scheme}{ip}:{port}" for ip, port in agg_workers] - cmd.extend(["--worker-urls"] + worker_urls) - - cmd.extend(["--host", "0.0.0.0", "--port", str(topology.frontend_port)]) - cmd.extend(self.get_frontend_args_list(config.frontend.args)) - logger.info("Router command: %s", shlex.join(cmd)) +@register_frontend("sglang") +class SGLangFrontend(StaticRouterFrontend): + """SGLang Model Gateway static router.""" - # Build env vars - env_to_set: dict[str, str] = {} - if config.frontend.env: - env_to_set.update(config.frontend.env) + type: ClassVar[str] = "sglang" + backend_type: ClassVar[str] = "sglang" + executable: ClassVar[tuple[str, ...]] = ("python", "-m", "sglang_router.launch_router") + pd_flag: ClassVar[str] = "--pd-disaggregation" + process_name: ClassVar[str] = "sglang_router" - proc = start_srun_process( - command=cmd, - nodelist=[node], - output=str(router_log), - container_image=str(runtime.container_image), - container_mounts=runtime.container_mounts, - env_to_set=env_to_set if env_to_set else None, - het_group=runtime.nodes.het_group_for(node), - ) + def worker_scheme(self, backend: Any, mode: str) -> str: + return "grpc" if backend.is_grpc_mode(mode) else "http" - processes.append( - ManagedProcess( - name=f"sglang_router_{idx}", - popen=proc, - log_file=router_log, - node=node, - critical=True, - ) - ) + def get_hostname_ip(self, node: str) -> str: + return get_hostname_ip(node) - return processes + def start_process(self, **kwargs: Any) -> Any: + return start_srun_process(**kwargs) diff --git a/src/srtctl/frontends/static_router.py b/src/srtctl/frontends/static_router.py new file mode 100644 index 000000000..268b7afdb --- /dev/null +++ b/src/srtctl/frontends/static_router.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared implementation for routers configured with static worker URLs.""" + +from __future__ import annotations + +import logging +import shlex +import threading +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, ClassVar + +from srtctl.core.health import WorkerHealthResult, check_static_router_health +from srtctl.core.slurm import get_hostname_ip, start_srun_process +from srtctl.frontends.base import build_setup_script_preamble + +if TYPE_CHECKING: + from srtctl.core.processes import ManagedProcess + from srtctl.core.runtime import RuntimeContext + from srtctl.core.topology import Process + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class RouterWorker: + """A logical backend endpoint exposed to a static router.""" + + mode: str + url: str + bootstrap_port: int | None = None + + +class StaticRouterFrontend: + """Base class for routers whose worker topology is supplied on the CLI.""" + + type: ClassVar[str] + backend_type: ClassVar[str] + executable: ClassVar[tuple[str, ...]] + pd_flag: ClassVar[str] + process_name: ClassVar[str] + + @property + def health_endpoint(self) -> str: + return "/workers" + + def parse_health( + self, + response_json: dict, + expected_prefill: int, + expected_decode: int, + ) -> WorkerHealthResult: + return check_static_router_health(response_json, expected_prefill, expected_decode) + + def get_frontend_args_list(self, args: dict[str, Any] | None) -> list[str]: + """Convert config values to CLI arguments, preserving repeated values.""" + if not args: + return [] + result: list[str] = [] + for key, value in args.items(): + flag = f"--{key}" + if value is True: + result.append(flag) + elif value is False or value is None: + continue + elif isinstance(value, list): + for item in value: + result.extend([flag, str(item)]) + else: + result.extend([flag, str(value)]) + return result + + def get_managed_frontend_args( + self, + config: Any, + backend: Any | None = None, + backend_processes: list[Process] | None = None, + ) -> list[str]: + """Return adapter-managed CLI arguments derived from srtctl config.""" + del backend, backend_processes + return [] + + def get_backend_health_urls(self, backend: Any, backend_processes: list[Process]) -> list[str]: + """Keep existing static-router readiness semantics unless an adapter opts in.""" + del backend, backend_processes + return [] + + def worker_scheme(self, backend: Any, mode: str) -> str: + """Return the protocol used to reach a worker endpoint.""" + return "http" + + def get_hostname_ip(self, node: str) -> str: + """Resolve a worker node to the address advertised to the router.""" + return get_hostname_ip(node) + + def worker_bootstrap_port(self, backend: Any, process: Process) -> int | None: + """Return the optional P/D bootstrap port advertised for a worker.""" + return process.bootstrap_port + + def start_process(self, **kwargs: Any) -> Any: + """Launch one router process. Split out for adapter-specific testing.""" + return start_srun_process(**kwargs) + + def collect_workers(self, backend: Any, backend_processes: list[Process]) -> list[RouterWorker]: + workers: list[RouterWorker] = [] + for process in backend_processes: + # An allocated HTTP port is the source of truth for whether a backend + # process is independently routable. Multi-node vLLM DP exposes one + # node-local pool per URL, while TP follower processes retain port 0. + if process.http_port <= 0: + continue + scheme = self.worker_scheme(backend, process.endpoint_mode) + workers.append( + RouterWorker( + mode=process.endpoint_mode, + url=f"{scheme}://{self.get_hostname_ip(process.node)}:{process.http_port}", + bootstrap_port=self.worker_bootstrap_port(backend, process), + ) + ) + return workers + + def build_router_command(self, workers: list[RouterWorker], host: str, port: int) -> list[str]: + """Build the router CLI for aggregate or prefill/decode topologies.""" + aggregate = [worker for worker in workers if worker.mode == "agg"] + prefills = [worker for worker in workers if worker.mode == "prefill"] + decodes = [worker for worker in workers if worker.mode == "decode"] + + cmd = list(self.executable) + if prefills or decodes: + if aggregate: + raise ValueError("Static router topology cannot mix aggregate and disaggregated workers") + if not prefills or not decodes: + raise ValueError("Disaggregated static router topology requires prefill and decode workers") + cmd.append(self.pd_flag) + for worker in prefills: + cmd.extend(["--prefill", worker.url]) + if worker.bootstrap_port is not None: + cmd.append(str(worker.bootstrap_port)) + for worker in decodes: + cmd.extend(["--decode", worker.url]) + else: + if not aggregate: + raise ValueError("Static router topology has no logical workers") + cmd.extend(["--worker-urls", *(worker.url for worker in aggregate)]) + + cmd.extend(["--host", host, "--port", str(port)]) + return cmd + + def start_frontends( + self, + topology: Any, + runtime: RuntimeContext, + config: Any, + backend: Any, + backend_processes: list[Process], + stop_event: threading.Event | None = None, + ) -> list[ManagedProcess]: + del stop_event # static routers return immediately after launch + from srtctl.core.processes import ManagedProcess + + configured_backend = getattr(getattr(config, "backend", None), "type", self.backend_type) + if configured_backend != self.backend_type: + raise ValueError( + f"frontend.type: {self.type} requires backend.type: {self.backend_type} (got {configured_backend!r})" + ) + + workers = self.collect_workers(backend, backend_processes) + processes: list[ManagedProcess] = [] + for idx, node in enumerate(topology.frontend_nodes): + router_log = runtime.log_dir / f"{node}_{self.type}_{idx}.out" + cmd = self.build_router_command(workers, "0.0.0.0", topology.frontend_port) + cmd.extend(self.get_managed_frontend_args(config, backend, backend_processes)) + cmd.extend(self.get_frontend_args_list(config.frontend.args)) + logger.info("Starting %s %d on %s: %s", self.type, idx, node, shlex.join(cmd)) + + container_image = getattr(config.frontend, "container_image", None) or str(runtime.container_image) + het_group_for = getattr(runtime.nodes, "het_group_for", lambda _node: None) + router_env = dict(getattr(runtime, "environment", {})) + router_env.update(config.frontend.env or {}) + proc = self.start_process( + command=cmd, + nodelist=[node], + output=str(router_log), + container_image=container_image, + container_mounts=runtime.container_mounts, + env_to_set=router_env or None, + bash_preamble=build_setup_script_preamble(getattr(config, "setup_script", None)), + het_group=het_group_for(node), + ) + processes.append( + ManagedProcess( + name=f"{self.process_name}_{idx}", + popen=proc, + log_file=router_log, + node=node, + critical=True, + ) + ) + return processes diff --git a/src/srtctl/frontends/trtllm_serve.py b/src/srtctl/frontends/trtllm_serve.py index f2e24e4c3..6c965a673 100644 --- a/src/srtctl/frontends/trtllm_serve.py +++ b/src/srtctl/frontends/trtllm_serve.py @@ -19,6 +19,7 @@ from srtctl.core.health import WorkerHealthResult, check_trtllm_serve_health, wait_for_health from srtctl.core.slurm import get_hostname_ip, start_srun_process +from srtctl.frontends.base import register_frontend if TYPE_CHECKING: from srtctl.core.processes import ManagedProcess @@ -28,6 +29,7 @@ logger = logging.getLogger(__name__) +@register_frontend("trtllm_serve") class TRTLLMServeFrontend: """trtllm-serve disaggregated frontend. @@ -53,6 +55,11 @@ def parse_health( """Parse trtllm-serve /health response (200 => ready).""" return check_trtllm_serve_health(response_json, expected_prefill, expected_decode) + def get_backend_health_urls(self, backend: Any, backend_processes: list["Process"]) -> list[str]: + """trtllm-serve gates its configured workers before its own health succeeds.""" + del backend, backend_processes + return [] + def get_frontend_args_list(self, args: dict[str, Any] | None) -> list[str]: """Convert frontend args dict to CLI arguments.""" if not args: diff --git a/src/srtctl/frontends/vllm.py b/src/srtctl/frontends/vllm.py index b3f817d75..725b974fc 100644 --- a/src/srtctl/frontends/vllm.py +++ b/src/srtctl/frontends/vllm.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Any from srtctl.core.health import WorkerHealthResult +from srtctl.frontends.base import register_frontend if TYPE_CHECKING: from srtctl.core.processes import ManagedProcess @@ -24,6 +25,7 @@ logger = logging.getLogger(__name__) +@register_frontend("vllm") class VLLMFrontend: """Direct vLLM OpenAI server frontend. @@ -55,6 +57,11 @@ def parse_health( decode_expected=expected_decode, ) + def get_backend_health_urls(self, backend: Any, backend_processes: list[Process]) -> list[str]: + """The direct frontend health endpoint is the vLLM backend itself.""" + del backend, backend_processes + return [] + def get_frontend_args_list(self, args: dict[str, Any] | None) -> list[str]: if not args: return [] diff --git a/src/srtctl/frontends/vllm_router.py b/src/srtctl/frontends/vllm_router.py new file mode 100644 index 000000000..84941ab10 --- /dev/null +++ b/src/srtctl/frontends/vllm_router.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""vLLM Router frontend.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, ClassVar + +from srtctl.frontends.base import register_frontend +from srtctl.frontends.static_router import StaticRouterFrontend + +if TYPE_CHECKING: + from srtctl.core.topology import Process + + +def node_local_data_parallel_size(backend: Any, backend_processes: list[Process]) -> int: + """Return Router's single node-local DP expansion factor.""" + grouped_processes: dict[tuple[str, int], list[Process]] = {} + for process in backend_processes: + if process.http_port > 0: + grouped_processes.setdefault((process.endpoint_mode, process.endpoint_index), []).append(process) + + local_dp_sizes: set[int] = set() + for (mode, _endpoint_index), processes in grouped_processes.items(): + global_dp_size = int(backend._get_dp_size(mode) or 1) + derive_local_dp_size = getattr(type(backend), "_get_local_dp_size", None) + if derive_local_dp_size is None: + # Keep lightweight protocol doubles usable while production vLLMProtocol + # derives this from TP*PP*PCP and the actual local GPU allocation. + if global_dp_size % len(processes) != 0: + raise ValueError( + f"vLLM Router {mode} data-parallel-size={global_dp_size} cannot be evenly split " + f"across {len(processes)} routed servers" + ) + process_local_dp_sizes = {global_dp_size // len(processes)} + else: + process_local_dp_sizes = { + derive_local_dp_size(backend, mode, len(process.gpu_indices)) for process in processes + } + if len(process_local_dp_sizes) != 1: + raise ValueError(f"vLLM Router {mode} endpoint has non-uniform node-local DP sizes") + local_dp_size = next(iter(process_local_dp_sizes)) + if global_dp_size != local_dp_size * len(processes): + raise ValueError( + f"vLLM Router {mode} data-parallel-size={global_dp_size} does not match " + f"{len(processes)} routed servers * local DP {local_dp_size}" + ) + local_dp_sizes.add(local_dp_size) + + if len(local_dp_sizes) > 1: + raise ValueError("vLLM Router requires the same node-local data-parallel size for every routed backend") + return next(iter(local_dp_sizes), 1) + + +@register_frontend("vllm-router") +class VLLMRouterFrontend(StaticRouterFrontend): + """Route requests to direct vLLM OpenAI-compatible worker endpoints.""" + + type: ClassVar[str] = "vllm-router" + backend_type: ClassVar[str] = "vllm" + executable: ClassVar[tuple[str, ...]] = ("vllm-router",) + pd_flag: ClassVar[str] = "--vllm-pd-disaggregation" + process_name: ClassVar[str] = "vllm_router" + + def get_backend_health_urls(self, backend: Any, backend_processes: list[Process]) -> list[str]: + """Return the exact logical vLLM endpoints advertised to Router. + + Router expands node-local DP pools internally, but its worker-count view + can become complete before every advertised HTTP server is accepting + requests. Polling each logical server closes that readiness race without + changing the semantics of other frontend adapters. + """ + return [f"{worker.url.rstrip('/')}/health" for worker in self.collect_workers(backend, backend_processes)] + + def get_managed_frontend_args( + self, + config: Any, + backend: Any | None = None, + backend_processes: list[Process] | None = None, + ) -> list[str]: + """Derive Router DP expansion and worker-readiness arguments.""" + frontend_args = config.frontend.args or {} + managed_args: list[str] = [] + + if backend is not None and backend_processes is not None: + local_dp_size = node_local_data_parallel_size(backend, backend_processes) + configured_dp_size = frontend_args.get( + "intra-node-data-parallel-size", + frontend_args.get("intra_node_data_parallel_size"), + ) + if configured_dp_size is not None and int(configured_dp_size) != local_dp_size: + raise ValueError( + "frontend.args.intra-node-data-parallel-size conflicts with the allocated vLLM topology: " + f"configured {configured_dp_size}, derived {local_dp_size}" + ) + if local_dp_size > 1 and configured_dp_size is None: + managed_args.extend(["--intra-node-data-parallel-size", str(local_dp_size)]) + + if "worker-startup-timeout-secs" not in frontend_args: + health_check = config.health_check + timeout_seconds = health_check.max_attempts * health_check.interval_seconds + managed_args.extend(["--worker-startup-timeout-secs", str(timeout_seconds)]) + return managed_args + + def worker_bootstrap_port(self, backend: Any, process: Process) -> int | None: + """Advertise vLLM's NIXL side-channel port to the P/D router.""" + return process.nixl_port diff --git a/tests/README.md b/tests/README.md index bbac07dd7..7da42e221 100644 --- a/tests/README.md +++ b/tests/README.md @@ -47,7 +47,7 @@ Tests the core `allocate_endpoints` and `endpoints_to_processes` functions: Tests health check response parsing for different backends: - `check_dynamo_health` - Dynamo `/metrics` response parsing -- `check_sglang_router_health` - SGLang `/workers` response parsing +- `check_static_router_health` - SGLang/vLLM Router `/workers` response parsing - Error handling for malformed responses - Aggregated mode (workers count as decode) diff --git a/tests/test_benchmarks.py b/tests/test_benchmarks.py index 386168331..5ea6d12f7 100644 --- a/tests/test_benchmarks.py +++ b/tests/test_benchmarks.py @@ -358,6 +358,28 @@ def test_aggregated_worker_endpoint_uses_http_port_without_dynamo(self): assert "SRT_DECODE_ENDPOINTS" not in env assert env["AIPERF_SERVER_METRICS_URLS"] == "http://ip-node-a:6100/metrics" + def test_vllm_router_exports_each_node_local_dp_backend(self): + """Multi-node DEP8 metrics use two node-local server URLs, not eight rank URLs.""" + from unittest.mock import patch + + from srtctl.benchmarks.custom import CustomBenchmarkRunner + from srtctl.core.topology import Process + + processes = [ + Process("node-a", frozenset(range(4)), 7500, 6100, "agg", 0, node_rank=0), + Process("node-b", frozenset(range(4)), 7501, 6100, "agg", 0, node_rank=4), + ] + stage = self._benchmark_stage("vllm-router", processes) + + with patch( + "srtctl.cli.mixins.benchmark_stage.get_hostname_ip", + side_effect=lambda node, interface: f"ip-{node}", + ): + env = stage._get_benchmark_env(CustomBenchmarkRunner()) + + assert env["SRT_AGG_ENDPOINTS"] == "ip-node-a:6100,ip-node-b:6100" + assert env["AIPERF_SERVER_METRICS_URLS"] == ("http://ip-node-a:6100/metrics,http://ip-node-b:6100/metrics") + def test_worker_endpoint_order_keeps_colocated_logical_workers_aligned(self): from unittest.mock import patch diff --git a/tests/test_configs.py b/tests/test_configs.py index 1443a1ae8..98d5f3248 100644 --- a/tests/test_configs.py +++ b/tests/test_configs.py @@ -541,6 +541,48 @@ def test_grpc_mode_enabled_per_mode(self): assert config.is_grpc_mode("decode") is True assert config.is_grpc_mode("agg") is False + @pytest.mark.parametrize( + ("frontend_type", "mode", "expected"), + [ + ("sglang", "prefill", True), + ("sglang", "decode", True), + ("sglang", "agg", False), + ("dynamo", "decode", False), + ], + ) + def test_static_router_pd_workers_skip_local_fake_bootstrap_warmup( + self, frontend_type: str, mode: str, expected: bool + ) -> None: + """Only native P/D routing relies on the router-level warmup.""" + from unittest.mock import MagicMock, patch + + from srtctl.core.topology import Process + + process = Process( + node="node0", + gpu_indices=frozenset(range(8)), + sys_port=8081, + http_port=6100, + endpoint_mode=mode, + endpoint_index=0, + node_rank=0, + bootstrap_port=7200, + ) + 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"): + command = SGLangProtocol().build_worker_command( + process=process, + endpoint_processes=[process], + runtime=runtime, + frontend_type=frontend_type, + ) + + assert ("--skip-server-warmup" in command) is expected + class TestServedModelName: """Tests for served_model_name property extraction from backend configs.""" @@ -815,6 +857,26 @@ def test_telemetry_container_aliases_resolve(self): assert resolved["telemetry"]["dcgm_exporter"]["container_image"] == "/path/to/dcgm.sqsh" assert resolved["telemetry"]["node_exporter"]["container_image"] == "/path/to/node.sqsh" + def test_router_container_alias_resolves(self): + from srtctl.core.config import resolve_config_with_defaults + + user_config = { + "name": "test", + "model": {"path": "/model", "container": "worker", "precision": "fp8"}, + "resources": {"gpu_type": "h100", "gpus_per_node": 8, "agg_nodes": 1}, + "frontend": {"type": "vllm-router", "container_image": "router"}, + } + cluster_config = { + "containers": { + "worker": "/path/to/worker.sqsh", + "router": "/path/to/router.sqsh", + } + } + + resolved = resolve_config_with_defaults(user_config, cluster_config) + + assert resolved["frontend"]["container_image"] == "/path/to/router.sqsh" + def test_telemetry_literal_paths_pass_through(self): from srtctl.core.config import resolve_config_with_defaults @@ -1891,6 +1953,11 @@ def test_dp_mode_detection(self): assert backend._is_dp_mode("prefill") is False assert backend._is_dp_mode("decode") is False + backend_dp_one = VLLMProtocol( + vllm_config=VLLMServerConfig(prefill={"tensor-parallel-size": 4, "data-parallel-size": 1}) + ) + assert backend_dp_one._is_dp_mode("prefill") is False + # DP mode detected when data-parallel-size is set backend_dp = VLLMProtocol( vllm_config=VLLMServerConfig( @@ -1902,8 +1969,17 @@ 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_dp_size_must_be_positive(self): + """A zero DP size must not silently fall back to vLLM's default.""" + from marshmallow import ValidationError + + from srtctl.backends import VLLMProtocol, VLLMServerConfig + + with pytest.raises(ValidationError, match="must be a positive integer"): + VLLMProtocol(vllm_config=VLLMServerConfig(aggregated={"data-parallel-size": 0})) + + def test_dp_mode_creates_per_rank_processes(self): + """Test that DP mode creates one process per DP rank instead of per node.""" from srtctl.backends import VLLMProtocol, VLLMServerConfig from srtctl.core.topology import Endpoint @@ -1924,7 +2000,7 @@ def test_dp_mode_creates_per_gpu_processes(self): processes = backend.endpoints_to_processes([endpoint]) - # Should create 16 processes (1 per GPU), not 2 (1 per node) + # Should create 16 processes (1 per DP rank), not 2 (1 per node) assert len(processes) == 16 # Each process should have exactly 1 GPU @@ -1947,6 +2023,29 @@ 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_per_gpu_remains_a_compatibility_alias_for_per_rank(self): + """Existing recipes retain the original one-process-per-rank behavior.""" + from srtctl.backends import VLLMProtocol, VLLMServerConfig + from srtctl.core.topology import Endpoint + + config = VLLMServerConfig(aggregated={"data-parallel-size": 2}) + endpoint = Endpoint( + mode="agg", + index=0, + nodes=("node0",), + gpu_indices=frozenset(range(2)), + gpus_per_node=2, + ) + + canonical = VLLMProtocol(dp_launch_mode="per_rank", vllm_config=config) + compatibility = VLLMProtocol(dp_launch_mode="per_gpu", vllm_config=config) + + canonical_processes = canonical.endpoints_to_processes([endpoint]) + compatibility_processes = compatibility.endpoints_to_processes([endpoint]) + + assert [process.node_rank for process in compatibility_processes] == [0, 1] + assert compatibility_processes == canonical_processes + 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 @@ -2034,19 +2133,30 @@ def test_dp_per_node_mode_rejects_dp_size_mismatch(self): gpus_per_node=4, ) - with pytest.raises(ValueError, match="data-parallel-size=7"): + with pytest.raises(ValueError, match=r"DP\*TP\*PP\*PCP=7\*1=7 GPUs"): backend.endpoints_to_processes([endpoint]) - def test_dp_per_node_mode_rejects_headless(self): - """Headless node processes cannot satisfy per-node Dynamo health expectations.""" + @pytest.mark.parametrize( + "managed_field", + [ + "data-parallel-size-local", + "data-parallel-start-rank", + "data-parallel-address", + "data-parallel-rpc-port", + "data-parallel-hybrid-lb", + "headless", + ], + ) + def test_dp_per_node_mode_rejects_manually_managed_topology(self, managed_field): + """Slurm-derived DP topology cannot be contradicted by raw vLLM flags.""" from marshmallow import ValidationError from srtctl.backends import VLLMProtocol, VLLMServerConfig - with pytest.raises(ValidationError, match="remove headless"): + with pytest.raises(ValidationError, match="srt-slurm derives the node-local DP topology"): VLLMProtocol( dp_launch_mode="per_node", - vllm_config=VLLMServerConfig(decode={"data-parallel-size": 8, "headless": True}), + vllm_config=VLLMServerConfig(decode={"data-parallel-size": 8, managed_field: True}), ) def test_direct_vllm_dp_mode_keeps_single_process(self): @@ -2118,6 +2228,257 @@ def test_direct_vllm_command_preserves_current_main_device_binding(self): assert "--request-plane" not in cmd assert "dynamo.vllm" not in cmd + def test_vllm_router_keeps_one_direct_server_per_logical_endpoint(self): + """vLLM Router uses direct private servers rather than Dynamo runtimes.""" + from srtctl.backends import VLLMProtocol + from srtctl.core.topology import Endpoint + + backend = VLLMProtocol() + endpoints = [ + Endpoint( + mode="agg", + index=index, + nodes=(node,), + gpu_indices=frozenset(range(8)), + gpus_per_node=8, + ) + for index, node in enumerate(("node0", "node1")) + ] + + processes = backend.endpoints_to_processes(endpoints, frontend_type="vllm-router") + + assert len(processes) == 2 + assert all(process.is_leader for process in processes) + assert len({process.http_port for process in processes}) == 1 # ports may repeat on distinct nodes + + def test_vllm_router_uses_one_backend_url_for_single_node_dep4(self): + """Router expands one direct backend URL into four node-local DP ranks.""" + from srtctl.backends import VLLMProtocol, VLLMServerConfig + from srtctl.core.topology import Endpoint + + backend = VLLMProtocol( + vllm_config=VLLMServerConfig( + aggregated={"data-parallel-size": 4, "enable-expert-parallel": True}, + ), + ) + endpoint = Endpoint( + mode="agg", + index=0, + nodes=("node0",), + gpu_indices=frozenset(range(4)), + gpus_per_node=4, + ) + + processes = backend.endpoints_to_processes([endpoint], frontend_type="vllm-router") + + assert len(processes) == 1 + assert processes[0].gpu_indices == frozenset(range(4)) + assert processes[0].http_port > 0 + + def test_vllm_router_single_node_tp2_dp2_preserves_native_vllm_args(self): + """Explicit per_node does not inject hybrid flags into a native one-node server.""" + from pathlib import Path + from unittest.mock import MagicMock, patch + + from srtctl.backends import VLLMProtocol, VLLMServerConfig + from srtctl.core.topology import Endpoint + + backend = VLLMProtocol( + dp_launch_mode="per_node", + vllm_config=VLLMServerConfig(aggregated={"tensor-parallel-size": 2, "data-parallel-size": 2}), + ) + endpoint = Endpoint( + mode="agg", + index=0, + nodes=("node0",), + gpu_indices=frozenset(range(4)), + gpus_per_node=4, + ) + processes = backend.endpoints_to_processes([endpoint], frontend_type="vllm-router") + runtime = MagicMock(model_path=Path("/model"), is_hf_model=False, frontend_port=8000) + + with patch("srtctl.core.slurm.get_hostname_ip", return_value="10.0.0.1"): + command = backend.build_worker_command( + process=processes[0], + endpoint_processes=processes, + runtime=runtime, + frontend_type="vllm-router", + ) + + assert command[command.index("--tensor-parallel-size") + 1] == "2" + assert command[command.index("--data-parallel-size") + 1] == "2" + assert "--data-parallel-size-local" not in command + assert "--data-parallel-start-rank" not in command + assert "--data-parallel-hybrid-lb" not in command + + def test_vllm_router_multinode_dep8_uses_hybrid_node_local_pools(self): + """Two DEP8 nodes expose two DP4 HTTP pools sharing one global coordinator.""" + from pathlib import Path + from unittest.mock import MagicMock, patch + + from srtctl.backends import VLLMProtocol, VLLMServerConfig + from srtctl.core.topology import Endpoint + + backend = VLLMProtocol( + dp_launch_mode="per_node", + vllm_config=VLLMServerConfig( + aggregated={ + "data-parallel-size": 8, + "enable-expert-parallel": True, + }, + ), + ) + endpoint = Endpoint( + mode="agg", + index=0, + nodes=("node0", "node1"), + gpu_indices=frozenset(range(4)), + gpus_per_node=4, + ) + processes = backend.endpoints_to_processes([endpoint], frontend_type="vllm-router") + runtime = MagicMock() + runtime.model_path = Path("/model") + runtime.is_hf_model = False + runtime.frontend_port = 8000 + + with patch("srtctl.core.slurm.get_hostname_ip", return_value="10.0.0.1"): + commands = [ + backend.build_worker_command( + process=process, + endpoint_processes=processes, + runtime=runtime, + frontend_type="vllm-router", + ) + for process in processes + ] + + assert len(processes) == 2 + assert [process.node_rank for process in processes] == [0, 4] + assert all(process.http_port > 0 for process in processes) + assert len({process.dp_rpc_port for process in processes}) == 1 + for start_rank, command in zip((0, 4), commands, strict=True): + assert command[command.index("--data-parallel-size") + 1] == "8" + assert command[command.index("--data-parallel-size-local") + 1] == "4" + assert command[command.index("--data-parallel-start-rank") + 1] == str(start_rank) + assert command[command.index("--data-parallel-address") + 1] == "10.0.0.1" + assert "--data-parallel-hybrid-lb" in command + + def test_vllm_router_multinode_tp2_dp4_derives_local_dp2(self): + """Hybrid flags count DP replicas rather than raw GPUs on each node.""" + from pathlib import Path + from unittest.mock import MagicMock, patch + + from srtctl.backends import VLLMProtocol, VLLMServerConfig + from srtctl.core.topology import Endpoint + + backend = VLLMProtocol( + dp_launch_mode="per_node", + vllm_config=VLLMServerConfig(aggregated={"tensor-parallel-size": 2, "data-parallel-size": 4}), + ) + endpoint = Endpoint( + mode="agg", + index=0, + nodes=("node0", "node1"), + gpu_indices=frozenset(range(4)), + gpus_per_node=4, + ) + processes = backend.endpoints_to_processes([endpoint], frontend_type="vllm-router") + runtime = MagicMock(model_path=Path("/model"), is_hf_model=False, frontend_port=8000) + + with patch("srtctl.core.slurm.get_hostname_ip", return_value="10.0.0.1"): + commands = [ + backend.build_worker_command( + process=process, + endpoint_processes=processes, + runtime=runtime, + frontend_type="vllm-router", + ) + for process in processes + ] + + assert [process.node_rank for process in processes] == [0, 2] + assert [command[command.index("--data-parallel-size-local") + 1] for command in commands] == ["2", "2"] + assert [command[command.index("--data-parallel-start-rank") + 1] for command in commands] == ["0", "2"] + + def test_vllm_router_worker_uses_private_port_and_pd_connector(self): + """Disaggregated vLLM Router workers are direct servers with KV transfer.""" + from pathlib import Path + from unittest.mock import MagicMock, patch + + from srtctl.backends import VLLMProtocol, VLLMServerConfig + from srtctl.core.topology import Process + + backend = VLLMProtocol( + connector="nixl", + vllm_config=VLLMServerConfig(prefill={"tensor-parallel-size": 8}), + ) + process = Process( + node="node0", + gpu_indices=frozenset(range(8)), + sys_port=8081, + http_port=30123, + endpoint_mode="prefill", + endpoint_index=0, + node_rank=0, + bootstrap_port=30001, + ) + runtime = MagicMock() + runtime.model_path = Path("/model") + runtime.is_hf_model = False + runtime.frontend_port = 8000 + + with patch("srtctl.core.slurm.get_hostname_ip", return_value="10.0.0.1"): + cmd = backend.build_worker_command( + process=process, + endpoint_processes=[process], + runtime=runtime, + frontend_type="vllm-router", + ) + + assert cmd[:3] == ["vllm", "serve", "/model"] + assert cmd[cmd.index("--port") + 1] == "30123" + assert "dynamo.vllm" not in cmd + kv_config = json.loads(cmd[cmd.index("--kv-transfer-config") + 1]) + assert kv_config == {"kv_connector": "NixlConnector", "kv_role": "kv_both"} + + def test_vllm_router_stable_release_uses_legacy_cuda_binding(self): + """Stable vLLM builds can avoid the newer --device-ids CLI.""" + from pathlib import Path + from unittest.mock import MagicMock, patch + + from srtctl.backends import VLLMProtocol, VLLMServerConfig + from srtctl.core.topology import Process + + backend = VLLMProtocol( + set_cuda_visible_devices=True, + vllm_config=VLLMServerConfig(decode={"tensor-parallel-size": 4}), + ) + process = Process( + node="node0", + gpu_indices=frozenset(range(4)), + sys_port=8081, + http_port=30123, + endpoint_mode="decode", + endpoint_index=0, + node_rank=0, + ) + runtime = MagicMock() + runtime.model_path = Path("/model") + runtime.is_hf_model = False + runtime.frontend_port = 8000 + + with patch("srtctl.core.slurm.get_hostname_ip", return_value="10.0.0.1"): + cmd = backend.build_worker_command( + process=process, + endpoint_processes=[process], + runtime=runtime, + frontend_type="vllm-router", + ) + + assert cmd[:3] == ["vllm", "serve", "/model"] + assert "--device-ids" not in cmd + assert backend.should_set_cuda_visible_devices(process) + def test_direct_vllm_command_keeps_iteration_profiler_config(self): """Direct vllm serve retains main's profiling-derived server option.""" from pathlib import Path @@ -2295,10 +2656,6 @@ def test_dp_per_node_hybrid_command_targets_local_rank_range(self): 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, }, ), @@ -2349,7 +2706,7 @@ def test_dp_per_node_forces_hybrid_lb_for_follower(self): backend = VLLMProtocol( dp_launch_mode="per_node", - vllm_config=VLLMServerConfig(decode={"data-parallel-size": 8, "data_parallel_hybrid_lb": False}), + vllm_config=VLLMServerConfig(decode={"data-parallel-size": 8}), ) leader = Process( node="node0", @@ -2405,7 +2762,7 @@ def test_standard_tp_mode_still_works(self): processes = backend.endpoints_to_processes([endpoint]) - # Should create 2 processes (1 per node), not 16 (1 per GPU) + # Should create 2 processes (1 per node), not 16 (1 per DP rank) assert len(processes) == 2 assert processes[0].node == "node0" assert processes[1].node == "node1" diff --git a/tests/test_frontend_topology.py b/tests/test_frontend_topology.py index 73cc95047..ede820a83 100644 --- a/tests/test_frontend_topology.py +++ b/tests/test_frontend_topology.py @@ -4,6 +4,7 @@ """Tests for frontend topology logic (nginx + multiple frontends).""" from pathlib import Path +from types import SimpleNamespace from unittest.mock import MagicMock, patch from srtctl.cli.do_sweep import SweepOrchestrator @@ -285,8 +286,9 @@ def test_single_node_starts_one_dynamo_frontend(self, mock_mixin_srun, mock_dyna assert processes[0].node == "node0" @patch("srtctl.frontends.sglang.start_srun_process") + @patch("srtctl.frontends.sglang.get_hostname_ip", return_value="10.0.0.1") @patch("srtctl.cli.mixins.frontend_stage.start_srun_process") - def test_single_node_starts_one_sglang_router(self, mock_mixin_srun, mock_sglang_srun): + def test_single_node_starts_one_sglang_router(self, mock_mixin_srun, _mock_ip, mock_sglang_srun): """Single node starts one sglang router, no nginx.""" mock_mixin_srun.return_value = MagicMock() mock_sglang_srun.return_value = MagicMock() @@ -294,7 +296,16 @@ def test_single_node_starts_one_sglang_router(self, mock_mixin_srun, mock_sglang config = make_config(enable_multiple_frontends=True, frontend_type="sglang") runtime = make_runtime(["node0"]) orchestrator = SweepOrchestrator(config=config, runtime=runtime) - orchestrator._backend_processes = [] # No workers for this test + orchestrator.__dict__["backend_processes"] = [ + SimpleNamespace( + is_leader=True, + endpoint_mode=mode, + node="node0", + http_port=30000 + index, + bootstrap_port=30010 if mode == "prefill" else None, + ) + for index, mode in enumerate(("prefill", "decode")) + ] registry = MagicMock() processes = orchestrator.start_frontend(registry) @@ -389,8 +400,9 @@ def test_multi_node_nginx_ulimit_when_opt_in(self, mock_mixin_srun, mock_dynamo_ assert "ulimit -n 1048576" in nginx_cmd[2] @patch("srtctl.frontends.sglang.start_srun_process") + @patch("srtctl.frontends.sglang.get_hostname_ip", return_value="10.0.0.1") @patch("srtctl.cli.mixins.frontend_stage.start_srun_process") - def test_multi_node_sglang_with_nginx(self, mock_mixin_srun, mock_sglang_srun, tmp_path): + def test_multi_node_sglang_with_nginx(self, mock_mixin_srun, _mock_ip, mock_sglang_srun, tmp_path): """Multi-node with sglang router starts nginx + routers.""" mock_mixin_srun.return_value = MagicMock() mock_sglang_srun.return_value = MagicMock() @@ -413,7 +425,16 @@ def test_multi_node_sglang_with_nginx(self, mock_mixin_srun, mock_sglang_srun, t environment=runtime.environment, ) orchestrator = SweepOrchestrator(config=config, runtime=runtime) - orchestrator._backend_processes = [] + orchestrator.__dict__["backend_processes"] = [ + SimpleNamespace( + is_leader=True, + endpoint_mode=mode, + node="node0", + http_port=30000 + index, + bootstrap_port=30010 if mode == "prefill" else None, + ) + for index, mode in enumerate(("prefill", "decode")) + ] registry = MagicMock() processes = orchestrator.start_frontend(registry) diff --git a/tests/test_health.py b/tests/test_health.py index b894104ef..2212e9afe 100644 --- a/tests/test_health.py +++ b/tests/test_health.py @@ -1,14 +1,48 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Tests for health check parsing (Dynamo and SGLang router).""" +"""Tests for health check parsing (Dynamo and static routers).""" + +import threading +from unittest.mock import MagicMock, patch from srtctl.core.health import ( WorkerHealthResult, check_dynamo_health, - check_sglang_router_health, + check_static_router_health, + wait_for_http_endpoints, ) + +def test_wait_for_http_endpoints_keeps_2p2d_blocked_while_one_base_is_unavailable() -> None: + urls = [f"http://{mode}{index}/health" for mode in ("p", "d") for index in range(2)] + responses = [ + MagicMock(status_code=503), + MagicMock(status_code=200), + MagicMock(status_code=200), + MagicMock(status_code=200), + *[MagicMock(status_code=200) for _ in urls], + ] + + with ( + patch("srtctl.core.health.requests.get", side_effect=responses) as get, + patch("srtctl.core.health.time.sleep"), + ): + assert wait_for_http_endpoints(urls, timeout=10.0) + + assert [call.args[0] for call in get.call_args_list] == [*urls, *urls] + + +def test_wait_for_http_endpoints_honors_stop_event() -> None: + stop_event = threading.Event() + stop_event.set() + + with patch("srtctl.core.health.requests.get") as get: + assert not wait_for_http_endpoints(["http://p/health"], stop_event=stop_event) + + get.assert_not_called() + + # ============================================================================ # Dynamo Health Check Tests # ============================================================================ @@ -203,7 +237,7 @@ def test_all_workers_ready_realistic_format(self): }, } - result = check_sglang_router_health(response, expected_prefill=1, expected_decode=2) + result = check_static_router_health(response, expected_prefill=1, expected_decode=2) assert result.ready is True assert result.prefill_ready == 1 @@ -222,7 +256,7 @@ def test_all_workers_ready(self): }, } - result = check_sglang_router_health(response, expected_prefill=4, expected_decode=8) + result = check_static_router_health(response, expected_prefill=4, expected_decode=8) assert result.ready is True assert result.prefill_ready == 4 @@ -241,7 +275,7 @@ def test_more_workers_than_expected(self): }, } - result = check_sglang_router_health(response, expected_prefill=4, expected_decode=8) + result = check_static_router_health(response, expected_prefill=4, expected_decode=8) assert result.ready is True assert result.prefill_ready == 6 @@ -259,7 +293,7 @@ def test_missing_prefill_workers(self): }, } - result = check_sglang_router_health(response, expected_prefill=4, expected_decode=8) + result = check_static_router_health(response, expected_prefill=4, expected_decode=8) assert result.ready is False assert result.prefill_ready == 2 @@ -278,7 +312,7 @@ def test_missing_decode_workers(self): }, } - result = check_sglang_router_health(response, expected_prefill=4, expected_decode=8) + result = check_static_router_health(response, expected_prefill=4, expected_decode=8) assert result.ready is False assert result.decode_ready == 3 @@ -297,7 +331,7 @@ def test_zero_workers(self): }, } - result = check_sglang_router_health(response, expected_prefill=2, expected_decode=4) + result = check_static_router_health(response, expected_prefill=2, expected_decode=4) assert result.ready is False assert result.prefill_ready == 0 @@ -323,7 +357,7 @@ def test_regular_workers_count_as_decode(self): } # Aggregated: expect 0 prefill, 4 decode (regular counts as decode) - result = check_sglang_router_health(response, expected_prefill=0, expected_decode=4) + result = check_static_router_health(response, expected_prefill=0, expected_decode=4) assert result.ready is True assert result.decode_ready == 4 @@ -342,7 +376,7 @@ def test_aggregated_workers_as_decode(self): } # Aggregated: expect 0 prefill, N decode - result = check_sglang_router_health(response, expected_prefill=0, expected_decode=4) + result = check_static_router_health(response, expected_prefill=0, expected_decode=4) assert result.ready is True assert result.decode_ready == 4 @@ -360,7 +394,7 @@ def test_mixed_decode_and_regular(self): } # Both decode and regular should count - result = check_sglang_router_health(response, expected_prefill=0, expected_decode=4) + result = check_static_router_health(response, expected_prefill=0, expected_decode=4) assert result.ready is True assert result.decode_ready == 4 # 2 decode + 2 regular @@ -373,7 +407,7 @@ def test_missing_stats_key(self): """Response missing 'stats' key.""" response = {"workers": []} - result = check_sglang_router_health(response, expected_prefill=1, expected_decode=1) + result = check_static_router_health(response, expected_prefill=1, expected_decode=1) assert result.ready is False assert "stats" in result.message @@ -382,7 +416,7 @@ def test_empty_response(self): """Empty response dict.""" response = {} - result = check_sglang_router_health(response, expected_prefill=1, expected_decode=1) + result = check_static_router_health(response, expected_prefill=1, expected_decode=1) assert result.ready is False @@ -390,7 +424,7 @@ def test_missing_count_fields_defaults_to_zero(self): """Missing count fields default to 0.""" response = {"stats": {}} - result = check_sglang_router_health(response, expected_prefill=1, expected_decode=1) + result = check_static_router_health(response, expected_prefill=1, expected_decode=1) assert result.ready is False assert result.prefill_ready == 0 diff --git a/tests/test_health_expectations.py b/tests/test_health_expectations.py index 11fcfb23b..a8994abcd 100644 --- a/tests/test_health_expectations.py +++ b/tests/test_health_expectations.py @@ -16,10 +16,29 @@ def _config( num_decode=0, num_agg=0, vllm_config=None, - dp_launch_mode="per_gpu", + dp_launch_mode="per_rank", ): """Build a duck-typed stand-in for SrtConfig with only the fields the helpers read.""" - backend = SimpleNamespace(type=backend_type, vllm_config=vllm_config, dp_launch_mode=dp_launch_mode) + + def is_dp_mode(mode): + mode_name = "aggregated" if mode == "agg" else mode + mode_config = getattr(vllm_config, mode_name, None) if vllm_config else None + return bool(mode_config and (mode_config.get("data-parallel-size") or mode_config.get("data_parallel_size"))) + + def get_dp_size(mode): + mode_name = "aggregated" if mode == "agg" else mode + mode_config = getattr(vllm_config, mode_name, None) if vllm_config else None + if not mode_config: + return None + return mode_config.get("data-parallel-size") or mode_config.get("data_parallel_size") + + backend = SimpleNamespace( + type=backend_type, + vllm_config=vllm_config, + dp_launch_mode=dp_launch_mode, + _is_dp_mode=is_dp_mode, + _get_dp_size=get_dp_size, + ) return SimpleNamespace( frontend=SimpleNamespace(type=frontend_type), backend=backend, @@ -27,12 +46,36 @@ def _config( ) -def _processes(*, prefill=0, decode=0, agg=0): +def _processes(*, prefill=0, decode=0, agg=0, gpus_per_process=1): """Build backend-process stand-ins grouped by endpoint mode.""" return [ - *(SimpleNamespace(endpoint_mode="prefill") for _ in range(prefill)), - *(SimpleNamespace(endpoint_mode="decode") for _ in range(decode)), - *(SimpleNamespace(endpoint_mode="agg") for _ in range(agg)), + *( + SimpleNamespace( + endpoint_mode="prefill", + endpoint_index=index, + http_port=6100, + gpu_indices=frozenset(range(gpus_per_process)), + ) + for index in range(prefill) + ), + *( + SimpleNamespace( + endpoint_mode="decode", + endpoint_index=index, + http_port=6100, + gpu_indices=frozenset(range(gpus_per_process)), + ) + for index in range(decode) + ), + *( + SimpleNamespace( + endpoint_mode="agg", + endpoint_index=index, + http_port=6100, + gpu_indices=frozenset(range(gpus_per_process)), + ) + for index in range(agg) + ), ] @@ -122,6 +165,24 @@ def test_non_dynamo_frontend_uses_logical_worker_counts(): assert count_desc == "6P + 1D" +def test_vllm_router_counts_dp_workers_expanded_from_backend_urls(): + """Router health waits for four ranks behind each of one P and two D URLs.""" + vllm_config = SimpleNamespace( + prefill={"data-parallel-size": 4}, + decode={"data-parallel-size": 4}, + aggregated=None, + ) + config = _config("vllm-router", "vllm", num_prefill=1, num_decode=2, vllm_config=vllm_config) + + n_prefill, n_decode, count_desc, num_workers = _get_health_expectations( + config, + _processes(prefill=1, decode=2, gpus_per_process=4), + ) + + assert (n_prefill, n_decode, num_workers) == (4, 8, 12) + assert count_desc == "4P + 8D Router DP workers; logical workers: 1P + 2D" + + def test_dynamo_non_vllm_backend_uses_logical_worker_counts(): """Dynamo + sglang has no DP-rank fan-out in these units; stay logical.""" config = _config("dynamo", "sglang", num_prefill=6, num_decode=1) diff --git a/tests/test_static_router_frontends.py b/tests/test_static_router_frontends.py new file mode 100644 index 000000000..15a25f948 --- /dev/null +++ b/tests/test_static_router_frontends.py @@ -0,0 +1,536 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for native static-router frontend adapters.""" + +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from srtctl.frontends import SGLangFrontend, VLLMRouterFrontend, get_frontend +from srtctl.frontends.static_router import RouterWorker + + +def test_registry_exposes_native_router_names() -> None: + assert isinstance(get_frontend("sglang"), SGLangFrontend) + assert isinstance(get_frontend("vllm-router"), VLLMRouterFrontend) + + +@pytest.mark.parametrize("frontend", [SGLangFrontend(), VLLMRouterFrontend()]) +def test_aggregate_command_advertises_all_logical_workers(frontend) -> None: + command = frontend.build_router_command( + [ + RouterWorker("agg", "http://10.0.0.1:30000"), + RouterWorker("agg", "http://10.0.0.2:30000"), + ], + "0.0.0.0", + 8000, + ) + + assert command[-4:] == ["--host", "0.0.0.0", "--port", "8000"] + worker_urls = command[command.index("--worker-urls") + 1 : -4] + assert worker_urls == ["http://10.0.0.1:30000", "http://10.0.0.2:30000"] + + +@pytest.mark.parametrize( + ("frontend", "pd_flag"), + [ + (SGLangFrontend(), "--pd-disaggregation"), + (VLLMRouterFrontend(), "--vllm-pd-disaggregation"), + ], +) +def test_disaggregated_command_preserves_modes_and_bootstrap(frontend, pd_flag: str) -> None: + command = frontend.build_router_command( + [ + RouterWorker("prefill", "http://10.0.0.1:30000", 30001), + RouterWorker("decode", "http://10.0.0.2:30000"), + ], + "0.0.0.0", + 8000, + ) + + assert pd_flag in command + assert command[command.index("--prefill") + 1 : command.index("--decode")] == [ + "http://10.0.0.1:30000", + "30001", + ] + assert command[command.index("--decode") + 1] == "http://10.0.0.2:30000" + + +def test_router_command_rejects_incomplete_or_mixed_topology() -> None: + frontend = VLLMRouterFrontend() + with pytest.raises(ValueError, match="requires prefill and decode"): + frontend.build_router_command([RouterWorker("prefill", "http://p:1")], "0.0.0.0", 8000) + with pytest.raises(ValueError, match="cannot mix"): + frontend.build_router_command( + [ + RouterWorker("agg", "http://a:1"), + RouterWorker("prefill", "http://p:1"), + RouterWorker("decode", "http://d:1"), + ], + "0.0.0.0", + 8000, + ) + + +def test_frontend_args_repeat_list_values() -> None: + frontend = VLLMRouterFrontend() + assert frontend.get_frontend_args_list({"routing-logic": ["round_robin", "session"]}) == [ + "--routing-logic", + "round_robin", + "--routing-logic", + "session", + ] + + +def test_vllm_router_advertises_nixl_side_channel_port() -> None: + frontend = VLLMRouterFrontend() + process = SimpleNamespace( + is_leader=True, + endpoint_mode="prefill", + node="node1", + http_port=30000, + bootstrap_port=12000, + nixl_port=13000, + ) + + with patch.object(frontend, "get_hostname_ip", return_value="10.0.0.1"): + workers = frontend.collect_workers(MagicMock(), [process]) + + assert workers == [RouterWorker("prefill", "http://10.0.0.1:30000", 13000)] + + +def test_vllm_router_health_gates_every_advertised_2p2d_logical_worker() -> None: + frontend = VLLMRouterFrontend() + backend = MagicMock() + processes = [ + SimpleNamespace(endpoint_mode="prefill", node="p0", http_port=6100, nixl_port=5400), + SimpleNamespace(endpoint_mode="prefill", node="p1", http_port=6100, nixl_port=5401), + SimpleNamespace(endpoint_mode="decode", node="d0", http_port=6100, nixl_port=5500), + SimpleNamespace(endpoint_mode="decode", node="d1", http_port=6100, nixl_port=5501), + SimpleNamespace(endpoint_mode="decode", node="tp-follower", http_port=0, nixl_port=5502), + ] + + with patch.object(frontend, "get_hostname_ip", side_effect=lambda node: f"ip-{node}"): + urls = frontend.get_backend_health_urls(backend, processes) + + assert urls == [ + "http://ip-p0:6100/health", + "http://ip-p1:6100/health", + "http://ip-d0:6100/health", + "http://ip-d1:6100/health", + ] + + +def test_vllm_router_derives_dep4_expansion_for_1p2d() -> None: + """One P URL and two D URLs are each expanded to four ranks by Router.""" + frontend = VLLMRouterFrontend() + backend = MagicMock() + backend._is_dp_mode.return_value = True + backend._get_dp_size.return_value = 4 + processes = [ + SimpleNamespace( + endpoint_mode="prefill", + endpoint_index=0, + node="prefill", + gpu_indices=frozenset(range(4)), + http_port=6100, + nixl_port=5400, + node_rank=0, + ), + SimpleNamespace( + endpoint_mode="decode", + endpoint_index=0, + node="decode0", + gpu_indices=frozenset(range(4)), + http_port=6100, + nixl_port=5500, + node_rank=0, + ), + SimpleNamespace( + endpoint_mode="decode", + endpoint_index=1, + node="decode1", + gpu_indices=frozenset(range(4)), + http_port=6100, + nixl_port=5504, + node_rank=0, + ), + ] + config = SimpleNamespace( + frontend=SimpleNamespace(args={}), + health_check=SimpleNamespace(max_attempts=360, interval_seconds=10), + ) + + with patch.object(frontend, "get_hostname_ip", side_effect=lambda node: f"ip-{node}"): + workers = frontend.collect_workers(backend, processes) + command = frontend.build_router_command(workers, "0.0.0.0", 8000) + + assert len([worker for worker in workers if worker.mode == "prefill"]) == 1 + assert len([worker for worker in workers if worker.mode == "decode"]) == 2 + assert command.count("--prefill") == 1 + assert command.count("--decode") == 2 + assert frontend.get_managed_frontend_args(config, backend, processes) == [ + "--intra-node-data-parallel-size", + "4", + "--worker-startup-timeout-secs", + "3600", + ] + + +def test_vllm_router_launch_uses_router_container_env_and_only_leaders() -> None: + frontend = VLLMRouterFrontend() + runtime = SimpleNamespace( + log_dir=Path("/logs"), + container_image=Path("/worker.sqsh"), + container_mounts={"/host": "/container"}, + environment={"GLOBAL": "value", "ROUTER_LOG": "info"}, + nodes=SimpleNamespace(het_group_for=lambda node: 1), + ) + config = SimpleNamespace( + backend=SimpleNamespace(type="vllm"), + health_check=SimpleNamespace(max_attempts=360, interval_seconds=10), + frontend=SimpleNamespace( + args={"routing-logic": "session"}, + env={"ROUTER_LOG": "debug"}, + container_image="docker://router:test", + ), + setup_script="router-deps.sh", + ) + topology = SimpleNamespace(frontend_nodes=["node0"], frontend_port=8180) + workers = [ + SimpleNamespace( + is_leader=True, + endpoint_mode="agg", + endpoint_index=0, + node="node1", + gpu_indices=frozenset(range(8)), + http_port=30000, + bootstrap_port=None, + nixl_port=None, + ), + SimpleNamespace( + is_leader=False, + endpoint_mode="agg", + endpoint_index=0, + node="node2", + gpu_indices=frozenset(range(8)), + http_port=0, + bootstrap_port=None, + nixl_port=None, + ), + ] + + backend = MagicMock() + backend._is_dp_mode.return_value = False + backend._get_dp_size.return_value = None + + with ( + patch.object(frontend, "get_hostname_ip", return_value="10.0.0.1"), + patch.object(frontend, "start_process", return_value=MagicMock()) as start, + ): + processes = frontend.start_frontends(topology, runtime, config, backend, workers) + + kwargs = start.call_args.kwargs + assert kwargs["output"] == "/logs/node0_vllm-router_0.out" + assert kwargs["container_image"] == "docker://router:test" + assert kwargs["env_to_set"] == {"GLOBAL": "value", "ROUTER_LOG": "debug"} + assert kwargs["het_group"] == 1 + assert "/configs/${setup_script}" in kwargs["bash_preamble"] + assert kwargs["command"].count("http://10.0.0.1:30000") == 1 + assert "--routing-logic" in kwargs["command"] + timeout_index = kwargs["command"].index("--worker-startup-timeout-secs") + assert kwargs["command"][timeout_index + 1] == "3600" + assert processes[0].log_file == Path("/logs/node0_vllm-router_0.out") + + +def test_vllm_router_explicit_worker_startup_timeout_overrides_managed_value() -> None: + frontend = VLLMRouterFrontend() + config = SimpleNamespace( + health_check=SimpleNamespace(max_attempts=360, interval_seconds=10), + frontend=SimpleNamespace(args={"worker-startup-timeout-secs": 7200}), + ) + + command = [ + *frontend.get_managed_frontend_args(config), + *frontend.get_frontend_args_list(config.frontend.args), + ] + + assert command == ["--worker-startup-timeout-secs", "7200"] + + +def test_router_rejects_backend_mismatch_before_launch() -> None: + frontend = VLLMRouterFrontend() + config = SimpleNamespace( + backend=SimpleNamespace(type="sglang"), + frontend=SimpleNamespace(args=None, env=None, container_image=None), + ) + topology = SimpleNamespace(frontend_nodes=["node0"], frontend_port=8180) + runtime = SimpleNamespace(log_dir=Path("/logs"), container_image=Path("/worker.sqsh")) + + with pytest.raises(ValueError, match="requires backend.type: vllm"): + frontend.start_frontends(topology, runtime, config, MagicMock(), []) + + +def test_schema_rejects_router_backend_mismatch() -> None: + from marshmallow import ValidationError + + from srtctl.backends import SGLangProtocol + from srtctl.core.schema import FrontendConfig, ResourceConfig, SrtConfig + + with pytest.raises(ValidationError, match="vllm-router requires backend.type: vllm"): + SrtConfig( + name="bad-router-pair", + model={"path": "model", "container": "image", "precision": "fp8"}, + resources=ResourceConfig(gpu_type="h100", gpus_per_node=8, agg_nodes=1, agg_workers=1), + frontend=FrontendConfig(type="vllm-router", enable_multiple_frontends=False), + backend=SGLangProtocol(), + ) + + +def test_vllm_router_accepts_many_single_node_endpoints() -> None: + from srtctl.backends import VLLMProtocol, VLLMServerConfig + from srtctl.core.schema import FrontendConfig, ResourceConfig, SrtConfig + + config = SrtConfig( + name="multi-endpoint-router", + model={"path": "model", "container": "image", "precision": "fp8"}, + resources=ResourceConfig( + gpu_type="h100", + gpus_per_node=8, + agg_nodes=4, + agg_workers=4, + ), + frontend=FrontendConfig(type="vllm-router", enable_multiple_frontends=False), + backend=VLLMProtocol(vllm_config=VLLMServerConfig(aggregated={"tensor-parallel-size": 8})), + ) + + assert config.resources.gpus_per_agg == 8 + + +def test_vllm_router_rejects_multinode_tp_only_endpoint() -> None: + from marshmallow import ValidationError + + from srtctl.backends import VLLMProtocol + from srtctl.core.schema import FrontendConfig, ResourceConfig, SrtConfig + + with pytest.raises(ValidationError, match="multi-node TP-only"): + SrtConfig( + name="multi-node-endpoint", + model={"path": "model", "container": "image", "precision": "fp8"}, + resources=ResourceConfig( + gpu_type="h100", + gpus_per_node=8, + prefill_nodes=2, + prefill_workers=1, + decode_nodes=1, + decode_workers=1, + ), + frontend=FrontendConfig(type="vllm-router", enable_multiple_frontends=False), + backend=VLLMProtocol(dp_launch_mode="per_node"), + ) + + +def test_vllm_router_accepts_multinode_dep8_endpoint() -> None: + from srtctl.backends import VLLMProtocol, VLLMServerConfig + from srtctl.core.schema import FrontendConfig, ResourceConfig, SrtConfig + + config = SrtConfig( + name="multi-node-dep8", + model={"path": "model", "container": "image", "precision": "fp8"}, + resources=ResourceConfig( + gpu_type="gb200", + gpus_per_node=4, + agg_nodes=2, + agg_workers=1, + ), + frontend=FrontendConfig(type="vllm-router", enable_multiple_frontends=False), + backend=VLLMProtocol( + dp_launch_mode="per_node", + vllm_config=VLLMServerConfig( + aggregated={"data-parallel-size": 8, "enable-expert-parallel": True}, + ), + ), + ) + + assert config.resources.gpus_per_agg == 8 + + +def test_vllm_router_accepts_native_single_node_tp2_dp2() -> None: + """A single vllm serve owns its complete native TP x DP topology.""" + from srtctl.backends import VLLMProtocol, VLLMServerConfig + from srtctl.core.schema import FrontendConfig, ResourceConfig, SrtConfig + + config = SrtConfig( + name="native-tp2-dp2", + model={"path": "model", "container": "image", "precision": "fp8"}, + resources=ResourceConfig( + gpu_type="gb200", + gpus_per_node=4, + agg_nodes=1, + agg_workers=1, + ), + frontend=FrontendConfig(type="vllm-router", enable_multiple_frontends=False), + backend=VLLMProtocol( + vllm_config=VLLMServerConfig( + aggregated={ + "tensor-parallel-size": 2, + "data-parallel-size": 2, + "enable-expert-parallel": True, + } + ), + ), + ) + + assert config.backend._get_model_parallel_size("agg") == 2 + assert config.backend._get_local_dp_size("agg", 4) == 2 + + +def test_vllm_router_rejects_parallelism_allocation_mismatch() -> None: + """Reject a recipe before Slurm when vLLM cannot consume its GPU allocation.""" + from marshmallow import ValidationError + + from srtctl.backends import VLLMProtocol, VLLMServerConfig + from srtctl.core.schema import FrontendConfig, ResourceConfig, SrtConfig + + with pytest.raises(ValidationError, match=r"DP\*TP\*PP\*PCP=2\*2=4 GPUs.*allocate 8 GPUs"): + SrtConfig( + name="invalid-native-world", + model={"path": "model", "container": "image", "precision": "fp8"}, + resources=ResourceConfig( + gpu_type="h100", + gpus_per_node=8, + agg_nodes=1, + agg_workers=1, + ), + frontend=FrontendConfig(type="vllm-router", enable_multiple_frontends=False), + backend=VLLMProtocol( + vllm_config=VLLMServerConfig(aggregated={"tensor-parallel-size": 2, "data-parallel-size": 2}) + ), + ) + + +def test_direct_vllm_validates_native_parallelism_allocation() -> None: + """Direct vLLM uses the same native world-size invariant as Router workers.""" + from marshmallow import ValidationError + + from srtctl.backends import VLLMProtocol, VLLMServerConfig + from srtctl.core.schema import FrontendConfig, ResourceConfig, SrtConfig + + valid = SrtConfig( + name="direct-native-tp2-dp2", + model={"path": "model", "container": "image", "precision": "fp8"}, + resources=ResourceConfig(gpu_type="gb200", gpus_per_node=4, agg_nodes=1, agg_workers=1), + frontend=FrontendConfig(type="vllm", enable_multiple_frontends=False), + backend=VLLMProtocol( + vllm_config=VLLMServerConfig(aggregated={"tensor-parallel-size": 2, "data-parallel-size": 2}) + ), + ) + assert valid.resources.gpus_per_agg == 4 + + with pytest.raises(ValidationError, match=r"direct vLLM parallelism requires.*4 GPUs.*allocate 8 GPUs"): + SrtConfig( + name="invalid-direct-native-world", + model={"path": "model", "container": "image", "precision": "fp8"}, + resources=ResourceConfig(gpu_type="h100", gpus_per_node=8, agg_nodes=1, agg_workers=1), + frontend=FrontendConfig(type="vllm", enable_multiple_frontends=False), + backend=VLLMProtocol( + vllm_config=VLLMServerConfig(aggregated={"tensor-parallel-size": 2, "data-parallel-size": 2}) + ), + ) + + +def test_vllm_router_accepts_multinode_tp2_dp4_endpoint() -> None: + """Hybrid mode derives two local DP replicas per four-GPU node.""" + from srtctl.backends import VLLMProtocol, VLLMServerConfig + from srtctl.core.schema import FrontendConfig, ResourceConfig, SrtConfig + + config = SrtConfig( + name="multi-node-tp2-dp4", + model={"path": "model", "container": "image", "precision": "fp8"}, + resources=ResourceConfig( + gpu_type="gb200", + gpus_per_node=4, + agg_nodes=2, + agg_workers=1, + ), + frontend=FrontendConfig(type="vllm-router", enable_multiple_frontends=False), + backend=VLLMProtocol( + dp_launch_mode="per_node", + vllm_config=VLLMServerConfig( + aggregated={ + "tensor-parallel-size": 2, + "data-parallel-size": 4, + "enable-expert-parallel": True, + } + ), + ), + ) + + assert config.backend._get_local_dp_size("agg", 4) == 2 + + +def test_vllm_router_rejects_different_local_dp_sizes_across_pd_pools() -> None: + """Router has one DP expansion factor, so every advertised pool must agree.""" + from marshmallow import ValidationError + + from srtctl.backends import VLLMProtocol, VLLMServerConfig + from srtctl.core.schema import FrontendConfig, ResourceConfig, SrtConfig + + with pytest.raises(ValidationError, match="same node-local data-parallel size.*prefill=4, decode=2"): + SrtConfig( + name="mismatched-pd-local-dp", + model={"path": "model", "container": "image", "precision": "fp8"}, + resources=ResourceConfig( + gpu_type="gb200", + gpus_per_node=4, + prefill_nodes=1, + prefill_workers=1, + decode_nodes=1, + decode_workers=1, + ), + frontend=FrontendConfig(type="vllm-router", enable_multiple_frontends=False), + backend=VLLMProtocol( + vllm_config=VLLMServerConfig( + prefill={"tensor-parallel-size": 1, "data-parallel-size": 4}, + decode={"tensor-parallel-size": 2, "data-parallel-size": 2}, + ) + ), + ) + + +def test_sgl_router_rejects_non_divisible_tp_dp_layout() -> None: + from marshmallow import ValidationError + + from srtctl.backends import SGLangProtocol, SGLangServerConfig + from srtctl.core.schema import FrontendConfig, ResourceConfig, SrtConfig + + with pytest.raises(ValidationError, match="tp-size=1 must be divisible by dp-size=8"): + SrtConfig( + name="invalid-sglang-dpa", + model={"path": "model", "container": "image", "precision": "fp8"}, + resources=ResourceConfig(gpu_type="h100", gpus_per_node=8, agg_nodes=1, agg_workers=1), + frontend=FrontendConfig(type="sglang", enable_multiple_frontends=False), + backend=SGLangProtocol( + sglang_config=SGLangServerConfig(aggregated={"tp-size": 1, "dp-size": 8, "enable-dp-attention": True}) + ), + ) + + +def test_sgl_router_accepts_divisible_tp_dp_layout() -> None: + from srtctl.backends import SGLangProtocol, SGLangServerConfig + from srtctl.core.schema import FrontendConfig, ResourceConfig, SrtConfig + + config = SrtConfig( + name="valid-sglang-dpa", + model={"path": "model", "container": "image", "precision": "fp8"}, + resources=ResourceConfig(gpu_type="h100", gpus_per_node=8, agg_nodes=1, agg_workers=1), + frontend=FrontendConfig(type="sglang", enable_multiple_frontends=False), + backend=SGLangProtocol( + sglang_config=SGLangServerConfig(aggregated={"tp-size": 8, "dp-size": 8, "enable-dp-attention": True}) + ), + ) + + assert config.backend.sglang_config.aggregated["tp-size"] == 8