From 8081c5bd1a08fee1562b72e16a7fcf71bc1ecc75 Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Tue, 22 Sep 2026 13:08:25 -0500 Subject: [PATCH] feat(atom): integrate native ATOM and AToMesh --- CLAUDE.md | 1 + README.md | 2 +- docs/config-reference.md | 13 ++ docs/schema-reference.md | 16 +- examples/README.md | 6 + examples/atom/atomesh-disagg.yaml | 32 ++++ src/srtctl/backends/__init__.py | 6 +- src/srtctl/backends/atom.py | 218 +++++++++++++++++++++++++ src/srtctl/backends/base.py | 1 + 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 +++++ tests/test_atom_atomesh.py | 262 ++++++++++++++++++++++++++++++ tests/test_benchmarks.py | 39 +++++ tests/test_frontends.py | 3 +- tests/test_port_allocator.py | 2 +- tests/test_schema_docs.py | 11 +- 20 files changed, 685 insertions(+), 18 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..6b70ed702 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -377,6 +377,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..a0b5f9354 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 SGLang, vLLM, TensorRT LLM and AMD's ATOM. 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/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..07082f0af --- /dev/null +++ b/src/srtctl/backends/atom.py @@ -0,0 +1,218 @@ +# 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_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..991fcd10b 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 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/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..664418da0 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 @@ -1942,6 +1944,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_frontends.py b/tests/test_frontends.py index 6d8aa8fb1..ccfa089cb 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"}, 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: