From 05a3cec4de67b7529bf896168b770bffeda81efe Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Tue, 22 Sep 2026 13:08:25 -0500 Subject: [PATCH 1/3] feat(atom): integrate native ATOM and AToMesh --- CLAUDE.md | 2 + README.md | 2 +- docs/config-reference.md | 13 + docs/schema-reference.md | 16 +- docs/vllm-router.md | 15 +- examples/README.md | 6 + examples/atom/atomesh-disagg.yaml | 32 +++ src/srtctl/backends/__init__.py | 6 +- src/srtctl/backends/atom.py | 224 +++++++++++++++ src/srtctl/backends/base.py | 7 + src/srtctl/backends/mocker.py | 6 + src/srtctl/backends/sglang.py | 29 +- src/srtctl/backends/trtllm.py | 6 + src/srtctl/backends/vllm.py | 6 + .../benchmarks/scripts/lm-eval/bench.sh | 88 ++++++ src/srtctl/cli/do_sweep.py | 6 +- src/srtctl/cli/mixins/worker_stage.py | 13 + src/srtctl/core/config.py | 1 + src/srtctl/core/roles.py | 1 + src/srtctl/core/schema.py | 22 +- src/srtctl/core/schema_docs.py | 20 +- src/srtctl/frontends/__init__.py | 2 + src/srtctl/frontends/atomesh.py | 45 +++ src/srtctl/frontends/vllm_router.py | 29 +- tests/test_atom_atomesh.py | 262 ++++++++++++++++++ tests/test_benchmarks.py | 67 +++++ tests/test_configs.py | 42 ++- tests/test_frontends.py | 10 +- tests/test_port_allocator.py | 2 +- tests/test_schema_docs.py | 11 +- tests/test_slurm.py | 32 ++- tests/test_vllm_router_frontend.py | 63 +++++ 32 files changed, 1044 insertions(+), 42 deletions(-) create mode 100644 examples/atom/atomesh-disagg.yaml create mode 100644 src/srtctl/backends/atom.py create mode 100644 src/srtctl/frontends/atomesh.py create mode 100644 tests/test_atom_atomesh.py diff --git a/CLAUDE.md b/CLAUDE.md index e8c34876d..451893f92 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -369,6 +369,7 @@ with patch.dict(os.environ, H100Rack.slurm_env()): - `endpoints_to_processes()` - Physical process mapping; every port through `NodePortAllocator` - `build_worker_command(process, runtime)` - Command construction - `get_process_environment(process)` - Per-process env derived from `Process` ports (side channels, scan bases) + - `get_frontend_integration_environment(mode, frontend_type, frontend_args)` - frontend integration defaults, or `{}` - `mooncake_kv_store` / `get_mooncake_worker_env(...)` - the Mooncake block and its worker env; `None` / `{}` without one - `failover` / `get_failover_environment(...)` - shadow engine recovery; `None` / `{}` without it - `should_set_visible_devices()` - `True` unless the engine takes its devices on the command line; the variable is the cluster's `visible_devices_env` @@ -377,6 +378,7 @@ with patch.dict(os.environ, H100Rack.slurm_env()): 4. Add polymorphic deserialization in `BackendConfigField` in `schema.py` **Current backends:** +- **ATOM**: Native ROCm servers behind AToMesh, with one Slurm node per logical worker and allocator-owned Mooncake handshake ports - **SGLang**: Per-process srun launching, supports prefill/decode/aggregated modes - **TRTLLM**: MPI-style launching (one srun per endpoint with all nodes), prefill/decode only - **vLLM**: Per-process srun launching, prefill/decode/aggregated, `per_node` DP; `frontend_type` selects Dynamo registration or a direct `vllm serve` server, and `_CONNECTOR_MAP` owns the KV connector table diff --git a/README.md b/README.md index 856af1295..6044d25e3 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # srtctl -Command-line tool for distributed LLM inference benchmarks on SLURM clusters using TensorRT LLM, SGLang and vLLM. Replace complex shell scripts and 50+ CLI flags with a declarative `schema: 2` YAML recipe: `engine:` names the engine, `roles:` describes each worker role, and `services:` covers everything launched next to the workers. +Command-line tool for distributed LLM inference benchmarks on SLURM clusters using ATOM, TensorRT LLM, SGLang and vLLM. Replace complex shell scripts and 50+ CLI flags with a declarative `schema: 2` YAML recipe: `engine:` names the engine, `roles:` describes each worker role, and `services:` covers everything launched next to the workers. ## Quick Start diff --git a/docs/config-reference.md b/docs/config-reference.md index 3947e7587..dfac865c4 100644 --- a/docs/config-reference.md +++ b/docs/config-reference.md @@ -43,6 +43,19 @@ This page is the prose guide: what each block means, how the pieces interact, an ## Overview +### ATOM with AToMesh + +Use `engine: atom` with `frontend.type: atomesh` to launch native +`atom.entrypoints.openai_server` workers and the official AToMesh router. Both +aggregate workers and prefill/decode topologies use static HTTP endpoints; +disaggregated workers receive topology-owned Mooncake handshake ports. + +Engine flags belong under `roles.prefill.args`, `roles.decode.args`, or +`roles.agg.args` (schema v2). +srt-slurm owns the model path, HTTP port, tensor parallel size, and KV-transfer +contract, so recipes cannot override those arguments. See the complete +[ATOM/AToMesh recipe](../examples/atom/atomesh-disagg.yaml). + ```yaml schema: 2 # Required: recipe layout version name: "my-benchmark" # Required: job name diff --git a/docs/schema-reference.md b/docs/schema-reference.md index 60689cdd3..6c18b8cf4 100644 --- a/docs/schema-reference.md +++ b/docs/schema-reference.md @@ -13,7 +13,7 @@ Top-level keys of a recipe YAML. | `name` | str | required | | | `model` | [ModelConfig](#modelconfig) | required | | | `resources` | [ResourceConfig](#resourceconfig) | required | | -| `engine` | str \| mapping | required | The engine type (`sglang`, `trtllm`, `vllm`, `mocker`) as a string, or a mapping with `type` plus the engine-wide knobs listed under [Engine types](#engine-types). | +| `engine` | str \| mapping | required | The engine type (`atom`, `sglang`, `trtllm`, `vllm`, `mocker`) as a string, or a mapping with `type` plus the engine-wide knobs listed under [Engine types](#engine-types). | | `roles` | mapping of role -> [Role](#roles) | required | One block per worker role (`prefill`, `decode`, `agg`): topology, env, and engine args. | | `schema` | int | `2` | Recipe schema version. Write `schema: 2` for this layout. | | `slurm` | [SlurmConfig](#slurmconfig) | `SlurmConfig()` | | @@ -44,7 +44,7 @@ Three vocabularies are specific to the 2.0 layout. They are normalized into the ### engine -`engine: ` or `engine: {type: , ...}`. `type` is one of `sglang`, `trtllm`, `vllm`, `mocker`; the remaining keys are that engine's knobs, listed under [Engine types](#engine-types). +`engine: ` or `engine: {type: , ...}`. `type` is one of `atom`, `sglang`, `trtllm`, `vllm`, `mocker`; the remaining keys are that engine's knobs, listed under [Engine types](#engine-types). ### roles @@ -573,6 +573,18 @@ Ready when the service's log file contains a line matching the regular expressio `engine.type` selects one of the following; the remaining `engine` keys are that type's knobs. +### AtomProtocol + +`engine.type: atom` + +Launch ``atom.entrypoints.openai_server`` on ROCm workers. + +| Key | Type | Default | Description | +|---|---|---|---| +| `type` | one of `'atom'` | `'atom'` | | +| `connector` | one of `'mooncake'` | `'mooncake'` | | +| `mooncake_protocol` | one of `'rdma'`, `'tcp'` \| None | `None` | | + ### SGLangProtocol `engine.type: sglang` diff --git a/docs/vllm-router.md b/docs/vllm-router.md index 5ca9ff9a8..30b16625d 100644 --- a/docs/vllm-router.md +++ b/docs/vllm-router.md @@ -68,8 +68,10 @@ roles: ### Node-local data parallelism For DEP8 on two four-GPU nodes, upstream srt-slurm creates one hybrid-LB -`vllm serve` process on each node. Router receives both base URLs and srtctl -adds `--intra-node-data-parallel-size 4`, exposing all eight DP ranks. +`vllm serve` process on each node. Router receives both base URLs as two +node-local pools. srtctl does not apply `--intra-node-data-parallel-size` to +this multi-node endpoint because the second pool owns global ranks 4 through +7; vLLM's node-local hybrid load balancer selects among those ranks. ```yaml resources: @@ -137,10 +139,11 @@ roles: data-parallel-size: 4 ``` -Router has one global `--intra-node-data-parallel-size`, so every advertised -P/D base must represent the same number of local DP ranks. srtctl derives and -validates that value; do not set it manually unless it exactly matches the -allocated topology. +Router has one global `--intra-node-data-parallel-size`, so srtctl uses it only +when every logical endpoint has one advertised base URL. Multi-node hybrid-LB +endpoints remain unexpanded node-local pools because each later pool starts at +a nonzero global DP rank. Do not set the option manually unless it exactly +matches the allocated topology. ### MoRI-IO discovery diff --git a/examples/README.md b/examples/README.md index 7d403e1e4..487fa7bd9 100644 --- a/examples/README.md +++ b/examples/README.md @@ -62,3 +62,9 @@ for p in sorted(Path('examples').rglob('*.yaml')): print(validate_config_file(p) or f'ok {p}') " ``` + +### ATOM and AToMesh + +[`atom/atomesh-disagg.yaml`](atom/atomesh-disagg.yaml) launches native ATOM prefill +and decode workers with the AToMesh router. Each logical worker fits on one node; +the cluster configuration selects ROCm device visibility. diff --git a/examples/atom/atomesh-disagg.yaml b/examples/atom/atomesh-disagg.yaml new file mode 100644 index 000000000..669bb057b --- /dev/null +++ b/examples/atom/atomesh-disagg.yaml @@ -0,0 +1,32 @@ +# Native ATOM prefill/decode on separate Slurm nodes behind AToMesh. +# Configure visible_devices_env: ROCR_VISIBLE_DEVICES in the cluster srtslurm.yaml. +# The container must include both atom.entrypoints.openai_server and atomesh. +schema: 2 +name: qwen3-0.6b-atomesh-disagg +model: + path: hf:Qwen/Qwen3-0.6B + container: rocm/atom:latest + precision: bf16 +resources: + gpu_type: mi300x + gpus_per_node: 8 +engine: + type: atom + mooncake_protocol: tcp +frontend: + type: atomesh + enable_multiple_frontends: false +roles: + prefill: + nodes: 1 + workers: 1 + gpus: 1 + decode: + nodes: 1 + workers: 1 + gpus: 1 +benchmark: + type: sa-bench + isl: 128 + osl: 128 + concurrencies: "4x8" diff --git a/src/srtctl/backends/__init__.py b/src/srtctl/backends/__init__.py index 24890ed58..f184d8f09 100644 --- a/src/srtctl/backends/__init__.py +++ b/src/srtctl/backends/__init__.py @@ -9,6 +9,7 @@ - TRTLLM: TensorRT-LLM backend with prefill/decode disaggregation """ +from .atom import AtomProtocol, AtomServerConfig from .base import BackendProtocol, BackendType, SrunConfig from .mocker import MockerProtocol, MockerServerConfig from .sglang import MooncakeKVStoreConfig, SGLangProtocol, SGLangServerConfig @@ -16,9 +17,12 @@ from .vllm import VLLMFailoverConfig, VLLMMooncakeKVStoreConfig, VLLMProtocol, VLLMServerConfig # Union type for all backend configs -BackendConfig = SGLangProtocol | TRTLLMProtocol | VLLMProtocol | MockerProtocol +BackendConfig = AtomProtocol | SGLangProtocol | TRTLLMProtocol | VLLMProtocol | MockerProtocol __all__ = [ + # ATOM + "AtomProtocol", + "AtomServerConfig", "BackendConfig", # Base types "BackendProtocol", diff --git a/src/srtctl/backends/atom.py b/src/srtctl/backends/atom.py new file mode 100644 index 000000000..235ccb8e0 --- /dev/null +++ b/src/srtctl/backends/atom.py @@ -0,0 +1,224 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""ROCm ATOM inference backend.""" + +from __future__ import annotations + +import builtins +import json +from collections.abc import Sequence +from dataclasses import field +from pathlib import Path +from typing import TYPE_CHECKING, Any, ClassVar, Literal + +from marshmallow import Schema +from marshmallow_dataclass import dataclass + +from srtctl.ports import DYN_SYSTEM_PORT_BASE + +if TYPE_CHECKING: + from srtctl.backends.base import SrunConfig + from srtctl.core.runtime import RuntimeContext + from srtctl.core.schema import ProfilingConfig + from srtctl.core.topology import Endpoint, NodePortAllocator, Process + +WorkerMode = Literal["prefill", "decode", "agg"] + + +@dataclass(frozen=True) +class AtomServerConfig: + """Native ATOM CLI arguments for each serving role.""" + + prefill: dict[str, Any] | None = None + decode: dict[str, Any] | None = None + aggregated: dict[str, Any] | None = None + + Schema: ClassVar[type[Schema]] = Schema + + +@dataclass(frozen=True) +class AtomProtocol: + """Launch ``atom.entrypoints.openai_server`` on ROCm workers.""" + + type: Literal["atom"] = "atom" + prefill_environment: dict[str, str] = field(default_factory=dict) + decode_environment: dict[str, str] = field(default_factory=dict) + aggregated_environment: dict[str, str] = field(default_factory=dict) + atom_config: AtomServerConfig | None = None + connector: Literal["mooncake"] = "mooncake" + mooncake_protocol: Literal["rdma", "tcp"] | None = None + + Schema: ClassVar[builtins.type[Schema]] = Schema + + def get_srun_config(self) -> SrunConfig: + from srtctl.backends.base import SrunConfig + + return SrunConfig(mpi=None, oversubscribe=False, launch_per_endpoint=False) + + def get_config_for_mode(self, mode: WorkerMode) -> dict[str, Any]: + if self.atom_config is None: + return {} + values = { + "prefill": self.atom_config.prefill, + "decode": self.atom_config.decode, + "agg": self.atom_config.aggregated, + } + return dict(values.get(mode) or {}) + + def get_environment_for_mode(self, mode: WorkerMode) -> dict[str, str]: + values = { + "prefill": self.prefill_environment, + "decode": self.decode_environment, + "agg": self.aggregated_environment, + } + return dict(values.get(mode) or {}) + + def get_frontend_integration_environment( + self, mode: str, frontend_type: str, frontend_args: dict[str, Any] + ) -> dict[str, str]: + """Worker environment defaults required by the selected frontend integration.""" + return {} + + def get_process_environment(self, process: Process) -> dict[str, str]: + return {} + + def get_served_model_name(self, default: str) -> str: + return default + + @property + def mooncake_kv_store(self) -> None: + return None + + @property + def failover(self) -> None: + return None + + def get_mooncake_worker_env(self, infra_node_ip: str, local_hostname: str) -> dict[str, str]: + return {} + + def get_failover_environment(self, process: Process, job_id: str) -> dict[str, str]: + return {} + + def should_set_visible_devices(self) -> bool: + return True + + def allocate_endpoints( + self, + num_prefill: int, + num_decode: int, + num_agg: int, + gpus_per_prefill: int, + gpus_per_decode: int, + gpus_per_agg: int, + gpus_per_node: int, + available_nodes: Sequence[str], + spread_workers: bool = False, + ) -> list[Endpoint]: + from srtctl.core.topology import allocate_endpoints + + return allocate_endpoints( + num_prefill=num_prefill, + num_decode=num_decode, + num_agg=num_agg, + gpus_per_prefill=gpus_per_prefill, + gpus_per_decode=gpus_per_decode, + gpus_per_agg=gpus_per_agg, + gpus_per_node=gpus_per_node, + available_nodes=available_nodes, + spread_workers=spread_workers, + ) + + def endpoints_to_processes( + self, + endpoints: list[Endpoint], + base_sys_port: int = DYN_SYSTEM_PORT_BASE, + port_allocator: NodePortAllocator | None = None, + frontend_type: str = "atomesh", + dynamo_sidecar: bool = False, + ) -> list[Process]: + if dynamo_sidecar: + raise ValueError("ATOM does not support Dynamo sidecars") + from srtctl.core.topology import endpoints_to_processes + + return endpoints_to_processes(endpoints, base_sys_port=base_sys_port, port_allocator=port_allocator) + + def _kv_transfer_config(self, process: Process, worker_ip: str) -> str: + if process.endpoint_mode not in {"prefill", "decode"}: + raise ValueError("ATOM KV transfer is only valid for prefill/decode workers") + if process.nixl_port is None: + raise ValueError("ATOM P/D worker is missing its Mooncake handshake port") + payload = { + "kv_role": "kv_producer" if process.endpoint_mode == "prefill" else "kv_consumer", + "kv_connector": self.connector, + "proxy_ip": worker_ip, + "handshake_port": process.nixl_port, + } + if self.mooncake_protocol is not None: + payload["protocol"] = self.mooncake_protocol + return json.dumps(payload, separators=(",", ":")) + + def build_worker_command( + self, + process: Process, + endpoint_processes: list[Process], + runtime: RuntimeContext, + frontend_type: str = "atomesh", + nsys_prefix: list[str] | None = None, + dump_config_path: Path | None = None, + profiling: ProfilingConfig | None = None, + ) -> list[str]: + if frontend_type != "atomesh": + raise ValueError(f"backend.type: atom requires frontend.type: atomesh (got {frontend_type!r})") + if len({item.node for item in endpoint_processes}) != 1: + raise ValueError("ATOM currently requires each logical endpoint to fit on one Slurm node") + + from srtctl.core.slurm import get_hostname_ip + + worker_ip = get_hostname_ip(process.node, runtime.network_interface) + config = self.get_config_for_mode(process.endpoint_mode) + reserved = {"model", "host", "server-port", "tp", "tensor-parallel-size", "kv-transfer-config"} + overlap = reserved.intersection(_canonical_arg_key(key) for key in config) + if overlap: + raise ValueError(f"ATOM config cannot override srtctl-managed argument(s): {sorted(overlap)}") + + command = ["env", f"ATOM_HOST_IP={worker_ip}", *(nsys_prefix or [])] + command.extend( + [ + "python3", + "-m", + "atom.entrypoints.openai_server", + "--model", + runtime.worker_model_arg, + "--host", + "0.0.0.0", + "--server-port", + str(process.http_port), + "-tp", + str(len(process.gpu_indices)), + ] + ) + if process.endpoint_mode in {"prefill", "decode"}: + command.extend(["--kv-transfer-config", self._kv_transfer_config(process, worker_ip)]) + command.extend(_config_to_cli_args(config)) + return command + + +def _config_to_cli_args(config: dict[str, Any]) -> list[str]: + """Preserve ATOM's native CLI spelling, which mixes hyphens and underscores.""" + args: list[str] = [] + for key, value in sorted(config.items()): + flag = f"--{key}" + if value is True: + args.append(flag) + elif value is False or value is None: + continue + elif isinstance(value, list): + args.extend([flag, *(str(item) for item in value)]) + else: + args.extend([flag, str(value)]) + return args + + +def _canonical_arg_key(key: str) -> str: + return key.lstrip("-").replace("_", "-") diff --git a/src/srtctl/backends/base.py b/src/srtctl/backends/base.py index 2b0cba2c9..cf4183130 100644 --- a/src/srtctl/backends/base.py +++ b/src/srtctl/backends/base.py @@ -29,6 +29,7 @@ class BackendType(str, Enum): TRTLLM = "trtllm" VLLM = "vllm" MOCKER = "mocker" + ATOM = "atom" @dataclass @@ -155,6 +156,12 @@ def build_worker_command( """Build command to start a worker process.""" ... + def get_frontend_integration_environment( + self, mode: str, frontend_type: str, frontend_args: dict[str, Any] + ) -> dict[str, str]: + """Worker environment defaults required by the selected frontend integration.""" + ... + def get_process_environment(self, process: "Process") -> dict[str, str]: """Get process-specific environment variables. diff --git a/src/srtctl/backends/mocker.py b/src/srtctl/backends/mocker.py index 06db1ac56..79811ea55 100644 --- a/src/srtctl/backends/mocker.py +++ b/src/srtctl/backends/mocker.py @@ -166,6 +166,12 @@ def get_environment_for_mode(self, mode: WorkerMode) -> dict[str, str]: return dict(self.aggregated_environment) return {} + def get_frontend_integration_environment( + self, mode: str, frontend_type: str, frontend_args: dict[str, Any] + ) -> dict[str, str]: + """Worker environment defaults required by the selected frontend integration.""" + return {} + def get_process_environment(self, process: "Process") -> dict[str, str]: """Get process-specific environment variables. diff --git a/src/srtctl/backends/sglang.py b/src/srtctl/backends/sglang.py index f2de9f83b..c4b224b14 100644 --- a/src/srtctl/backends/sglang.py +++ b/src/srtctl/backends/sglang.py @@ -208,6 +208,31 @@ def get_process_environment(self, process: "Process") -> dict[str, str]: """ return {} + def get_frontend_integration_environment( + self, + mode: str, + frontend_type: str, + frontend_args: dict[str, Any], + ) -> dict[str, str]: + """Bridge SGLang Model Gateway's DP routing to P/D KV bootstrap. + + A DP-aware external router supplies ``routed_dp_rank`` independently + for prefill and decode. SGLang's P/D bootstrap must therefore register + the actual selected prefill rank instead of deriving it from the + bootstrap-room modulo. The environment variable activates SGLang's + supported register/query path on both sides of the transfer. + """ + if frontend_type != "sglang-router" or mode not in ("prefill", "decode"): + return {} + if not frontend_args.get("dp-aware", frontend_args.get("dp_aware", False)): + return {} + + mode_config = self.get_config_for_mode(mode) + dp_size = mode_config.get("dp-size", mode_config.get("dp_size", 1)) + if int(dp_size or 1) <= 1: + return {} + return {"SGLANG_DISAGGREGATION_FORCE_QUERY_PREFILL_DP_RANK": "1"} + def get_mooncake_worker_env(self, infra_node_ip: str, local_hostname: str) -> dict[str, str]: """Get mooncake env vars to inject on a specific worker. @@ -394,7 +419,7 @@ def build_worker_command( is_multi_node = len(endpoint_nodes) > 1 # Get leader IP for distributed init - leader_ip = get_hostname_ip(endpoint_nodes[0]) + leader_ip = get_hostname_ip(endpoint_nodes[0], runtime.network_interface) # Direct frontends run the native server; Dynamo frontends run the registering worker. use_sglang = frontend.worker_launch == "direct" @@ -523,7 +548,7 @@ def _build_sidecar_command( endpoint_nodes = list(dict.fromkeys(candidate.node for candidate in endpoint_processes)) node_rank = endpoint_nodes.index(process.node) is_leader = node_rank == 0 - leader_ip = get_hostname_ip(endpoint_nodes[0]) + leader_ip = get_hostname_ip(endpoint_nodes[0], runtime.network_interface) grpc_port = sidecar_grpc_port(process) served_model_name = self.get_served_model_name(runtime.model_path.name) diff --git a/src/srtctl/backends/trtllm.py b/src/srtctl/backends/trtllm.py index 91a5549b4..0c2df0e8d 100644 --- a/src/srtctl/backends/trtllm.py +++ b/src/srtctl/backends/trtllm.py @@ -237,6 +237,12 @@ def get_environment_for_mode(self, mode: WorkerMode) -> dict[str, str]: env["TLLM_NUMA_AWARE_WORKER_AFFINITY"] = "0" return env + def get_frontend_integration_environment( + self, mode: str, frontend_type: str, frontend_args: dict[str, Any] + ) -> dict[str, str]: + """Worker environment defaults required by the selected frontend integration.""" + return {} + def get_process_environment(self, process: "Process") -> dict[str, str]: """Get process-specific environment variables. diff --git a/src/srtctl/backends/vllm.py b/src/srtctl/backends/vllm.py index 048627a9b..dfa8f0431 100644 --- a/src/srtctl/backends/vllm.py +++ b/src/srtctl/backends/vllm.py @@ -600,6 +600,12 @@ def _discovery_extra_config(self, process: Process, runtime: RuntimeContext) -> "read_mode": True, } + def get_frontend_integration_environment( + self, mode: str, frontend_type: str, frontend_args: dict[str, Any] + ) -> dict[str, str]: + """Worker environment defaults required by the selected frontend integration.""" + return {} + def get_process_environment(self, process: Process) -> dict[str, str]: """Get process-specific environment variables for vLLM workers. diff --git a/src/srtctl/benchmarks/scripts/lm-eval/bench.sh b/src/srtctl/benchmarks/scripts/lm-eval/bench.sh index a10e4e7d3..6e52147b1 100755 --- a/src/srtctl/benchmarks/scripts/lm-eval/bench.sh +++ b/src/srtctl/benchmarks/scripts/lm-eval/bench.sh @@ -17,6 +17,83 @@ PORT=$(echo "$ENDPOINT" | sed -E 's|.*:([0-9]+).*|\1|') echo "lm-eval Config: endpoint=${ENDPOINT}; host=${HOST}; port=${PORT}; workspace=${INFMAX_WORKSPACE}" +# Serving images commonly make their system Python environment read-only. The +# InferenceX eval harness installs a pinned lm-eval runtime before executing, so +# give it a job-local writable environment while retaining the serving image's +# already-installed framework dependencies. Prepending the venv to PATH keeps +# benchmark_lib.sh's existing `python3 -m ...` interface unchanged. +LM_EVAL_RUNTIME_DIR="${SRTCTL_LM_EVAL_RUNTIME_DIR:-${TMPDIR:-/tmp}/srtctl-lm-eval-${SLURM_JOB_ID:-$$}}" +LM_EVAL_VENV="${LM_EVAL_RUNTIME_DIR}/venv" +LM_EVAL_CACHE_DIR="${SRTCTL_LM_EVAL_CACHE_DIR:-${LM_EVAL_RUNTIME_DIR}/cache}" +LM_EVAL_RESULT_DIR="${SRTCTL_LM_EVAL_RESULT_DIR:-}" + +# Serving containers often expose the host user's home and shared model cache +# read-only. lm-eval still needs writable Hugging Face and XDG caches for task +# datasets, even when model weights are already present. Keep all client-side +# downloads inside the disposable job-local runtime instead of mutating the +# serving cache or relying on $HOME/.cache. +export XDG_CACHE_HOME="${LM_EVAL_CACHE_DIR}/xdg" +export HF_HOME="${LM_EVAL_CACHE_DIR}/huggingface" +export HF_HUB_CACHE="${HF_HOME}/hub" +export HUGGINGFACE_HUB_CACHE="${HF_HUB_CACHE}" +export HF_DATASETS_CACHE="${HF_HOME}/datasets" +mkdir -p "${XDG_CACHE_HOME}" "${HF_HUB_CACHE}" "${HF_DATASETS_CACHE}" + +# The Slurm step can receive a reduced PATH even when the serving image keeps +# its ROCm/PyTorch environment under /opt/venv. lm-eval imports torch for its +# API client helpers, so preserve the framework environment's site-packages in +# the disposable eval venv instead of downloading a second (and potentially +# incompatible) torch build. Prefer an explicit override, then the standard +# SGLang image environment, and finally the Python visible on PATH. +LM_EVAL_FRAMEWORK_PYTHON="${SRTCTL_LM_EVAL_FRAMEWORK_PYTHON:-}" +if [[ -z "${LM_EVAL_FRAMEWORK_PYTHON}" ]]; then + for candidate in /opt/venv/bin/python3 "$(command -v python3)"; do + if [[ -x "${candidate}" ]] && "${candidate}" -c 'import torch' >/dev/null 2>&1; then + LM_EVAL_FRAMEWORK_PYTHON="${candidate}" + break + fi + done +fi +if [[ -z "${LM_EVAL_FRAMEWORK_PYTHON}" ]]; then + echo "ERROR: no serving-image Python with torch is available for lm-eval" >&2 + exit 1 +fi +LM_EVAL_FRAMEWORK_SITE_PACKAGES="$("${LM_EVAL_FRAMEWORK_PYTHON}" - <<'PY' +import sys + +print("\n".join(path for path in sys.path if "site-packages" in path)) +PY +)" + +if [[ ! -x "${LM_EVAL_VENV}/bin/python3" ]]; then + rm -rf "${LM_EVAL_VENV}" + mkdir -p "${LM_EVAL_RUNTIME_DIR}" + python3 -m venv --system-site-packages "${LM_EVAL_VENV}" +fi +LM_EVAL_SITE_PACKAGES="$("${LM_EVAL_VENV}/bin/python3" -c 'import site; print(site.getsitepackages()[0])')" +printf '%s\n' "${LM_EVAL_FRAMEWORK_SITE_PACKAGES}" > "${LM_EVAL_SITE_PACKAGES}/srtctl-framework.pth" +export PATH="${LM_EVAL_VENV}/bin:${PATH}" +hash -r + +if [[ "$(python3 -c 'import sys; print(sys.prefix)')" != "${LM_EVAL_VENV}" ]]; then + echo "ERROR: failed to activate writable lm-eval runtime at ${LM_EVAL_VENV}" >&2 + exit 1 +fi +echo "lm-eval Runtime: python=$(command -v python3); prefix=${LM_EVAL_VENV}" +python3 -c 'import torch' || { + echo "ERROR: job-local lm-eval runtime cannot import serving-image torch" >&2 + exit 1 +} + +# Some serving images seed virtual environments with a pip version old enough +# that it does not recognize --break-system-packages. InferenceX's shared eval +# installer passes that option, so make the job-local pip understand it before +# handing control to benchmark_lib.sh. This only mutates the disposable venv. +if ! python3 -m pip install --help 2>/dev/null | grep -q -- '--break-system-packages'; then + echo "lm-eval Runtime: upgrading job-local pip for --break-system-packages support" + python3 -m pip install --upgrade 'pip>=23.0' +fi + # Auto-discover the served model name from /v1/models if MODEL_NAME is not set. # This ensures we use the exact name the server recognizes, regardless of what # $MODEL (the HuggingFace ID from the workflow) is set to. @@ -69,6 +146,17 @@ cp -v meta_env.json /logs/eval_results/ 2>/dev/null || true cp -v results*.json /logs/eval_results/ 2>/dev/null || true cp -v sample*.jsonl /logs/eval_results/ 2>/dev/null || true +# Integrations that maintain a separate result tree can request an additional +# copy without coupling this benchmark to a particular host mount layout. The +# default remains /logs/eval_results for existing srt-slurm consumers. +if [[ -n "${LM_EVAL_RESULT_DIR}" ]]; then + mkdir -p "${LM_EVAL_RESULT_DIR}" + echo "Copying eval artifacts to ${LM_EVAL_RESULT_DIR}/..." + cp -v meta_env.json "${LM_EVAL_RESULT_DIR}/" 2>/dev/null || true + cp -v results*.json "${LM_EVAL_RESULT_DIR}/" 2>/dev/null || true + cp -v sample*.jsonl "${LM_EVAL_RESULT_DIR}/" 2>/dev/null || true +fi + if [[ "$eval_rc" -ne 0 ]]; then echo "lm-eval evaluation failed with exit code ${eval_rc}" exit "$eval_rc" diff --git a/src/srtctl/cli/do_sweep.py b/src/srtctl/cli/do_sweep.py index 33d07af38..32607a4ff 100644 --- a/src/srtctl/cli/do_sweep.py +++ b/src/srtctl/cli/do_sweep.py @@ -548,7 +548,11 @@ def _run_post_eval(self, stop_event: threading.Event) -> int: # Pass through eval-related env vars. InferenceX writes multi-node # metadata from these variables in append_lm_eval_summary(). The recipe # extends this list with post_eval.passthrough_env. - env_to_set = {} + # Post-eval replaces the configured benchmark runner with lm-eval, but + # it is still a benchmark process. Preserve the recipe's benchmark + # environment so integrations can pass runner-specific settings such + # as an additional artifact sink through this substituted path. + env_to_set = {key: self.runtime.format_string(value) for key, value in self.config.benchmark.env.items()} for var in [ *self.config.post_eval.passthrough_env, "RUN_EVAL", diff --git a/src/srtctl/cli/mixins/worker_stage.py b/src/srtctl/cli/mixins/worker_stage.py index 0a441438f..a51252c78 100644 --- a/src/srtctl/cli/mixins/worker_stage.py +++ b/src/srtctl/cli/mixins/worker_stage.py @@ -222,6 +222,16 @@ def _profiling_selects_process(self, process: "Process") -> bool: process.node_rank, ) + def _apply_frontend_integration_env(self, env_to_set: dict[str, str], mode: str) -> None: + """Add backend/frontend integration defaults without overriding recipes.""" + integration_env = self.backend.get_frontend_integration_environment( + mode, + self.config.frontend.type, + dict(self.config.frontend.args or {}), + ) + for key, value in integration_env.items(): + env_to_set.setdefault(key, value) + def start_worker(self, process: "Process", endpoint_processes: list["Process"]) -> ManagedProcess: """Start a single worker process (one srun per node, used by SGLang).""" mode = process.endpoint_mode @@ -315,6 +325,7 @@ def __missing__(self, key: str) -> str: formatted_value = value.format_map(SafeDict(template_vars)) env_to_set[key] = formatted_value + self._apply_frontend_integration_env(env_to_set, mode) env_to_set.update(self._visible_device_environment(process)) # Add backend-specific process environment variables (e.g., unique ports) @@ -528,6 +539,8 @@ def start_endpoint_worker(self, endpoint_processes: list["Process"]) -> ManagedP ): env_to_set.setdefault("DYN_TRTLLM_KV_EVENT_HOSTS", ",".join(endpoint_nodes)) + self._apply_frontend_integration_env(env_to_set, mode) + force_mask = self.config.dynamo.sidecar and self.backend.type == "vllm" node_gpu_setup = "" if force_mask or self.backend.should_set_visible_devices(): diff --git a/src/srtctl/core/config.py b/src/srtctl/core/config.py index 8733fda97..f6aec0ab7 100755 --- a/src/srtctl/core/config.py +++ b/src/srtctl/core/config.py @@ -104,6 +104,7 @@ def load_cluster_config() -> dict[str, Any] | None: "sbatch_directives", "srun_options", "sglang_config", + "atom_config", "vllm_config", "trtllm_config", "mocker_config", diff --git a/src/srtctl/core/roles.py b/src/srtctl/core/roles.py index 082d25148..57c643e61 100644 --- a/src/srtctl/core/roles.py +++ b/src/srtctl/core/roles.py @@ -54,6 +54,7 @@ # backend.type -> the engine's per-mode CLI config key. ENGINE_CONFIG_KEY: dict[str, str] = { + "atom": "atom_config", "sglang": "sglang_config", "vllm": "vllm_config", "trtllm": "trtllm_config", diff --git a/src/srtctl/core/schema.py b/src/srtctl/core/schema.py index 9b1036638..7502b5f5f 100755 --- a/src/srtctl/core/schema.py +++ b/src/srtctl/core/schema.py @@ -35,6 +35,7 @@ from marshmallow_dataclass import dataclass from srtctl.backends import ( + AtomProtocol, BackendConfig, MockerProtocol, SGLangProtocol, @@ -448,7 +449,7 @@ def _deserialize( # Default to SGLang return SGLangProtocol() - if isinstance(value, SGLangProtocol | TRTLLMProtocol | VLLMProtocol | MockerProtocol): + if isinstance(value, AtomProtocol | SGLangProtocol | TRTLLMProtocol | VLLMProtocol | MockerProtocol): return value if not isinstance(value, dict): @@ -457,7 +458,9 @@ def _deserialize( # Get backend type from the value dict backend_type = value.get("type", "sglang") - if backend_type == "sglang": + if backend_type == "atom": + return AtomProtocol.Schema().load(value) + elif backend_type == "sglang": schema = SGLangProtocol.Schema() return schema.load(value) elif backend_type == "trtllm": @@ -471,13 +474,15 @@ def _deserialize( return schema.load(value) else: raise ValidationError( - f"Unknown backend type: {backend_type!r}. Supported types: sglang, trtllm, vllm, mocker" + f"Unknown backend type: {backend_type!r}. Supported types: atom, sglang, trtllm, vllm, mocker" ) def _serialize(self, value: Any | None, attr: str | None, obj: Any, **kwargs) -> Any: """Serialize backend config to dict.""" if value is None: return None + if isinstance(value, AtomProtocol): + return AtomProtocol.Schema().dump(value) if isinstance(value, SGLangProtocol): return SGLangProtocol.Schema().dump(value) if isinstance(value, TRTLLMProtocol): @@ -3228,6 +3233,17 @@ def from_yaml(cls, yaml_path: Path) -> "SrtConfig": def served_model_name(self) -> str: """Get the served model name from backend config or model path.""" default = Path(self.model.path).name + if isinstance(self.backend, AtomProtocol): + # ATOM advertises the literal --model argument; unlike SGLang/vLLM, + # it has no separate served-model-name alias. Match the worker's + # HF ID or container-visible path, including node-local staging. + model_path = os.path.expandvars(self.model.path) + if model_path.startswith("hf:"): + default = model_path[3:] + elif self.model.stage_dir: + default = str(Path(os.path.expandvars(self.model.stage_dir)) / Path(model_path).resolve().name) + else: + default = "/model" return self.backend.get_served_model_name(default) @property diff --git a/src/srtctl/core/schema_docs.py b/src/srtctl/core/schema_docs.py index 4de1bcb09..973bb9ea4 100644 --- a/src/srtctl/core/schema_docs.py +++ b/src/srtctl/core/schema_docs.py @@ -43,6 +43,8 @@ from typing import Annotated, Any, Literal, get_args, get_origin, get_type_hints from srtctl.backends import ( + AtomProtocol, + AtomServerConfig, MockerProtocol, MockerServerConfig, SGLangProtocol, @@ -79,6 +81,7 @@ # backend.type value -> dataclass. Order is the documentation order. BACKEND_TYPES: tuple[tuple[str, type], ...] = ( + ("atom", AtomProtocol), ("sglang", SGLangProtocol), ("trtllm", TRTLLMProtocol), ("vllm", VLLMProtocol), @@ -103,7 +106,7 @@ class FieldDoc: type_label="str \\| mapping", default="required", description=( - "The engine type (`sglang`, `trtllm`, `vllm`, `mocker`) as a string, or a mapping with `type` " + "The engine type (`atom`, `sglang`, `trtllm`, `vllm`, `mocker`) as a string, or a mapping with `type` " "plus the engine-wide knobs listed under [Engine types](#engine-types)." ), ), @@ -431,6 +434,10 @@ def _present(cls: type, mapping: dict[str, str]) -> dict[str, str]: VLLMMooncakeKVStoreConfig: frozenset({"device_names_by_gpu"}), } +# ATOM was added in v2; its normalized role fields are internal, not v1 API. +INTERNAL_FIELDS = {AtomProtocol: _present(AtomProtocol, _backend_legacy_fields(ENGINE_CONFIG_KEY["atom"]))} +INTERNAL_CLASSES = frozenset({AtomServerConfig}) + def _row(row: FieldDoc) -> str: return f"| `{row.key}` | {row.type_label} | {row.default} | {_cell(row.description)} |" @@ -449,7 +456,7 @@ def _render_table(cls: type, *, only_legacy: bool = False) -> list[str]: legacy = _legacy_keys(cls) out = _table_header() for row in field_docs(cls): - if row.key in _V2_SERVICE_OPTION_FIELDS.get(cls, ()): + if row.key in _V2_SERVICE_OPTION_FIELDS.get(cls, ()) or row.key in INTERNAL_FIELDS.get(cls, {}): continue if (row.key in legacy) == only_legacy: out.append(_row(row)) @@ -576,7 +583,8 @@ def render_schema_reference() -> str: lines.append("") lines.extend(_render_authoring_surface()) - nested = [cls for cls in _walk(SrtConfig, skip=_BACKEND_CLASSES) if cls not in LEGACY_CLASSES] + hidden_classes = LEGACY_CLASSES | INTERNAL_CLASSES + nested = [cls for cls in _walk(SrtConfig, skip=_BACKEND_CLASSES) if cls not in hidden_classes] if nested: lines.extend(["## Recipe sections", ""]) for cls in nested: @@ -599,7 +607,7 @@ def render_schema_reference() -> str: lines.extend(_render_table(cls)) lines.append("") for extra in _walk(cls, skip=_BACKEND_CLASSES | set(nested)): - if extra not in engine_nested and extra not in LEGACY_CLASSES: + if extra not in engine_nested and extra not in hidden_classes: engine_nested.append(extra) for cls in engine_nested: lines.extend(_render_class_section(cls, level=3)) @@ -657,7 +665,7 @@ def render_legacy_reference() -> str: lines.append(f"| `{prefix}.{key}` | {_cell(replacement)} |") backend_rows: dict[str, str] = {} for _, cls in BACKEND_TYPES: - for key, replacement in LEGACY_FIELDS[cls].items(): + for key, replacement in LEGACY_FIELDS.get(cls, {}).items(): backend_rows.setdefault(key, replacement) for key, replacement in backend_rows.items(): lines.append(f"| `backend.{key}` | {_cell(replacement)} |") @@ -704,6 +712,8 @@ def render_legacy_reference() -> str: ) server_configs: list[type] = [] for type_name, cls in BACKEND_TYPES: + if cls not in LEGACY_FIELDS: + continue lines.extend([f"### {cls.__name__}", "", f"`backend.type: {type_name}`", ""]) lines.extend(_render_table(cls, only_legacy=True)) lines.append("") diff --git a/src/srtctl/frontends/__init__.py b/src/srtctl/frontends/__init__.py index ec4920d98..b5e462cf5 100644 --- a/src/srtctl/frontends/__init__.py +++ b/src/srtctl/frontends/__init__.py @@ -17,6 +17,7 @@ - vllm-router: Official vLLM Router with static aggregate or P/D workers """ +from srtctl.frontends.atomesh import AtomeshFrontend from srtctl.frontends.base import ( FRONTEND_NONE, FrontendProtocol, @@ -34,6 +35,7 @@ __all__ = [ "FRONTEND_NONE", + "AtomeshFrontend", "DynamicFrontend", "DynamoFrontend", "FrontendProtocol", diff --git a/src/srtctl/frontends/atomesh.py b/src/srtctl/frontends/atomesh.py new file mode 100644 index 000000000..b9ad5f21b --- /dev/null +++ b/src/srtctl/frontends/atomesh.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Official AToMesh router frontend for native ATOM workers.""" + +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.runtime import RuntimeContext + from srtctl.core.topology import Process + + +@register_frontend("atomesh") +class AtomeshFrontend(StaticRouterFrontend): + """Route aggregate or disaggregated traffic to native ATOM servers.""" + + type: ClassVar[str] = "atomesh" + required_backend: ClassVar[str] = "atom" + executable: ClassVar[tuple[str, ...]] = ("atomesh", "launch") + pd_flag: ClassVar[str] = "--pd-disaggregation" + process_name: ClassVar[str] = "atomesh" + wait_for_workers_before_start: ClassVar[bool] = True + + def worker_metrics_port(self, process: Process, runtime: RuntimeContext) -> None: + """Native ATOM workers do not expose the Prometheus metrics endpoint.""" + + def worker_bootstrap_port(self, backend: Any, process: Process) -> int | None: + """ATOM exposes transfer topology through ``/kv_transfer_info``.""" + return None + + def get_managed_frontend_args( + self, + config: Any, + backend: Any, + backend_processes: list[Process], + ) -> list[str]: + normalized = {str(key).replace("_", "-") for key in (config.frontend.args or {})} + if "backend" in normalized: + raise ValueError("frontend.args.backend is managed by srtctl for atomesh") + return ["--backend", "atom"] diff --git a/src/srtctl/frontends/vllm_router.py b/src/srtctl/frontends/vllm_router.py index 4b34259c8..ed0b4dcbb 100644 --- a/src/srtctl/frontends/vllm_router.py +++ b/src/srtctl/frontends/vllm_router.py @@ -42,7 +42,18 @@ def routed_process_dp_size(backend: Any, process: Process) -> int: def node_local_data_parallel_size(backend: Any, backend_processes: list[Process]) -> int: """Return Router's single DP expansion factor for all advertised URLs.""" - routed_sizes = {routed_process_dp_size(backend, process) for process in backend_processes if process.http_port > 0} + routable = [process for process in backend_processes if process.http_port > 0] + process_count_by_endpoint: dict[tuple[str, int], int] = {} + for process in routable: + endpoint = (process.endpoint_mode, process.endpoint_index) + process_count_by_endpoint[endpoint] = process_count_by_endpoint.get(endpoint, 0) + 1 + + # Hybrid-LB pools on later nodes have nonzero DP-rank offsets. + # Let vLLM route locally; Router expansion would restart ranks at zero. + if any(count > 1 for count in process_count_by_endpoint.values()): + return 1 + + routed_sizes = {routed_process_dp_size(backend, process) for process in routable} if len(routed_sizes) > 1: sizes = ", ".join(str(size) for size in sorted(routed_sizes)) raise ValueError(f"vLLM Router requires one uniform node-local DP expansion factor; derived {sizes}") @@ -83,6 +94,7 @@ def validate(self, config: Any) -> None: ) expansion_by_mode: dict[str, int] = {} + has_multinode_pools = False for mode, gpu_count in endpoint_gpu_counts.items(): if gpu_count <= 0: continue @@ -109,12 +121,16 @@ def validate(self, config: Any) -> None: ) local_gpu_count = min(gpu_count, resources.gpus_per_node) + if replica_size <= local_gpu_count and gpu_count > resources.gpus_per_node: + has_multinode_pools = True if replica_size > local_gpu_count: expansion_by_mode[mode] = 1 else: expansion_by_mode[mode] = backend._get_local_dp_size(mode, local_gpu_count) - expansions = set(expansion_by_mode.values()) + # A single global Router expansion cannot express later pools' DP-rank offsets. + # Match node_local_data_parallel_size: any multi-node pool disables expansion. + expansions = {1} if has_multinode_pools else set(expansion_by_mode.values()) if len(expansions) > 1: detail = ", ".join(f"{mode}={size}" for mode, size in expansion_by_mode.items()) raise ValueError( @@ -300,15 +316,12 @@ def health_expectations(self, config: Any, processes: list[Process] | None) -> t return logical_prefill, logical_decode, f"{worker_desc}, registering with the Router over ZMQ discovery" if processes is None: return logical_prefill, logical_decode, worker_desc + expansion = node_local_data_parallel_size(config.backend, processes) n_prefill = sum( - routed_process_dp_size(config.backend, process) - for process in processes - if process.endpoint_mode == "prefill" and process.http_port > 0 + expansion for process in processes if process.endpoint_mode == "prefill" and process.http_port > 0 ) n_decode = sum( - routed_process_dp_size(config.backend, process) - for process in processes - if process.endpoint_mode in {"decode", "agg"} and process.http_port > 0 + expansion for process in processes if process.endpoint_mode in {"decode", "agg"} and process.http_port > 0 ) return n_prefill, n_decode, f"{n_prefill}P + {n_decode}D Router workers; logical workers: {worker_desc}" diff --git a/tests/test_atom_atomesh.py b/tests/test_atom_atomesh.py new file mode 100644 index 000000000..937800b6e --- /dev/null +++ b/tests/test_atom_atomesh.py @@ -0,0 +1,262 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 SemiAnalysis LLC. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native ATOM backend and AToMesh frontend contracts.""" + +import json +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import yaml +from marshmallow import ValidationError + +from srtctl.backends import AtomProtocol, AtomServerConfig +from srtctl.cli.do_sweep import SweepOrchestrator +from srtctl.core.config import load_config, resolve_config_with_defaults +from srtctl.core.runtime import Nodes, RuntimeContext +from srtctl.core.schema import SrtConfig +from srtctl.core.topology import Process +from srtctl.frontends import AtomeshFrontend + +WORKER_IP = "10.0.0.20" + + +def _config() -> dict: + return { + "schema": 2, + "engine": "atom", + "name": "atom-atomesh", + "model": { + "path": "hf:Qwen/Qwen3-0.6B", + "container": "rocm/atom:latest", + "precision": "bf16", + }, + "resources": { + "gpu_type": "mi300x", + "gpus_per_node": 8, + }, + "roles": { + "prefill": {"nodes": 1, "workers": 1}, + "decode": {"nodes": 1, "workers": 1, "args": {"gpu-memory-utilization": 0.9}}, + }, + "frontend": {"type": "atomesh", "enable_multiple_frontends": False}, + } + + +def _load(data: dict) -> SrtConfig: + return SrtConfig.Schema().load(resolve_config_with_defaults(data, None)) + + +def _runtime() -> SimpleNamespace: + return SimpleNamespace(worker_model_arg="/model", network_interface=None) + + +def _build(backend: AtomProtocol, process: Process, runtime=None, **kwargs) -> list[str]: + with patch("srtctl.core.slurm.get_hostname_ip", return_value=WORKER_IP): + return backend.build_worker_command(process, [process], runtime or _runtime(), **kwargs) + + +def test_atomesh_frontend_requires_atom_backend() -> None: + data = _config() + data["engine"] = "sglang" + + with pytest.raises(ValidationError, match="frontend.type: atomesh requires backend.type: atom"): + _load(data) + + +def test_atom_worker_requires_atomesh_frontend() -> None: + process = Process("node0", frozenset(range(8)), 7500, 6100, "agg", 0) + + with pytest.raises(ValueError, match="requires frontend.type: atomesh"): + AtomProtocol().build_worker_command(process, [process], _runtime(), frontend_type="dynamo") + + +def test_v2_atom_roles_build_prefill_decode_workers(tmp_path: Path) -> None: + """Load a v2 recipe through the real orchestrator topology path.""" + runtime = SimpleNamespace( + nodes=Nodes(head="node0", bench="node0", infra="node0", worker=("node0", "node1")), + worker_model_arg="Qwen/Qwen3-0.6B", + network_interface="hsn0", + ) + path = tmp_path / "recipe.yaml" + path.write_text(yaml.safe_dump(_config())) + config = load_config(path) + orchestrator = SweepOrchestrator(config=config, runtime=runtime) + + with patch("srtctl.core.slurm.get_hostname_ip", return_value=WORKER_IP): + processes = orchestrator.backend_processes + launches = [config.backend.build_worker_command(process, [process], runtime) for process in processes] + + assert [(process.node, process.endpoint_mode) for process in processes] == [ + ("node0", "prefill"), + ("node1", "decode"), + ] + prefill, decode = launches + assert json.loads(prefill[prefill.index("--kv-transfer-config") + 1])["kv_role"] == "kv_producer" + assert decode[decode.index("--gpu-memory-utilization") + 1] == "0.9" + assert json.loads(decode[decode.index("--kv-transfer-config") + 1])["kv_role"] == "kv_consumer" + + +@pytest.mark.parametrize("layout", ["hf", "mounted", "staged"]) +def test_atom_served_name_matches_worker_model_argument(tmp_path: Path, layout: str) -> None: + """ATOM advertises its literal --model, so served_model_name must track the runtime's worker_model_arg.""" + data = _config() + stage_dir = tmp_path / "scratch" / "models" + if layout == "hf": + data["model"]["path"] = "hf:deepseek-ai/DeepSeek-V4-Pro" + else: + model_dir = tmp_path / "DeepSeek-V4-Pro" + model_dir.mkdir() + data["model"]["path"] = str(model_dir) + if layout == "staged": + data["model"]["stage_dir"] = str(stage_dir) + config = _load(data) + nodes = Nodes(head="node0", bench="node0", infra="node0", worker=("node0",)) + + with ( + patch("srtctl.core.runtime.Nodes.from_slurm", return_value=nodes), + patch("srtctl.core.runtime.get_srtslurm_setting", side_effect=lambda name, default=None: default), + patch("srtctl.core.runtime.get_hostname_ip", return_value=WORKER_IP), + ): + runtime = RuntimeContext.from_config(config, job_id="42", log_dir_base=tmp_path / "outputs") + process = Process("node0", frozenset(range(8)), 7500, 6100, "prefill", 0, nixl_port=5400) + command = _build(config.backend, process, runtime) + + expected = { + "hf": "deepseek-ai/DeepSeek-V4-Pro", + "mounted": "/model", + "staged": str(stage_dir / "DeepSeek-V4-Pro"), + }[layout] + assert command[command.index("--model") + 1] == expected + assert config.served_model_name == expected + + +def test_atom_builds_native_aggregate_command() -> None: + """Recipe flags keep ATOM's mixed hyphen/underscore spelling and follow the managed arguments.""" + backend = AtomProtocol( + atom_config=AtomServerConfig( + aggregated={ + "trust-remote-code": True, + "gpu-memory-utilization": 0.9, + "kv_cache_dtype": "fp8", + "no-enable_prefix_caching": True, + "disable-log-stats": False, + } + ) + ) + process = Process("node0", frozenset(range(8)), 7500, 6100, "agg", 0, nixl_port=5400) + + command = _build(backend, process, SimpleNamespace(worker_model_arg="/model", network_interface="hsn0")) + + assert command == [ + "env", + f"ATOM_HOST_IP={WORKER_IP}", + "python3", + "-m", + "atom.entrypoints.openai_server", + "--model", + "/model", + "--host", + "0.0.0.0", + "--server-port", + "6100", + "-tp", + "8", + "--gpu-memory-utilization", + "0.9", + "--kv_cache_dtype", + "fp8", + "--no-enable_prefix_caching", + "--trust-remote-code", + ] + + +@pytest.mark.parametrize("key", ["tensor_parallel_size", "--server-port", "model"]) +def test_atom_rejects_recipe_overrides_of_managed_arguments(key: str) -> None: + """Reserved flags are matched after normalizing dashes and underscores.""" + backend = AtomProtocol(atom_config=AtomServerConfig(aggregated={key: 4})) + process = Process("node0", frozenset(range(8)), 7500, 6100, "agg", 0) + + with pytest.raises(ValueError, match="srtctl-managed argument"): + _build(backend, process) + + +@pytest.mark.parametrize(("protocol", "extra"), [(None, {}), ("tcp", {"protocol": "tcp"})]) +def test_atom_pd_worker_emits_mooncake_kv_transfer_config(protocol: str | None, extra: dict) -> None: + backend = AtomProtocol(mooncake_protocol=protocol) + process = Process("node0", frozenset(range(4)), 7500, 6100, "prefill", 0, nixl_port=6301) + + command = _build(backend, process) + + payload = json.loads(command[command.index("--kv-transfer-config") + 1]) + assert payload == { + "kv_role": "kv_producer", + "kv_connector": "mooncake", + "proxy_ip": WORKER_IP, + "handshake_port": 6301, + **extra, + } + + +def test_atom_rejects_cross_node_model_parallel_endpoint() -> None: + """ATOM cannot coordinate one logical worker across two Slurm nodes.""" + leader = Process("node0", frozenset(range(4)), 7500, 6100, "prefill", 0, nixl_port=6301) + peer = Process("node1", frozenset(range(4)), 7501, 6101, "prefill", 0, nixl_port=6302) + + with pytest.raises(ValueError, match="fit on one Slurm node"): + AtomProtocol().build_worker_command(leader, [leader, peer], _runtime()) + + +def test_atom_pd_worker_requires_mooncake_handshake_port() -> None: + """Fail before launch when the allocated P/D endpoint has no transfer port.""" + process = Process("node0", frozenset(range(4)), 7500, 6100, "decode", 0) + + with pytest.raises(ValueError, match="missing its Mooncake handshake port"): + _build(AtomProtocol(), process) + + +def test_atomesh_builds_static_pd_command_without_bootstrap_ports() -> None: + """ATOM publishes transfer topology itself, so no handshake port follows --prefill.""" + frontend = AtomeshFrontend() + backend = AtomProtocol() + processes = [ + Process("node0", frozenset(range(8)), 7500, 6100, "prefill", 0, nixl_port=6301), + Process("node1", frozenset(range(8)), 7500, 6101, "decode", 0, nixl_port=6302), + ] + config = SimpleNamespace(frontend=SimpleNamespace(args={})) + + with patch("srtctl.frontends.static_router.get_hostname_ip", side_effect=["10.0.0.20", "10.0.0.21"]): + workers = frontend.collect_workers(backend, processes, "fabric0") + command = frontend.build_router_command(workers, "0.0.0.0", 8000, backend) + command.extend(frontend.get_managed_frontend_args(config, backend, processes)) + + assert command == [ + "atomesh", + "launch", + "--pd-disaggregation", + "--prefill", + "http://10.0.0.20:6100", + "--decode", + "http://10.0.0.21:6101", + "--host", + "0.0.0.0", + "--port", + "8000", + "--backend", + "atom", + ] + + +def test_atomesh_rejects_recipe_backend_argument() -> None: + config = SimpleNamespace(frontend=SimpleNamespace(args={"backend": "vllm"})) + + with pytest.raises(ValueError, match="managed by srtctl"): + AtomeshFrontend().get_managed_frontend_args(config, AtomProtocol(), []) + + +def test_atomesh_worker_metrics_are_not_scraped() -> None: + """The frontend owns scrape targets; ATOM has no worker metrics endpoint.""" + process = Process("node0", frozenset(range(8)), 7500, 6100, "agg", 0) + assert AtomeshFrontend().worker_metrics_port(process, _runtime()) is None diff --git a/tests/test_benchmarks.py b/tests/test_benchmarks.py index 559ef6ead..cb5b26c75 100644 --- a/tests/test_benchmarks.py +++ b/tests/test_benchmarks.py @@ -3,6 +3,8 @@ """Tests for benchmark runners.""" +from pathlib import Path + import pytest from srtctl.benchmarks import get_runner, list_benchmarks @@ -1857,6 +1859,34 @@ def capture_srun(**kwargs): assert env_to_set["MODEL"] == "test-model" assert env_to_set["MODEL_NAME"] == "test-model" + def test_benchmark_env_passthrough(self): + """Eval-only substitution preserves the configured benchmark env.""" + import os + import threading + from unittest.mock import MagicMock, patch + + orch = self._make_orchestrator() + orch.config.benchmark.env["SRTCTL_LM_EVAL_RESULT_DIR"] = "/results/{job_id}/eval" + stop = threading.Event() + + mock_proc = MagicMock() + mock_proc.poll.return_value = 0 + mock_proc.returncode = 0 + captured_kwargs = {} + + def capture_srun(**kwargs): + captured_kwargs.update(kwargs) + return mock_proc + + with ( + patch.dict(os.environ, {"EVAL_ONLY": "false"}, clear=False), + patch("srtctl.cli.do_sweep.wait_for_port", return_value=True), + patch("srtctl.cli.do_sweep.start_srun_process", side_effect=capture_srun), + ): + orch._run_post_eval(stop) + + assert captured_kwargs["env_to_set"]["SRTCTL_LM_EVAL_RESULT_DIR"] == "/results/12345/eval" + def test_eval_conc_from_env(self): """EVAL_CONC from env takes priority over benchmark concurrencies.""" import os @@ -1942,6 +1972,43 @@ class TestSweepRunEvalIntegration: def _make_orchestrator(): return TestRunPostEval._make_orchestrator() + @pytest.mark.parametrize(("frontend_type", "needs_infra"), [("atomesh", False), ("dynamo", True)]) + def test_only_dynamo_starts_head_infrastructure( + self, frontend_type: str, needs_infra: bool, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + """An ATOM/AToMesh run reaches the benchmark without a discovery plane; Dynamo still starts it.""" + from dataclasses import replace + from unittest.mock import MagicMock, patch + + from srtctl.backends import AtomProtocol + from srtctl.core.schema import FrontendConfig + + orch = self._make_orchestrator() + orch.config = replace( + orch.config, + frontend=FrontendConfig(type=frontend_type), + backend=AtomProtocol() if frontend_type == "atomesh" else orch.config.backend, + ) + orch.runtime = replace(orch.runtime, log_dir=tmp_path / "logs") + orch.runtime.log_dir.mkdir() + monkeypatch.setenv("EVAL_ONLY", "false") + monkeypatch.setenv("RUN_EVAL", "false") + with ( + patch.object(orch, "start_head_infrastructure", return_value=MagicMock()) as head, + patch.object(orch, "start_all_workers", return_value={}), + patch.object(orch, "start_frontend", return_value=[]), + patch.object(orch, "run_benchmark", return_value=0) as benchmark, + patch.object(orch, "run_postprocess"), + patch("srtctl.cli.do_sweep.StatusReporter"), + ): + assert orch.run() == 0 + + benchmark.assert_called_once() + if needs_infra: + head.assert_called_once() + else: + head.assert_not_called() + def test_run_eval_only_mode(self): """EVAL_ONLY=true skips benchmark and runs _run_post_eval.""" import os diff --git a/tests/test_configs.py b/tests/test_configs.py index 267e55aea..c5885902b 100644 --- a/tests/test_configs.py +++ b/tests/test_configs.py @@ -566,6 +566,41 @@ def test_get_environment_for_mode(self): assert config.get_environment_for_mode("decode") == {"DECODE_VAR": "1"} assert config.get_environment_for_mode("agg") == {} + @pytest.mark.parametrize( + ("mode", "frontend_type", "frontend_args", "dp_size", "expected"), + [ + ("prefill", "sglang-router", {"dp-aware": True}, 8, True), + ("decode", "sglang-router", {"dp_aware": True}, 8, True), + ("prefill", "sglang-router", {"dp-aware": True}, 1, False), + ("prefill", "sglang-router", {}, 8, False), + ("prefill", "dynamo", {"dp-aware": True}, 8, False), + ("agg", "sglang-router", {"dp-aware": True}, 8, False), + ], + ) + def test_dp_aware_router_bootstrap_rank_environment( + self, + mode, + frontend_type, + frontend_args, + dp_size, + expected, + ): + config = SGLangProtocol( + sglang_config=SGLangServerConfig( + prefill={"dp-size": dp_size}, + decode={"dp-size": dp_size}, + aggregated={"dp-size": dp_size}, + ) + ) + + environment = config.get_frontend_integration_environment( + mode, + frontend_type, + frontend_args, + ) + + assert (environment.get("SGLANG_DISAGGREGATION_FORCE_QUERY_PREFILL_DP_RANK") == "1") is expected + def test_kv_events_config_global_bool(self): """Test kv_events_config=True enables prefill+decode+aggregated with defaults.""" config = SGLangProtocol(kv_events_config=True) @@ -678,9 +713,11 @@ def test_worker_command_passes_the_allocated_nccl_port(self): runtime = MagicMock() runtime.model_path = Path("/model") runtime.is_hf_model = False - with patch("srtctl.core.slurm.get_hostname_ip", return_value="10.0.0.1"): + runtime.network_interface = "management0" + with patch("srtctl.core.slurm.get_hostname_ip", return_value="10.0.0.1") as resolve_ip: command = backend.build_worker_command(processes[1], [processes[1]], runtime) + resolve_ip.assert_called_once_with("node0", "management0") assert command[command.index("--nccl-port") + 1] == str(SGLANG_NCCL_PORT_BASE + 1) @@ -3014,7 +3051,7 @@ def test_direct_vllm_command_preserves_current_main_device_binding(self): assert "dynamo.vllm" not in cmd def test_vllm_router_can_use_environment_device_binding(self): - """Stable vLLM builds can avoid the newer --device-ids CLI.""" + """set_visible_devices swaps --device-ids for the worker-stage environment mask.""" from pathlib import Path from unittest.mock import MagicMock, patch @@ -3049,7 +3086,6 @@ def test_vllm_router_can_use_environment_device_binding(self): assert cmd[:3] == ["vllm", "serve", "/model"] assert "--device-ids" not in cmd - assert backend.should_set_visible_devices() @pytest.mark.parametrize( ("mode", "role"), diff --git a/tests/test_frontends.py b/tests/test_frontends.py index 6d8aa8fb1..acd4363a2 100644 --- a/tests/test_frontends.py +++ b/tests/test_frontends.py @@ -67,6 +67,7 @@ class TestFrontendRegistry: def test_registry_lists_every_frontend_type(self): assert list_frontend_types() == [ + "atomesh", "dynamo", "none", "sglang", @@ -206,7 +207,7 @@ def test_schema_rejects_unknown_type_at_load(self): from srtctl.backends import SGLangProtocol from srtctl.core.schema import FrontendConfig, ResourceConfig, SrtConfig - with pytest.raises(ValidationError, match="Unknown frontend.type 'toy-router'.*Available: dynamo, none"): + with pytest.raises(ValidationError, match="Unknown frontend.type 'toy-router'.*Available: atomesh, dynamo, none"): SrtConfig( name="toy", model={"path": "model", "container": "image", "precision": "fp8"}, @@ -595,7 +596,9 @@ def test_grpc_scheme_when_enabled(self, mock_get_ip, mock_srun): @patch("srtctl.frontends.sglang.get_hostname_ip") def test_disaggregated_mode_command(self, mock_get_ip, mock_srun): """Disaggregated mode uses --pd-disaggregation with --prefill and --decode.""" - mock_get_ip.side_effect = lambda node, interface=None: f"10.0.0.{node[-1]}" + mock_get_ip.side_effect = lambda node, interface=None: ( + f"10.0.0.{node[-1]}" if interface == "eth0" else f"192.168.0.{node[-1]}" + ) mock_srun.return_value = MagicMock() frontend = SGLangRouterFrontend() @@ -609,6 +612,7 @@ def test_disaggregated_mode_command(self, mock_get_ip, mock_srun): backend.is_grpc_mode.return_value = False runtime = MagicMock() + runtime.network_interface = "eth0" runtime.log_dir = MagicMock() runtime.log_dir.__truediv__ = lambda self, x: f"/logs/{x}" runtime.container_image = "/container.sqsh" @@ -631,6 +635,8 @@ def test_disaggregated_mode_command(self, mock_get_ip, mock_srun): assert "--decode" in cmd # Bootstrap port should be included assert "30001" in cmd + assert "http://10.0.0.1:30000" in cmd + assert "http://10.0.0.2:30000" in cmd @patch("srtctl.frontends.sglang.start_srun_process") @patch("srtctl.frontends.sglang.get_hostname_ip") diff --git a/tests/test_port_allocator.py b/tests/test_port_allocator.py index e2616f69c..08c16a8f7 100644 --- a/tests/test_port_allocator.py +++ b/tests/test_port_allocator.py @@ -31,7 +31,7 @@ PortKind, ) -TOPOLOGY_EXAMPLE_DIRS = ("examples/sglang", "examples/vllm", "examples/trtllm", "examples/mocker") +TOPOLOGY_EXAMPLE_DIRS = ("examples/atom", "examples/sglang", "examples/vllm", "examples/trtllm", "examples/mocker") class TestNodePortAllocator: diff --git a/tests/test_schema_docs.py b/tests/test_schema_docs.py index fd77c2044..e9386df5b 100644 --- a/tests/test_schema_docs.py +++ b/tests/test_schema_docs.py @@ -15,6 +15,8 @@ BACKEND_TYPES, DEFAULT_LEGACY_OUTPUT, DEFAULT_OUTPUT, + INTERNAL_CLASSES, + INTERNAL_FIELDS, LEGACY_CLASSES, LEGACY_FIELDS, LEGACY_TOP_LEVEL, @@ -64,11 +66,11 @@ def test_schema_reference_documents_only_the_2_0_layout() -> None: recipe_table = text[text.index("## Recipe\n") : text.index("## Authoring surface")] for key in LEGACY_TOP_LEVEL: assert f"| `{key}` |" not in recipe_table, f"legacy top-level key {key} leaked into schema-reference.md" - for cls, mapping in LEGACY_FIELDS.items(): + for cls, mapping in (LEGACY_FIELDS | INTERNAL_FIELDS).items(): section = _section(text, cls.__name__) for key in mapping: assert f"| `{key}` |" not in section, f"legacy key {cls.__name__}.{key} leaked into schema-reference.md" - for cls in LEGACY_CLASSES: + for cls in LEGACY_CLASSES | INTERNAL_CLASSES: assert f"### {cls.__name__}" not in text, f"legacy class {cls.__name__} leaked into schema-reference.md" for needle in ( "## Authoring surface", @@ -77,6 +79,7 @@ def test_schema_reference_documents_only_the_2_0_layout() -> None: "### placement", "`colocate`", "## Engine types", + "`engine.type: atom`", "`engine.type: sglang`", "## Cluster config", "[legacy-v1.md](legacy-v1.md)", @@ -100,8 +103,8 @@ def test_legacy_reference_documents_every_v1_key() -> None: assert f"### {cls.__name__}" in text, cls.__name__ for needle in ("## v1 keys and what replaced them", "## backend", "## infra", "srtctl migrate"): assert needle in text, needle - for type_name, _ in BACKEND_TYPES: - assert f"`backend.type: {type_name}`" in text + for type_name, cls in BACKEND_TYPES: + assert (f"`backend.type: {type_name}`" in text) == (cls in LEGACY_FIELDS) def test_mooncake_device_mapping_is_documented_as_a_v2_service_option() -> None: diff --git a/tests/test_slurm.py b/tests/test_slurm.py index 23d2f4499..feca2c455 100644 --- a/tests/test_slurm.py +++ b/tests/test_slurm.py @@ -180,6 +180,7 @@ def test_worker_stage_wraps_nonfatal_fingerprint_hook(tmp_path: Path) -> None: backend.build_worker_command.return_value = ["python3", "-m", "worker"] backend.get_environment_for_mode.return_value = {} backend.get_process_environment.return_value = {} + backend.get_frontend_integration_environment.return_value = {} backend.type = "vllm" backend.failover = None backend.mooncake_kv_store = None @@ -187,7 +188,7 @@ def test_worker_stage_wraps_nonfatal_fingerprint_hook(tmp_path: Path) -> None: mixin = WorkerStageMixin() mixin.config = SimpleNamespace( setup_script="setup.sh", - frontend=SimpleNamespace(type="sglang"), + frontend=SimpleNamespace(type="sglang", args={}), dynamo=SimpleNamespace(install=False, sidecar=False, request_plane="nats", event_plane="zmq"), observability=ObservabilityConfig(), profiling=SimpleNamespace(enabled=False, is_nsys=False), @@ -243,13 +244,14 @@ def _remap_worker_mixin(tmp_path: Path, *, frontend_type: str, dynamo_install: b backend.build_worker_command.return_value = ["python3", "-m", "worker"] backend.get_environment_for_mode.return_value = {} backend.get_process_environment.return_value = {} + backend.get_frontend_integration_environment.return_value = {} backend.failover = None backend.mooncake_kv_store = None mixin = WorkerStageMixin() mixin.config = SimpleNamespace( setup_script=None, - frontend=SimpleNamespace(type=frontend_type), + frontend=SimpleNamespace(type=frontend_type, args={}), dynamo=SimpleNamespace( install=dynamo_install, sidecar=False, @@ -428,6 +430,29 @@ def test_worker_stage_no_remap_root_when_dynamo_install_false(tmp_path: Path) -> assert mock_srun.call_args.kwargs["srun_export_env"] is None +@pytest.mark.parametrize(("recipe_value", "expected"), [(None, "1"), ("0", "0")]) +def test_worker_stage_applies_frontend_integration_environment( + tmp_path: Path, + recipe_value: str | None, + expected: str, +) -> None: + mixin, process = _remap_worker_mixin(tmp_path, frontend_type="sglang", dynamo_install=False) + mixin.config.frontend.args = {"dp-aware": True} + integration_key = "SGLANG_DISAGGREGATION_FORCE_QUERY_PREFILL_DP_RANK" + mixin.backend.get_frontend_integration_environment.return_value = {integration_key: "1"} + mixin.backend.get_environment_for_mode.return_value = ( + {} if recipe_value is None else {integration_key: recipe_value} + ) + + with ( + patch("srtctl.cli.mixins.worker_stage.generate_capture_script", return_value="fingerprint || true"), + patch("srtctl.cli.mixins.worker_stage.start_srun_process", return_value=MagicMock()) as mock_srun, + ): + mixin.start_worker(process, [process]) + + assert mock_srun.call_args.kwargs["env_to_set"][integration_key] == expected + + # ---- Event-plane propagation (DYN_EVENT_PLANE) ---- @@ -678,13 +703,14 @@ def test_worker_stage_unsets_vllm_port_for_multinode_endpoint(tmp_path: Path) -> backend.build_worker_command.return_value = ["python3", "-m", "worker"] backend.get_environment_for_mode.return_value = {} backend.get_process_environment.return_value = {} + backend.get_frontend_integration_environment.return_value = {} backend.failover = None backend.mooncake_kv_store = None mixin = WorkerStageMixin() mixin.config = SimpleNamespace( setup_script=None, - frontend=SimpleNamespace(type="sglang"), + frontend=SimpleNamespace(type="sglang", args={}), dynamo=SimpleNamespace(install=False, sidecar=False, request_plane="nats", event_plane=None), observability=ObservabilityConfig(), profiling=SimpleNamespace(enabled=False, is_nsys=False), diff --git a/tests/test_vllm_router_frontend.py b/tests/test_vllm_router_frontend.py index 2b4223495..76f1e0311 100644 --- a/tests/test_vllm_router_frontend.py +++ b/tests/test_vllm_router_frontend.py @@ -219,6 +219,69 @@ def test_dep4_expansion_and_health_counts_follow_upstream_per_node_topology() -> ) +def test_multinode_dep8_routes_node_local_hybrid_pools_without_rank_reexpansion() -> None: + """A later node owns global ranks 4..7, not another local 0..3 namespace.""" + backend = VLLMProtocol( + vllm_config=VLLMServerConfig( + prefill={"data-parallel-size": 8}, + decode={"data-parallel-size": 8}, + ) + ) + processes = [ + Process("p0", frozenset(range(4)), 7500, 6100, "prefill", 0, node_rank=0), + Process("p1", frozenset(range(4)), 7501, 6100, "prefill", 0, node_rank=4), + Process("d0", frozenset(range(4)), 7502, 6100, "decode", 0, node_rank=0), + Process("d1", frozenset(range(4)), 7503, 6100, "decode", 0, node_rank=4), + ] + config = SimpleNamespace( + frontend=SimpleNamespace(type="vllm-router"), + backend=backend, + resources=SimpleNamespace(num_prefill=1, num_decode=1, num_agg=0), + ) + + assert node_local_data_parallel_size(backend, processes) == 1 + assert _get_health_expectations(config, processes) == ( + 2, + 2, + "2P + 2D Router workers; logical workers: 1P + 1D", + 4, + ) + + +@pytest.mark.parametrize("expansion", [1, 4]) +def test_multinode_hybrid_pool_schema_matches_router_expansion(expansion: int) -> None: + """The declared expansion must match the unexpanded node-local pool URLs.""" + data = { + "name": "hybrid-pools", + "model": {"path": "/model", "container": "vllm", "precision": "bf16"}, + "resources": { + "gpus_per_node": 4, + "prefill_nodes": 2, + "decode_nodes": 2, + "prefill_workers": 1, + "decode_workers": 1, + }, + "backend": { + "type": "vllm", + "vllm_config": { + "prefill": {"data-parallel-size": 8}, + "decode": {"data-parallel-size": 8}, + }, + }, + "frontend": { + "type": "vllm-router", + "enable_multiple_frontends": False, + "args": {"intra-node-data-parallel-size": expansion}, + }, + } + if expansion == 1: + config = SrtConfig.Schema().load(data) + assert config.frontend.args["intra-node-data-parallel-size"] == 1 + else: + with pytest.raises(ValidationError, match="conflicts with the allocated vLLM topology"): + SrtConfig.Schema().load(data) + + def test_cross_node_model_parallel_base_is_not_dp_expanded() -> None: backend = VLLMProtocol( vllm_config=VLLMServerConfig(aggregated={"data-parallel-size": 2, "tensor-parallel-size": 8}) From 8242106eaa41782c81906a9b45c294b39b46ee00 Mon Sep 17 00:00:00 2001 From: adibarra <93070681+adibarra@users.noreply.github.com> Date: Tue, 22 Sep 2026 15:13:14 -0500 Subject: [PATCH 2/3] feat(atom): support direct aggregate serving MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 新增 ATOM 单聚合 worker 的直接服务模式，复用 SGLang 生命周期及原生端口接口。 Signed-off-by: adibarra <93070681+adibarra@users.noreply.github.com> --- docs/config-reference.md | 8 +- examples/atom/direct.yaml | 25 +++++ src/srtctl/backends/atom.py | 11 +- src/srtctl/frontends/__init__.py | 2 + src/srtctl/frontends/atom.py | 28 +++++ src/srtctl/frontends/direct.py | 131 ++++++++++++++++++++++ src/srtctl/frontends/sglang_direct.py | 151 ++------------------------ tests/test_atom_atomesh.py | 4 +- tests/test_atom_direct_frontend.py | 90 +++++++++++++++ tests/test_frontends.py | 3 +- 10 files changed, 302 insertions(+), 151 deletions(-) create mode 100644 examples/atom/direct.yaml create mode 100644 src/srtctl/frontends/atom.py create mode 100644 src/srtctl/frontends/direct.py create mode 100644 tests/test_atom_direct_frontend.py diff --git a/docs/config-reference.md b/docs/config-reference.md index dfac865c4..fe635c8b2 100644 --- a/docs/config-reference.md +++ b/docs/config-reference.md @@ -43,7 +43,13 @@ This page is the prose guide: what each block means, how the pieces interact, an ## Overview -### ATOM with AToMesh +### ATOM direct serving and AToMesh + +For one aggregate worker, use `engine: atom` with `frontend.type: atom` and +`frontend.enable_multiple_frontends: false`. The ATOM worker binds the public +port directly; no router container or AToMesh installation is needed. Readiness +requires `/health` to succeed and `/v1/models` to list a model. See the +[direct ATOM recipe](../examples/atom/direct.yaml). Use `engine: atom` with `frontend.type: atomesh` to launch native `atom.entrypoints.openai_server` workers and the official AToMesh router. Both diff --git a/examples/atom/direct.yaml b/examples/atom/direct.yaml new file mode 100644 index 000000000..ee551ccf2 --- /dev/null +++ b/examples/atom/direct.yaml @@ -0,0 +1,25 @@ +# One ATOM aggregate worker owns the public endpoint; no router is launched. +# Configure visible_devices_env: ROCR_VISIBLE_DEVICES in srtslurm.yaml. +schema: 2 +name: qwen3-0.6b-atom-direct +model: + path: hf:Qwen/Qwen3-0.6B + container: rocm/atom:latest + precision: bf16 +resources: + gpu_type: mi300x + gpus_per_node: 8 +engine: atom +frontend: + type: atom + enable_multiple_frontends: false +roles: + agg: + nodes: 1 + workers: 1 + gpus: 1 +benchmark: + type: sa-bench + isl: 128 + osl: 128 + concurrencies: "4" diff --git a/src/srtctl/backends/atom.py b/src/srtctl/backends/atom.py index 235ccb8e0..f295d37cb 100644 --- a/src/srtctl/backends/atom.py +++ b/src/srtctl/backends/atom.py @@ -168,8 +168,13 @@ def build_worker_command( dump_config_path: Path | None = None, profiling: ProfilingConfig | None = None, ) -> list[str]: - if frontend_type != "atomesh": - raise ValueError(f"backend.type: atom requires frontend.type: atomesh (got {frontend_type!r})") + from srtctl.frontends import get_frontend + + frontend = get_frontend(frontend_type) + if frontend.required_backend != self.type: + raise ValueError(f"backend.type: atom requires an ATOM frontend (got {frontend_type!r})") + public_endpoint = frontend.worker_api_port(process.endpoint_mode) == "public" + port = runtime.frontend_port if public_endpoint else process.http_port if len({item.node for item in endpoint_processes}) != 1: raise ValueError("ATOM currently requires each logical endpoint to fit on one Slurm node") @@ -193,7 +198,7 @@ def build_worker_command( "--host", "0.0.0.0", "--server-port", - str(process.http_port), + str(port), "-tp", str(len(process.gpu_indices)), ] diff --git a/src/srtctl/frontends/__init__.py b/src/srtctl/frontends/__init__.py index b5e462cf5..5b9396d1e 100644 --- a/src/srtctl/frontends/__init__.py +++ b/src/srtctl/frontends/__init__.py @@ -17,6 +17,7 @@ - vllm-router: Official vLLM Router with static aggregate or P/D workers """ +from srtctl.frontends.atom import AtomFrontend from srtctl.frontends.atomesh import AtomeshFrontend from srtctl.frontends.base import ( FRONTEND_NONE, @@ -35,6 +36,7 @@ __all__ = [ "FRONTEND_NONE", + "AtomFrontend", "AtomeshFrontend", "DynamicFrontend", "DynamoFrontend", diff --git a/src/srtctl/frontends/atom.py b/src/srtctl/frontends/atom.py new file mode 100644 index 000000000..a5ee260d6 --- /dev/null +++ b/src/srtctl/frontends/atom.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 SemiAnalysis LLC. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Direct ATOM OpenAI server for a single aggregate worker.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from srtctl.frontends.base import register_frontend +from srtctl.frontends.direct import DirectServerFrontend + +if TYPE_CHECKING: + from srtctl.core.runtime import RuntimeContext + from srtctl.core.topology import Process + + +@register_frontend("atom") +class AtomFrontend(DirectServerFrontend): + required_backend = "atom" + server_name = "atom.entrypoints.openai_server" + router_hint = "Use frontend.type: atomesh for replicas or prefill/decode." + + def worker_metrics_port(self, process: Process, runtime: RuntimeContext) -> None: + return None + + def profiling_control_port(self, process: Process, config: Any, runtime: RuntimeContext) -> None: + return None diff --git a/src/srtctl/frontends/direct.py b/src/srtctl/frontends/direct.py new file mode 100644 index 000000000..dbebdb91b --- /dev/null +++ b/src/srtctl/frontends/direct.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared lifecycle for a single direct OpenAI server.""" + +from __future__ import annotations + +import logging +import threading +from typing import TYPE_CHECKING, Any, ClassVar, Literal + +from srtctl.core.health import WorkerHealthResult, probe_direct_server +from srtctl.frontends.base import agg_leader_nodes, logical_health_expectations + +if TYPE_CHECKING: + from srtctl.core.processes import ManagedProcess + from srtctl.core.runtime import RuntimeContext + from srtctl.core.topology import Process + from srtctl.services.implicit import EffectiveService + +logger = logging.getLogger(__name__) + + +class DirectServerFrontend: + """One aggregate worker owns the public OpenAI endpoint without a router.""" + + required_backend: ClassVar[str] + server_name: ClassVar[str] + router_hint: ClassVar[str] + worker_launch: ClassVar[Literal["dynamo", "direct"]] = "direct" + expands_node_local_dp: ClassVar[bool] = False + + @property + def type(self) -> str: + return self.required_backend + + def worker_api_port(self, mode: str) -> Literal["public", "allocated"]: + """The one server is the endpoint, so it binds the public port.""" + return "public" + + metrics_path: ClassVar[str] = "/metrics" + + def worker_metrics_port(self, process: Process, runtime: RuntimeContext) -> int | None: + """The aggregate leader binds the public port; its followers serve nothing.""" + if process.endpoint_mode == "agg" and process.is_leader: + return runtime.frontend_port + return None + + def worker_endpoint_port(self, process: Process, config: Any, runtime: RuntimeContext) -> int | None: + return runtime.frontend_port if process.is_leader else None + + def profiling_control_port(self, process: Process, config: Any, runtime: RuntimeContext) -> int | None: + """The leader's server on the public port carries the control routes; followers have none.""" + return runtime.frontend_port if process.is_leader else None + + def profiling_control_is_leader_only(self, config: Any) -> bool: + return False + + def direct_endpoint_nodes(self, processes: list[Process]) -> list[str]: + return agg_leader_nodes(processes) + + def worker_ready_port(self, process: Process) -> int: + return process.sys_port + + def probe_ready( + self, host: str, port: int, expected_prefill: int, expected_decode: int, config: Any + ) -> WorkerHealthResult: + """The worker's own /health, then /v1/models must list the model.""" + return probe_direct_server(host, port) + + def health_expectations(self, config: Any, processes: list[Process] | None) -> tuple[int, int, str]: + return logical_health_expectations(config) + + def validate(self, config: Any) -> None: + """Reject layouts that need a router before allocating a job.""" + prefix = f"frontend.type: {self.type}" + if config.frontend.enable_multiple_frontends: + raise ValueError( + f"{prefix} binds {self.server_name} directly; set frontend.enable_multiple_frontends: false" + ) + if config.resources.is_disaggregated: + raise ValueError( + f"{prefix} supports one aggregate worker only, not a prefill/decode layout. {self.router_hint}" + ) + if config.resources.num_agg != 1: + raise ValueError( + f"{prefix} supports exactly one aggregate worker, got {config.resources.num_agg}. {self.router_hint}" + ) + if config.dynamo.sidecar: + raise ValueError(f"{prefix} does not support dynamo.sidecar; use frontend.type: dynamo") + + def get_backend_health_urls( + self, + backend: Any, + backend_processes: list[Process], + network_interface: str | None = None, + ) -> list[str]: + return [] + + def implied_services(self, config: Any) -> list[EffectiveService]: + return [] + + def frontend_metrics_port(self, frontend_args: dict[str, Any] | None) -> int | None: + return None + + def start_frontends( + self, + topology: Any, + runtime: RuntimeContext, + config: Any, + backend: Any, + backend_processes: list[Process], + stop_event: threading.Event | None = None, + ) -> list[ManagedProcess]: + if config.backend.type != self.required_backend: + raise ValueError( + f"frontend.type: {self.type} requires engine {self.required_backend} (got {config.backend.type!r})" + ) + self.validate(config) + if topology.uses_nginx or len(topology.frontend_nodes) != 1: + raise ValueError( + f"frontend.type: {self.type} binds {self.server_name} directly to the public port; " + "set frontend.enable_multiple_frontends: false" + ) + logger.info( + "frontend.type=%s: no separate frontend process; %s owns port %d", + self.type, + self.server_name, + topology.public_port, + ) + return [] diff --git a/src/srtctl/frontends/sglang_direct.py b/src/srtctl/frontends/sglang_direct.py index 64f308f6b..a22336f78 100644 --- a/src/srtctl/frontends/sglang_direct.py +++ b/src/srtctl/frontends/sglang_direct.py @@ -1,151 +1,14 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Direct SGLang frontend (`frontend.type: sglang`). +"""Direct SGLang server; use sglang-router for replicas or prefill/decode.""" -For a single aggregate SGLang worker the OpenAI-compatible HTTP server is the -worker process itself (`sglang.launch_server` bound to the public port). No -router process starts. Use `frontend.type: sglang-router` for several replicas -or a prefill/decode layout, and `frontend.type: dynamo` for KV-aware routing. -""" - -from __future__ import annotations - -import logging -import threading -from typing import TYPE_CHECKING, Any, ClassVar, Literal - -from srtctl.core.health import WorkerHealthResult, probe_direct_server -from srtctl.frontends.base import agg_leader_nodes, logical_health_expectations, register_frontend - -if TYPE_CHECKING: - from srtctl.core.processes import ManagedProcess - from srtctl.core.runtime import RuntimeContext - from srtctl.core.topology import Process - from srtctl.services.implicit import EffectiveService - -logger = logging.getLogger(__name__) +from srtctl.frontends.base import register_frontend +from srtctl.frontends.direct import DirectServerFrontend @register_frontend("sglang") -class SGLangFrontend: - """Direct SGLang OpenAI server frontend. - - Intentionally narrow: exactly one aggregate worker, which binds the public - port itself. Readiness is the worker's ``/health`` plus ``/v1/models``. - """ - - required_backend: ClassVar[str | None] = "sglang" - worker_launch: ClassVar[Literal["dynamo", "direct"]] = "direct" - expands_node_local_dp: ClassVar[bool] = False - - @property - def type(self) -> str: - return "sglang" - - def worker_api_port(self, mode: str) -> Literal["public", "allocated"]: - """The one ``sglang.launch_server`` is the endpoint, so it binds the public port.""" - return "public" - - metrics_path: ClassVar[str] = "/metrics" - - def worker_metrics_port(self, process: Process, runtime: RuntimeContext) -> int | None: - """The aggregate leader binds the public port; its followers serve nothing.""" - if process.endpoint_mode == "agg" and process.is_leader: - return runtime.frontend_port - return None - - def worker_endpoint_port(self, process: Process, config: Any, runtime: RuntimeContext) -> int | None: - return runtime.frontend_port if process.is_leader else None - - def profiling_control_port(self, process: Process, config: Any, runtime: RuntimeContext) -> int | None: - """The leader's server on the public port carries the control routes; followers have none.""" - return runtime.frontend_port if process.is_leader else None - - def profiling_control_is_leader_only(self, config: Any) -> bool: - return False - - def direct_endpoint_nodes(self, processes: list[Process]) -> list[str]: - return agg_leader_nodes(processes) - - def worker_ready_port(self, process: Process) -> int: - return process.sys_port - - def probe_ready( - self, host: str, port: int, expected_prefill: int, expected_decode: int, config: Any - ) -> WorkerHealthResult: - """The worker's own /health, then /v1/models must list the model.""" - return probe_direct_server(host, port) - - def health_expectations(self, config: Any, processes: list[Process] | None) -> tuple[int, int, str]: - return logical_health_expectations(config) - - def validate(self, config: Any) -> None: - """One aggregate ``sglang.launch_server`` owns the public port. - - Several replicas or a prefill/decode layout need ``sglang-router`` (or - ``dynamo``); a schema 2 recipe that still says ``sglang`` for those is an - old router recipe and is rejected rather than silently run unbalanced. - """ - if config.frontend.enable_multiple_frontends: - raise ValueError( - "frontend.type: sglang binds sglang.launch_server directly; set frontend.enable_multiple_frontends: false" - ) - if config.resources.is_disaggregated: - raise ValueError( - "frontend.type: sglang supports one aggregate worker only, not a prefill/decode layout. " - "The SGLang router is frontend.type: sglang-router (renamed in 2.0; `srtctl migrate` rewrites " - "schema 1 recipes)." - ) - if config.resources.num_agg != 1: - raise ValueError( - f"frontend.type: sglang supports exactly one aggregate worker, got {config.resources.num_agg}. " - "sglang.launch_server owns the public port directly and there is no router to balance " - "replicas. Use frontend.type: sglang-router (the SGLang Model Gateway, renamed in 2.0) or dynamo." - ) - if config.dynamo.sidecar: - raise ValueError("frontend.type: sglang does not support dynamo.sidecar; use frontend.type: dynamo") - - def get_backend_health_urls( - self, - backend: Any, - backend_processes: list[Process], - network_interface: str | None = None, - ) -> list[str]: - return [] - - def implied_services(self, config: Any) -> list[EffectiveService]: - return [] - - def frontend_metrics_port(self, frontend_args: dict[str, Any] | None) -> int | None: - return None - - def start_frontends( - self, - topology: Any, - runtime: RuntimeContext, - config: Any, - backend: Any, - backend_processes: list[Process], - stop_event: threading.Event | None = None, - ) -> list[ManagedProcess]: - if config.backend.type != "sglang": - raise ValueError(f"frontend.type: sglang requires engine sglang (got {config.backend.type!r})") - if topology.uses_nginx or len(topology.frontend_nodes) != 1: - raise ValueError( - "frontend.type: sglang binds sglang.launch_server directly to the public port; " - "set frontend.enable_multiple_frontends: false" - ) - if config.resources.is_disaggregated: - raise ValueError("frontend.type: sglang supports one aggregate worker only; use sglang-router or dynamo") - if config.resources.num_agg != 1: - raise ValueError( - f"frontend.type: sglang supports exactly one aggregate worker, got {config.resources.num_agg}; " - "use frontend.type: sglang-router or dynamo to balance between replicas" - ) - - logger.info( - "frontend.type=sglang: no separate frontend process; sglang.launch_server owns port %d", - topology.public_port, - ) - return [] +class SGLangFrontend(DirectServerFrontend): + required_backend = "sglang" + server_name = "sglang.launch_server" + router_hint = "Use frontend.type: sglang-router (the SGLang Model Gateway, renamed in 2.0) or dynamo." diff --git a/tests/test_atom_atomesh.py b/tests/test_atom_atomesh.py index 937800b6e..2dcee3f18 100644 --- a/tests/test_atom_atomesh.py +++ b/tests/test_atom_atomesh.py @@ -66,10 +66,10 @@ def test_atomesh_frontend_requires_atom_backend() -> None: _load(data) -def test_atom_worker_requires_atomesh_frontend() -> None: +def test_atom_worker_rejects_incompatible_frontend() -> None: process = Process("node0", frozenset(range(8)), 7500, 6100, "agg", 0) - with pytest.raises(ValueError, match="requires frontend.type: atomesh"): + with pytest.raises(ValueError, match="requires an ATOM frontend"): AtomProtocol().build_worker_command(process, [process], _runtime(), frontend_type="dynamo") diff --git a/tests/test_atom_direct_frontend.py b/tests/test_atom_direct_frontend.py new file mode 100644 index 000000000..0fe8b9a1f --- /dev/null +++ b/tests/test_atom_direct_frontend.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 SemiAnalysis LLC. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Direct ATOM commands, topology validation, and readiness.""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import requests +from marshmallow import ValidationError + +from srtctl.core.config import resolve_config_with_defaults +from srtctl.core.schema import SrtConfig +from srtctl.core.topology import Process +from srtctl.frontends import get_frontend + + +def recipe(): + return { + "schema": 2, + "engine": "atom", + "name": "direct-test", + "model": {"path": "hf:test/model", "container": "test:tag", "precision": "bf16"}, + "resources": {"gpu_type": "mi300x", "gpus_per_node": 8}, + "frontend": {"type": "atom", "enable_multiple_frontends": False}, + "roles": {"agg": {"nodes": 1, "workers": 1, "gpus": 4, "args": {"kv_cache_dtype": "fp8"}}}, + } + + +def load(data): + return SrtConfig.Schema().load(resolve_config_with_defaults(data, None)) + + +def test_direct_worker_owns_public_port_without_router(): + config = load(recipe()) + frontend = get_frontend(config.frontend.type) + runtime = SimpleNamespace(worker_model_arg="test/model", network_interface=None, frontend_port=9017) + process = Process("node0", frozenset(range(4)), 7500, 6100, "agg", 0) + with patch("srtctl.core.slurm.get_hostname_ip", return_value="10.0.0.4"): + command = config.backend.build_worker_command(process, [process], runtime, frontend_type=config.frontend.type) + assert command == [ + "env", + "ATOM_HOST_IP=10.0.0.4", + "python3", + "-m", + "atom.entrypoints.openai_server", + "--model", + "test/model", + "--host", + "0.0.0.0", + "--server-port", + "9017", + "-tp", + "4", + "--kv_cache_dtype", + "fp8", + ] + topology = SimpleNamespace(uses_nginx=False, frontend_nodes=["node0"], public_port=9017) + assert frontend.start_frontends(topology, runtime, config, config.backend, [process]) == [] + assert frontend.worker_endpoint_port(process, config, runtime) == 9017 + assert frontend.direct_endpoint_nodes([process]) == ["node0"] + + +@pytest.mark.parametrize( + "change,error", + [ + ({"roles": {"agg": {"nodes": 1, "workers": 2, "gpus": 1}}}, "exactly one aggregate worker"), + ({"roles": {"prefill": {"nodes": 1, "workers": 1}, "decode": {"nodes": 1, "workers": 1}}}, "prefill/decode"), + ({"frontend": {"type": "atom", "enable_multiple_frontends": True}}, "enable_multiple_frontends: false"), + ({"engine": "sglang"}, "requires backend.type: atom"), + ], +) +def test_direct_layout_rejects_jobs_requiring_a_router(change, error): + with pytest.raises(ValidationError, match=error): + load({**recipe(), **change}) + + +@pytest.mark.parametrize("models,ready", [([], False), ([{"id": "test/model"}], True)]) +def test_direct_readiness_requires_loaded_model(models, ready): + health = requests.Response() + health.status_code = 200 + listing = requests.Response() + listing.status_code = 200 + import json + + listing._content = json.dumps({"data": models}).encode() + with patch("srtctl.core.health.requests.get", side_effect=[health, listing]): + result = get_frontend("atom").probe_ready("worker", 9017, 0, 1, load(recipe())) + assert result.ready is ready diff --git a/tests/test_frontends.py b/tests/test_frontends.py index acd4363a2..dade0cc02 100644 --- a/tests/test_frontends.py +++ b/tests/test_frontends.py @@ -67,6 +67,7 @@ class TestFrontendRegistry: def test_registry_lists_every_frontend_type(self): assert list_frontend_types() == [ + "atom", "atomesh", "dynamo", "none", @@ -207,7 +208,7 @@ def test_schema_rejects_unknown_type_at_load(self): from srtctl.backends import SGLangProtocol from srtctl.core.schema import FrontendConfig, ResourceConfig, SrtConfig - with pytest.raises(ValidationError, match="Unknown frontend.type 'toy-router'.*Available: atomesh, dynamo, none"): + with pytest.raises(ValidationError, match="Unknown frontend.type 'toy-router'.*Available: atom, atomesh, dynamo, none"): SrtConfig( name="toy", model={"path": "model", "container": "image", "precision": "fp8"}, From c29ef7c5d0732fbf7fc93aa4b7929b48f0bfa76a Mon Sep 17 00:00:00 2001 From: adibarra <93070681+adibarra@users.noreply.github.com> Date: Tue, 22 Sep 2026 15:15:06 -0500 Subject: [PATCH 3/3] chore(atom): preserve inherited NVIDIA copyright attribution MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 保留复用的 NVIDIA frontend 代码署名。 Signed-off-by: adibarra <93070681+adibarra@users.noreply.github.com> --- src/srtctl/frontends/atom.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/srtctl/frontends/atom.py b/src/srtctl/frontends/atom.py index a5ee260d6..aa49835da 100644 --- a/src/srtctl/frontends/atom.py +++ b/src/srtctl/frontends/atom.py @@ -1,3 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-FileCopyrightText: Copyright (c) 2026 SemiAnalysis LLC. All rights reserved. # SPDX-License-Identifier: Apache-2.0