From e9472b52a9fc31eb4ff3c5bc08c366df05976ae1 Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Fri, 9 Oct 2026 09:18:36 -0600 Subject: [PATCH 01/20] draft srt-slurm support --- conf/experimental/srt_slurm/recipe.yaml | 23 ++ conf/experimental/srt_slurm/scenario.toml | 40 +++ conf/experimental/srt_slurm/srtslurm.yaml | 4 + doc/workloads/index.rst | 1 + doc/workloads/srt_slurm.rst | 96 +++++++ src/cloudai/registration.py | 8 + src/cloudai/workloads/srt_slurm/__init__.py | 26 ++ .../workloads/srt_slurm/installable.py | 107 +++++++ .../srt_slurm/report_generation_strategy.py | 74 +++++ src/cloudai/workloads/srt_slurm/results.py | 157 +++++++++++ .../srt_slurm/slurm_command_gen_strategy.py | 157 +++++++++++ src/cloudai/workloads/srt_slurm/srt_slurm.py | 95 +++++++ tests/test_acceptance.py | 20 ++ tests/workloads/srt_slurm/__init__.py | 15 + tests/workloads/srt_slurm/test_installable.py | 102 +++++++ .../test_report_generation_strategy.py | 50 ++++ tests/workloads/srt_slurm/test_results.py | 95 +++++++ .../test_slurm_command_gen_strategy.py | 71 +++++ tests/workloads/srt_slurm/test_srt_slurm.py | 260 ++++++++++++++++++ 19 files changed, 1401 insertions(+) create mode 100644 conf/experimental/srt_slurm/recipe.yaml create mode 100644 conf/experimental/srt_slurm/scenario.toml create mode 100644 conf/experimental/srt_slurm/srtslurm.yaml create mode 100644 doc/workloads/srt_slurm.rst create mode 100644 src/cloudai/workloads/srt_slurm/__init__.py create mode 100644 src/cloudai/workloads/srt_slurm/installable.py create mode 100644 src/cloudai/workloads/srt_slurm/report_generation_strategy.py create mode 100644 src/cloudai/workloads/srt_slurm/results.py create mode 100644 src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py create mode 100644 src/cloudai/workloads/srt_slurm/srt_slurm.py create mode 100644 tests/workloads/srt_slurm/__init__.py create mode 100644 tests/workloads/srt_slurm/test_installable.py create mode 100644 tests/workloads/srt_slurm/test_report_generation_strategy.py create mode 100644 tests/workloads/srt_slurm/test_results.py create mode 100644 tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py create mode 100644 tests/workloads/srt_slurm/test_srt_slurm.py diff --git a/conf/experimental/srt_slurm/recipe.yaml b/conf/experimental/srt_slurm/recipe.yaml new file mode 100644 index 000000000..32f2980f0 --- /dev/null +++ b/conf/experimental/srt_slurm/recipe.yaml @@ -0,0 +1,23 @@ +schema: 2 +name: sglang-example +model: + path: Qwen/Qwen3-0.6B + container: docker.io/lmsysorg/sglang:latest +resources: + gpu_type: h100 + gpus_per_node: 8 +frontend: + type: sglang +engine: sglang +roles: + agg: + nodes: 1 + workers: 1 + gpus: 1 + args: + tensor-parallel-size: 1 +benchmark: + type: sa-bench + isl: 128 + osl: 128 + concurrencies: "1" diff --git a/conf/experimental/srt_slurm/scenario.toml b/conf/experimental/srt_slurm/scenario.toml new file mode 100644 index 000000000..202e006e6 --- /dev/null +++ b/conf/experimental/srt_slurm/scenario.toml @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +name = "srt_slurm" + +[[Tests]] +id = "srt.sglang" +num_nodes = 1 + +name = "srt_sglang" +description = "SGLang through an existing srt-slurm recipe" +test_template_name = "SRTSlurm" +agent = "grid_search" +agent_metrics = ["tput_per_gpu"] +serving_gpus = 1 + + [[Tests.git_repos]] + url = "https://github.com/NVIDIA/srt-slurm.git" + commit = "v2.2.1" + + [Tests.cmd_args] + recipe_path = "conf/experimental/srt_slurm/recipe.yaml" + cluster_config = "conf/experimental/srt_slurm/srtslurm.yaml" + result_format = "sa-bench" + + [Tests.cmd_args.recipe_overrides.benchmark] + concurrencies = ["1", "4", "16"] diff --git a/conf/experimental/srt_slurm/srtslurm.yaml b/conf/experimental/srt_slurm/srtslurm.yaml new file mode 100644 index 000000000..6441472a8 --- /dev/null +++ b/conf/experimental/srt_slurm/srtslurm.yaml @@ -0,0 +1,4 @@ +# Adapt these paths and Slurm defaults to your shared filesystem. +default_partition: gpu +gpus_per_node: 8 +srtctl_root: /shared/cloudai/install/srt-slurm__v2.2.1 diff --git a/doc/workloads/index.rst b/doc/workloads/index.rst index 916d596cc..05f6c19ef 100644 --- a/doc/workloads/index.rst +++ b/doc/workloads/index.rst @@ -30,6 +30,7 @@ Available Workloads ":doc:`nixl_perftest`", "✅", "❌", "❌", "❌" ":doc:`sleep`", "✅", "✅", "❌", "✅" ":doc:`sglang`", "✅", "❌", "❌", "❌" + ":doc:`srt_slurm`", "✅", "❌", "❌", "❌" ":doc:`slurm_container`", "✅", "❌", "❌", "❌" "Triton Inference", "✅", "❌", "❌", "❌" ":doc:`ucc`", "✅", "❌", "❌", "❌" diff --git a/doc/workloads/srt_slurm.rst b/doc/workloads/srt_slurm.rst new file mode 100644 index 000000000..8d4b41ec8 --- /dev/null +++ b/doc/workloads/srt_slurm.rst @@ -0,0 +1,96 @@ +SRT-Slurm +========= + +The ``SRTSlurm`` workload submits an existing recipe with ``srtctl apply`` and +monitors the resulting Slurm job directly. The pinned srtctl version validates +the recipe, including recipes without a ``schema`` field. Each CloudAI trial +passes selected ``recipe_overrides`` as typed ``--set KEY=VALUE`` arguments; +the original YAML remains unchanged. + +Configuration +------------- + +Use ``conf/experimental/srt_slurm/scenario.toml`` as a starting point. Adapt +its cluster configuration to your system, including ``srtctl_root`` pointing +to the installed checkout. Paths are resolved from the CloudAI working directory +and must be accessible on the shared filesystem. + +.. code-block:: toml + + name = "srt_sglang" + description = "SGLang with srt-slurm" + test_template_name = "SRTSlurm" + agent_metrics = ["tput_per_gpu"] + serving_gpus = 4 + + [[git_repos]] + url = "https://github.com/NVIDIA/srt-slurm.git" + commit = "v2.2.1" + + [cmd_args] + recipe_path = "/shared/recipes/sglang.yaml" + cluster_config = "/shared/config/srtslurm.yaml" + result_format = "sa-bench" + + [cmd_args.recipe_overrides.benchmark] + concurrencies = [[1, 4], [8, 16]] + +The first ``git_repos`` entry supplies srtctl; additional entries use ordinary +CloudAI installation. Pin a revision supporting ``srtctl apply --set``. +Lists in ``recipe_overrides`` are DSE dimensions. Use nested candidate lists +when the selected recipe value is itself a list. Recipe paths, cluster config +paths, and result format are excluded from DSE. Additional apply flags use +``extra_cmd_args``, for example ``extra_cmd_args = { no-preflight = "" }``. + +``cluster_config`` is required and takes precedence over ``SRTSLURM_CONFIG`` +in the submitting environment. Supply a complete config with the recipe's +aliases, mounts, and cluster settings. CloudAI exports merged system and test +environment variables, removes login-node virtual environments from ``PATH``, +unsets ``VIRTUAL_ENV``, and submits from the installed srt-slurm checkout. + +Installation runs ``make setup`` in a single-node, 15-minute ``srun`` using the +CloudAI system's partition, account, and extra srun arguments. This downloads +runtime tools for the compute-node architecture. A checkout-local bootstrap +config prevents interactive setup prompts; existing checkout configs are +preserved. Trial submission always uses the explicit ``cluster_config``. + +The recipe and cluster config own job resources and orchestration. Configure +nodes, time limits, container mounts, profiling, and setup/cleanup commands +there. CloudAI's scenario allocation directives do not rewrite the recipe. +Use normal per-trial Slurm CLI execution; single-sbatch, Slurm REST submission, +and CloudAI pre/post-test hooks are unsupported for this workload. + +Results +------- + +``result_format = "sa-bench"`` reads job-local ``sa-bench_*/results_*.json`` +files, falling back to ``benchmark-rollup.json`` when raw results are absent. +``default`` and ``output_throughput`` expose output tokens/s. Set +``serving_gpus`` to the actual serving GPU count to enable ``tput_per_gpu`` +from total token throughput, including input tokens. For disaggregated serving, +also set ``prefill_gpus`` and ``decode_gpus`` for the input/output per-GPU +divisors. Infrastructure and idle GPUs do not belong in these counts; shared +aggregate GPUs count once. A rollup lacking total throughput cannot provide +``tput_per_gpu``. Latency aliases such as ``median_ttft`` use seconds; +``ttft_median_ms`` uses milliseconds. ``median_intvty`` and ``p90_intvty`` +use the reciprocal of the corresponding TPOT, in tokens/s/user. + +``result_format = "agentx"`` reads InferenceX aggregate JSON from the submitted +job's ``logs/agentic`` directory. Each aggregate must have its corresponding +``conc_/aiperf_artifacts/profile_export.jsonl``. CloudAI reads the existing +aggregation rather than recomputing per-request statistics. Throughput uses +``request_metrics.throughput.per_gpu``; interactivity uses full-response ITL +when available, otherwise ITL. Latency aliases use the same units as sa-bench. + +An AgentX recipe uses ``benchmark.type: custom`` and its existing InferenceX +command. Supply a pinned InferenceX checkout with ``init_submodules = true`` +and the recipe's required ``INFMAX_WORKSPACE`` environment. CloudAI exports +``AGENTIC_OUTPUT_DIR=/logs/agentic`` for this result format. Keep that same value +in any recipe benchmark environment, and keep ``RESULT_DIR=/logs/agentic`` +if explicitly configured, so aggregate and profile outputs remain job-local. + +Reports write ``srt_slurm_results.csv`` with every concurrency. Scalar DSE +metrics come from one point: the highest per-GPU total throughput when available, +otherwise the highest output throughput. All latency metrics refer to that +same point. Missing, malformed, and nonfinite results do not mark a trial +successful; unavailable metrics return CloudAI's metric error value. diff --git a/src/cloudai/registration.py b/src/cloudai/registration.py index 883d5f8fd..7f8b9c626 100644 --- a/src/cloudai/registration.py +++ b/src/cloudai/registration.py @@ -183,6 +183,11 @@ def register_all(): SleepTestDefinition, ) from cloudai.workloads.slurm_container import SlurmContainerCommandGenStrategy, SlurmContainerTestDefinition + from cloudai.workloads.srt_slurm import ( + SRTSlurmCommandGenStrategy, + SRTSlurmReportGenerationStrategy, + SRTSlurmTestDefinition, + ) from cloudai.workloads.triton_inference import ( TritonInferenceReportGenerationStrategy, TritonInferenceSlurmCommandGenStrategy, @@ -266,6 +271,7 @@ def register_all(): Registry().add_command_gen_strategy(SlurmSystem, FioTestDefinition, FioSlurmCommandGenStrategy) Registry().add_command_gen_strategy(SlurmSystem, SglangTestDefinition, SglangSlurmCommandGenStrategy) Registry().add_command_gen_strategy(SlurmSystem, VllmTestDefinition, VllmSlurmCommandGenStrategy) + Registry().add_command_gen_strategy(SlurmSystem, SRTSlurmTestDefinition, SRTSlurmCommandGenStrategy) Registry().add_installer("slurm", SlurmInstaller) Registry().add_installer("standalone", StandaloneInstaller) @@ -298,6 +304,7 @@ def register_all(): Registry().add_test_definition("NixlEP", NixlEPTestDefinition) Registry().add_test_definition("AIDynamo", AIDynamoTestDefinition) Registry().add_test_definition("BashCmd", BashCmdTestDefinition) + Registry().add_test_definition("SRTSlurm", SRTSlurmTestDefinition) Registry().add_test_definition("NixlPerftest", NixlPerftestTestDefinition) Registry().add_test_definition("NIXLKVBench", NIXLKVBenchTestDefinition) Registry().add_test_definition("Aiconfigurator", AiconfiguratorTestDefinition) @@ -333,6 +340,7 @@ def register_all(): Registry().add_report(OSUBenchTestDefinition, OSUBenchReportGenerationStrategy) Registry().add_report(SglangTestDefinition, SGLangBenchReportGenerationStrategy) Registry().add_report(VllmTestDefinition, VLLMBenchReportGenerationStrategy) + Registry().add_report(SRTSlurmTestDefinition, SRTSlurmReportGenerationStrategy) Registry().add_scenario_report("per_test", PerTestReporter, ReportConfig(enable=True)) Registry().add_scenario_report("training", TrainingReporter, ReportConfig(enable=True)) diff --git a/src/cloudai/workloads/srt_slurm/__init__.py b/src/cloudai/workloads/srt_slurm/__init__.py new file mode 100644 index 000000000..202cda6e9 --- /dev/null +++ b/src/cloudai/workloads/srt_slurm/__init__.py @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .report_generation_strategy import SRTSlurmReportGenerationStrategy +from .slurm_command_gen_strategy import SRTSlurmCommandGenStrategy +from .srt_slurm import SRTSlurmCmdArgs, SRTSlurmTestDefinition + +__all__ = [ + "SRTSlurmCmdArgs", + "SRTSlurmCommandGenStrategy", + "SRTSlurmReportGenerationStrategy", + "SRTSlurmTestDefinition", +] diff --git a/src/cloudai/workloads/srt_slurm/installable.py b/src/cloudai/workloads/srt_slurm/installable.py new file mode 100644 index 000000000..f39391943 --- /dev/null +++ b/src/cloudai/workloads/srt_slurm/installable.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import shlex +import subprocess +from pathlib import Path +from typing import cast + +import yaml + +from cloudai.core import BaseInstaller, InstallStatusResult, PythonExecutable +from cloudai.systems.slurm import SlurmSystem + + +class SRTSlurmExecutable(PythonExecutable): + """Install srtctl and the binaries required on the compute nodes.""" + + def _runtime_status(self, installer: BaseInstaller) -> InstallStatusResult: + repo_path = self.git_repo.installed_path or installer.system.install_path / self.git_repo.repo_name + binaries = ["configs/nats-server", "configs/etcd", "bin/uv"] + try: + # Older srt-slurm revisions did not install a Tachometer scraper. + if "tachometer-scraper-download" in (repo_path / "Makefile").read_text(encoding="utf-8"): + binaries.append("bin/tachometer-scraper") + invalid = [ + name + for name in binaries + if not (repo_path / name).is_file() or not os.access(repo_path / name, os.X_OK) + ] + except OSError as error: + return InstallStatusResult(False, f"Unable to check srt-slurm runtime: {error}") + if invalid: + return InstallStatusResult(False, f"Missing or non-executable runtime binaries: {', '.join(invalid)}") + return InstallStatusResult(True, "srt-slurm runtime binaries installed") + + def is_installed(self, installer: BaseInstaller) -> InstallStatusResult: + for result in (self.git_repo.is_installed(installer), super().is_installed(installer)): + if not result.success: + return result + if self.venv_path is None or not os.access(self.venv_path / "bin" / "srtctl", os.X_OK): + return InstallStatusResult(False, "srtctl executable is missing from the Python environment") + repo_path = self.git_repo.installed_path or installer.system.install_path / self.git_repo.repo_name + if not (repo_path / "srtslurm.yaml").is_file(): + return InstallStatusResult(False, "srt-slurm setup configuration is missing") + return self._runtime_status(installer) + + def _prepare_cluster_config(self, installer: BaseInstaller, repo_path: Path) -> None: + # make setup prompts interactively without this file. This bootstrap + # config is checkout-local; trial submission always uses the supplied config. + config = repo_path / "srtslurm.yaml" + if not config.is_file(): + self._write_cluster_config(installer, config, repo_path) + + def _write_cluster_config(self, installer: BaseInstaller, config_path: Path, repo_path: Path) -> None: + system = cast(SlurmSystem, installer.system) + config = { + "default_account": system.account, + "default_partition": system.default_partition, + "gpus_per_node": system.gpus_per_node, + "srtctl_root": str(repo_path.resolve()), + } + config_path.write_text( + yaml.safe_dump({key: value for key, value in config.items() if value is not None}, sort_keys=False), + encoding="utf-8", + ) + + def install(self, installer: BaseInstaller) -> InstallStatusResult: + result = super().install(installer) + if not result.success: + return result + repo_path = self.git_repo.installed_path or installer.system.install_path / self.git_repo.repo_name + try: + self._prepare_cluster_config(installer, repo_path) + if self.is_installed(installer).success: + return InstallStatusResult(True, "srtctl and runtime binaries already installed") + system = cast(SlurmSystem, installer.system) + command = ["srun", "--nodes=1", "--ntasks=1", "--time=00:15:00", f"--partition={system.default_partition}"] + if system.account: + command.append(f"--account={system.account}") + command.extend(shlex.split(system.extra_srun_args or "")) + command.extend(["make", "setup"]) + completed = subprocess.run( + command, + cwd=repo_path, + stdin=subprocess.DEVNULL, + capture_output=True, + text=True, + ) + except OSError as error: + return InstallStatusResult(False, f"srt-slurm setup failed: {error}") + if completed.returncode != 0: + return InstallStatusResult(False, f"srt-slurm setup failed:\n{completed.stdout}\n{completed.stderr}") + return self.is_installed(installer) diff --git a/src/cloudai/workloads/srt_slurm/report_generation_strategy.py b/src/cloudai/workloads/srt_slurm/report_generation_strategy.py new file mode 100644 index 000000000..2ee62d92b --- /dev/null +++ b/src/cloudai/workloads/srt_slurm/report_generation_strategy.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import csv +from typing import ClassVar, cast + +from cloudai.core import METRIC_ERROR, MetricValue, ReportGenerationStrategy + +from .results import load_results +from .srt_slurm import SRTSlurmTestDefinition + + +class SRTSlurmReportGenerationStrategy(ReportGenerationStrategy): + """Write every concurrency; use the best throughput point for scalar DSE metrics.""" + + metrics: ClassVar[list[str]] = [ + "default", + "output_throughput", + "throughput_toks", + "request_throughput", + "tput_per_gpu", + "input_tput_per_gpu", + "output_tput_per_gpu", + "completed_requests", + *[ + f"{family}_{stat}_ms" + for family in ("ttft", "tpot", "itl", "e2el") + for stat in ("mean", "median", "p75", "p90", "p95", "p99") + ], + *[ + f"{stat}_{family}" + for family in ("ttft", "tpot", "itl", "e2el", "intvty") + for stat in ("mean", "median", "p75", "p90", "p95", "p99") + ], + ] + + def results(self) -> list[dict[str, float]]: + return load_results(cast(SRTSlurmTestDefinition, self.test_run.test), self.test_run) + + def can_handle_directory(self) -> bool: + return bool(self.results()) + + def generate_report(self) -> None: + results = self.results() + if not results: + return + fields = sorted({key for run in results for key in run}) + with (self.test_run.output_path / "srt_slurm_results.csv").open("w", newline="") as output: + writer = csv.DictWriter(output, fieldnames=fields) + writer.writeheader() + writer.writerows(results) + + def get_metric(self, metric: str) -> MetricValue: + results = self.results() + if not results: + return METRIC_ERROR + key = "output_throughput" if metric == "default" else metric + # All metrics come from the same point. AgentX already supplies the + # serving-GPU-normalized throughput; sa-bench keeps the original default. + best = max(results, key=lambda run: run.get("tput_per_gpu", run.get("output_throughput", 0))) + return best.get(key, METRIC_ERROR) diff --git a/src/cloudai/workloads/srt_slurm/results.py b/src/cloudai/workloads/srt_slurm/results.py new file mode 100644 index 000000000..2d9df819a --- /dev/null +++ b/src/cloudai/workloads/srt_slurm/results.py @@ -0,0 +1,157 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import math +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from cloudai.core import TestRun + +if TYPE_CHECKING: + from .srt_slurm import SRTSlurmTestDefinition + + +def read_json(path: Path) -> dict[str, Any]: + try: + data = json.loads(path.read_text(encoding="utf-8")) + except (OSError, UnicodeDecodeError, ValueError): + return {} + return data if isinstance(data, dict) else {} + + +def number(value: Any) -> float | None: + if isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value): + return float(value) + return None + + +def numeric_fields(data: dict[str, Any]) -> dict[str, float]: + return {key: value for key, raw in data.items() if (value := number(raw)) is not None} + + +def sa_bench_metrics(data: dict[str, Any], tdef: "SRTSlurmTestDefinition") -> dict[str, float]: + metrics = numeric_fields(data) + aliases = {"max_concurrency": "concurrency", "output_throughput": "throughput_toks"} + for old, new in aliases.items(): + if old in metrics: + metrics[new] = metrics[old] + if "throughput_toks" not in metrics: + return {} + metrics["output_throughput"] = metrics["throughput_toks"] + if tdef.serving_gpus and "total_token_throughput" in metrics: + metrics["tput_per_gpu"] = metrics["total_token_throughput"] / tdef.serving_gpus + output_gpus = tdef.decode_gpus or tdef.serving_gpus + metrics["output_tput_per_gpu"] = metrics["throughput_toks"] / output_gpus + input_gpus = tdef.prefill_gpus or tdef.serving_gpus + metrics["input_tput_per_gpu"] = (metrics["total_token_throughput"] - metrics["throughput_toks"]) / input_gpus + add_sa_latencies(data, metrics) + return metrics + + +def sa_latency_value(data: dict[str, Any], metrics: dict[str, float], family: str, stat: str) -> float | None: + value = metrics.get(f"{stat}_{family}_ms", metrics.get(f"{family}_{stat}_ms")) + if value is not None: + return value + percentiles = data.get(f"percentiles_{family}_ms", []) + if not stat.startswith("p") or not isinstance(percentiles, list): + return None + return next( + (number(v[1]) for v in percentiles if isinstance(v, (list, tuple)) and len(v) == 2 and v[0] == int(stat[1:])), + None, + ) + + +def add_sa_latencies(data: dict[str, Any], metrics: dict[str, float]) -> None: + for family in ("ttft", "tpot", "itl", "e2el"): + for stat in ("mean", "median", "p75", "p90", "p95", "p99"): + value = sa_latency_value(data, metrics, family, stat) + if value is not None: + metrics[f"{stat}_{family}"] = value / 1000 + metrics[f"{family}_{stat}_ms"] = value + if family == "tpot" and value > 0: + metrics[f"{stat}_intvty"] = 1000 / value + + +def agentx_metrics(data: dict[str, Any]) -> dict[str, float]: + request = data.get("request_metrics") + if not isinstance(request, dict): + return {} + throughput = request.get("throughput", {}) + per_gpu = throughput.get("per_gpu", {}) if isinstance(throughput, dict) else {} + if not isinstance(per_gpu, dict) or number(per_gpu.get("total_tput_tps")) is None: + return {} + metrics = numeric_fields(data) + if "conc" in metrics: + metrics["concurrency"] = metrics["conc"] + for source, dest in ( + ("total_tput_tps", "tput_per_gpu"), + ("output_tput_tps", "output_tput_per_gpu"), + ("input_tput_tps", "input_tput_per_gpu"), + ): + if (value := number(per_gpu.get(source))) is not None: + metrics[dest] = value + metrics.update(numeric_fields(throughput)) + for source, dest in (("total", "total_tput_tps"), ("input", "input_tput_tps"), ("output", "output_tput_tps")): + block = throughput.get(source, {}) + if isinstance(block, dict) and (value := number(block.get("tokens_per_second"))) is not None: + metrics[dest] = value + if "output_tput_tps" in metrics: + metrics["output_throughput"] = metrics["throughput_toks"] = metrics["output_tput_tps"] + add_agentx_latencies(request.get("latency", {}), metrics) + return metrics + + +def add_agentx_latencies(latency: Any, metrics: dict[str, float]) -> None: + if isinstance(latency, dict): + # InferenceX's published AgentX interactivity uses full-response ITL. + for family in ("ttft", "tpot", "itl", "e2el"): + stats = ( + latency.get("full_response_itl", latency.get("itl", {})) if family == "itl" else latency.get(family, {}) + ) + if not isinstance(stats, dict): + continue + for stat, value in numeric_fields(stats).items(): + metrics[f"{stat}_{family}"] = value + metrics[f"{family}_{stat}_ms"] = value * 1000 + if family == "itl" and value > 0 and stat != "std": + metrics[f"{stat}_intvty"] = 1 / value + + +def load_results(tdef: "SRTSlurmTestDefinition", tr: TestRun) -> list[dict[str, float]]: + """Read only the current job's outputs, retaining each concurrency.""" + logs = tdef.logs_path(tr) + if logs is None: + return [] + if tdef.cmd_args.result_format == "agentx": + results = [] + for path in sorted((logs / "agentic").glob("*.json")): + metrics = agentx_metrics(read_json(path)) + concurrency = metrics.get("concurrency") + if concurrency is None or not concurrency.is_integer(): + continue + profiles = logs / "agentic" / f"conc_{int(concurrency)}" / "aiperf_artifacts" + if not any(p.stat().st_size > 0 for p in profiles.rglob("profile_export.jsonl")): + continue + results.append(metrics) + return results + raw_files = sorted(logs.glob("sa-bench_*/results_*.json")) + if raw_files: + return [metrics for path in raw_files if (metrics := sa_bench_metrics(read_json(path), tdef))] + runs = read_json(logs / "benchmark-rollup.json").get("runs", []) + if not isinstance(runs, list): + return [] + return [metrics for run in runs if isinstance(run, dict) and (metrics := sa_bench_metrics(run, tdef))] diff --git a/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py b/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py new file mode 100644 index 000000000..f736f04ce --- /dev/null +++ b/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py @@ -0,0 +1,157 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import shlex +from pathlib import Path +from typing import Any, cast + +import yaml + +from cloudai.systems.slurm import SlurmCommandGenStrategy + +from .srt_slurm import SRTSlurmTestDefinition + + +class SRTSlurmCommandGenStrategy(SlurmCommandGenStrategy): + """Submit the recipe job directly and let CloudAI monitor its job ID.""" + + @staticmethod + def _load_recipe(path: Path) -> dict[str, Any]: + try: + data = yaml.safe_load(path.read_text(encoding="utf-8")) + except OSError as error: + raise ValueError(f"Unable to read srt-slurm recipe {path}: {error}") from error + except yaml.YAMLError as error: + raise ValueError(f"Unable to parse srt-slurm recipe {path}: {error}") from error + if not isinstance(data, dict): + raise ValueError(f"srt-slurm recipe {path} must contain a YAML mapping.") + return data + + @staticmethod + def _override_args(overrides: dict[str, Any], prefix: str = "") -> list[str]: + """Flatten DSE-selected values into typed srtctl --set arguments.""" + args: list[str] = [] + for key, value in overrides.items(): + segment = f'"{key}"' if any(char in key for char in ".[]") else key + path = f"{prefix}.{segment}" if prefix else segment + if isinstance(value, dict): + args.extend(SRTSlurmCommandGenStrategy._override_args(value, path)) + else: + # JSON is valid YAML and quotes strings such as "true" and "720". + args.extend(["--set", f"{path}={json.dumps(value)}"]) + return args + + @property + def tdef(self) -> SRTSlurmTestDefinition: + return cast(SRTSlurmTestDefinition, self.test_run.test) + + def _paths(self) -> dict[str, Path]: + output = self.test_run.output_path.absolute() + generated = output / "cloudai_generated" + return { + "generated": generated, + "recipe": self.tdef.cmd_args.recipe_path.expanduser().absolute(), + "nested_outputs": output / "srt_slurm_outputs", + "run_log": output / "run_srtctl.log", + "run_script": generated / "run_srtctl.sh", + "job_id": self.tdef.job_id_path(self.test_run).absolute(), + } + + def write_run_script(self) -> None: + paths = self._paths() + paths["generated"].mkdir(parents=True, exist_ok=True) + repo = self.tdef.git_repos[0] + repo_path = (repo.installed_path or self.system.install_path / repo.repo_name).absolute() + cluster_config = self.tdef.cmd_args.cluster_config.expanduser().absolute() + apply_command = shlex.join( + [ + str(self.tdef.srtctl_path), + "apply", + "-f", + str(paths["recipe"]), + "-o", + str(paths["nested_outputs"]), + "-y", + *self._override_args(self.tdef.cmd_args.recipe_overrides), + *self._extra_cmd_arg_parts(), + ] + ) + script = "\n".join( + [ + "#!/usr/bin/env bash", + "set -euo pipefail", + *[f"export {key}={shlex.quote(str(value))}" for key, value in self.final_env_vars.items()], + f"export SRTSLURM_CONFIG={shlex.quote(str(cluster_config))}", + *(["export AGENTIC_OUTPUT_DIR=/logs/agentic"] if self.tdef.cmd_args.result_format == "agentx" else []), + # uv also discovers Python through PATH after VIRTUAL_ENV is + # unset. Strip login-node venvs before Slurm exports the PATH. + 'IFS=: read -r -a SRT_PATH_PARTS <<< "$PATH"', + "SRT_CLEAN_PATH=()", + 'for SRT_PATH_PART in "${SRT_PATH_PARTS[@]}"; do', + ' if [[ ! -f "$SRT_PATH_PART/../pyvenv.cfg" ]]; then', + ' SRT_CLEAN_PATH+=("$SRT_PATH_PART")', + " fi", + "done", + 'export PATH="$(IFS=:; printf \'%s\' "${SRT_CLEAN_PATH[*]}")"', + "unset VIRTUAL_ENV", + f"mkdir -p {shlex.quote(str(paths['nested_outputs']))}", + f"rm -f {shlex.quote(str(paths['job_id']))}", + # Slurm inherits this directory. The early Dynamo prefetch must + # not discover the login-node CloudAI .venv in the caller's cwd. + f"cd {shlex.quote(str(repo_path))}", + f"if ! {apply_command} > {shlex.quote(str(paths['run_log']))} 2>&1; then", + f" cat {shlex.quote(str(paths['run_log']))} >&2", + " exit 1", + "fi", + f"cat {shlex.quote(str(paths['run_log']))} >&2", + "SRT_JOB_ID=\"$(sed -nE 's/.*Job ([0-9]+) submitted.*/\\1/p; s/.*Submitted batch job ([0-9]+).*/\\1/p' " + f'{shlex.quote(str(paths["run_log"]))} | tail -1)"', + '[[ "$SRT_JOB_ID" =~ ^[0-9]+$ ]] || { echo "Unable to find submitted srt-slurm job ID"; exit 1; }', + f'printf "%s\\n" "$SRT_JOB_ID" > {shlex.quote(str(paths["job_id"]))}', + 'echo "Submitted batch job $SRT_JOB_ID"', + "", + ] + ) + paths["run_script"].write_text(script, encoding="utf-8") + + def _extra_cmd_arg_parts(self) -> list[str]: + args = [] + for name, value in self.tdef.extra_cmd_args.items(): + args.append(f"--{name.replace('_', '-')}") + if value: + args.append(value) + return args + + def _container_mounts(self) -> list[str]: + return [] + + def gen_exec_command(self) -> str: + self.test_run.output_path.mkdir(parents=True, exist_ok=True) + self._load_recipe(self._paths()["recipe"]) + if not self.tdef.cmd_args.cluster_config.expanduser().is_file(): + raise ValueError("The supplied srt-slurm cluster_config does not exist.") + if self.system.uses_slurm_api: + raise ValueError("SRTSlurm requires Slurm CLI submission.") + if self.test_run.pre_test or self.test_run.post_test: + raise ValueError("SRTSlurm hooks must be configured in the recipe.") + self.write_run_script() + return self.gen_srun_command() + + def gen_srun_command(self) -> str: + if not self._paths()["run_script"].is_file(): + raise ValueError("SRTSlurm submits its own job; single-sbatch execution is unsupported.") + return shlex.join(["bash", str(self._paths()["run_script"])]) diff --git a/src/cloudai/workloads/srt_slurm/srt_slurm.py b/src/cloudai/workloads/srt_slurm/srt_slurm.py new file mode 100644 index 000000000..eb719a903 --- /dev/null +++ b/src/cloudai/workloads/srt_slurm/srt_slurm.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path +from typing import Any, Literal + +from pydantic import Field, PrivateAttr, model_validator + +from cloudai.core import GitRepo, Installable, JobStatusResult, TestRun +from cloudai.models.workload import CmdArgs, TestDefinition + +from .installable import SRTSlurmExecutable + + +class SRTSlurmCmdArgs(CmdArgs): + """An existing srt-slurm recipe and the values CloudAI may tune.""" + + recipe_path: Path + cluster_config: Path + result_format: Literal["sa-bench", "agentx"] = "sa-bench" + recipe_overrides: dict[str, Any] = Field(default_factory=dict) + + def excluded_cmd_args(self) -> list[str]: + return ["cmd_args.recipe_path", "cmd_args.cluster_config", "cmd_args.result_format"] + + +class SRTSlurmTestDefinition(TestDefinition): + """Run one declarative srt-slurm recipe as one CloudAI trial.""" + + cmd_args: SRTSlurmCmdArgs + serving_gpus: int | None = Field(default=None, gt=0) + prefill_gpus: int | None = Field(default=None, gt=0) + decode_gpus: int | None = Field(default=None, gt=0) + git_repos: list[GitRepo] = Field( + default_factory=list, min_length=1, validate_default=True, description="The first repository supplies srtctl." + ) + _srtctl: SRTSlurmExecutable | None = PrivateAttr(default=None) + + @model_validator(mode="after") + def exclude_recipe_structure_from_dse(self) -> "SRTSlurmTestDefinition": + self.dse_excluded_args = list(dict.fromkeys([*self.dse_excluded_args, *self.cmd_args.excluded_cmd_args()])) + return self + + @property + def srtctl(self) -> SRTSlurmExecutable: + if self._srtctl is None: + self._srtctl = SRTSlurmExecutable(git_repo=self.git_repos[0]) + return self._srtctl + + @property + def srtctl_path(self) -> Path: + if self.srtctl.venv_path is None: + raise RuntimeError("srt-slurm must be installed before generating its command.") + return self.srtctl.venv_path.absolute() / "bin" / "srtctl" + + @property + def installables(self) -> list[Installable]: + return [self.srtctl, *self.git_repos[1:]] + + @staticmethod + def job_id_path(tr: TestRun) -> Path: + return tr.output_path / "srt_job_id.txt" + + def rollup_path(self, tr: TestRun) -> Path | None: + try: + job_id = self.job_id_path(tr).read_text(encoding="utf-8").strip() + except OSError: + return None + if not job_id.isdigit(): + return None + return tr.output_path / "srt_slurm_outputs" / job_id / "logs" / "benchmark-rollup.json" + + def was_run_successful(self, tr: TestRun) -> JobStatusResult: + from .results import load_results + + if not load_results(self, tr): + return JobStatusResult(is_successful=False, error_message="No usable job-local srt-slurm results found.") + return JobStatusResult(is_successful=True) + + def logs_path(self, tr: TestRun) -> Path | None: + rollup = self.rollup_path(tr) + return rollup.parent if rollup else None diff --git a/tests/test_acceptance.py b/tests/test_acceptance.py index 3f6b38b50..4a2a2bea3 100644 --- a/tests/test_acceptance.py +++ b/tests/test_acceptance.py @@ -951,3 +951,23 @@ def test_sbatch_generation(slurm_system: SlurmSystem, test_req: tuple[TestRun, s curr_aiperf = aiperf_script.read_text().strip() ref_aiperf = (Path(__file__).parent / "ref_data" / "ai-dynamo-aiperf.sh").read_text().strip() assert curr_aiperf == ref_aiperf, "aiperf.sh does not match reference" + + +def test_srt_slurm_example_dry_run(slurm_system: SlurmSystem, tmp_path: Path): + from cloudai.core import BaseInstaller + from cloudai.test_scenario_parser import TestScenarioParser + from cloudai.workloads.srt_slurm import SRTSlurmTestDefinition + + scenario = TestScenarioParser(Path("conf/experimental/srt_slurm/scenario.toml"), slurm_system, {}, {}).parse() + tr = scenario.test_runs[0].apply_params_set({"recipe_overrides.benchmark.concurrencies": "4"}) + assert isinstance(tr.test, SRTSlurmTestDefinition) + tr.output_path = tmp_path / "trial" + BaseInstaller(slurm_system).mark_as_installed(tr.test.installables) + runner = SlurmRunner("dry-run", slurm_system, TestScenario(name="srt", test_runs=[tr]), tmp_path) + job = runner._submit_test(tr) + runner.on_job_submit(tr) + assert job.id == 0 + script = tr.output_path / "cloudai_generated" / "run_srtctl.sh" + assert script.is_file() + assert "--set" in script.read_text() + assert (tr.output_path / "test-run.toml").is_file() diff --git a/tests/workloads/srt_slurm/__init__.py b/tests/workloads/srt_slurm/__init__.py new file mode 100644 index 000000000..19b239dd6 --- /dev/null +++ b/tests/workloads/srt_slurm/__init__.py @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/workloads/srt_slurm/test_installable.py b/tests/workloads/srt_slurm/test_installable.py new file mode 100644 index 000000000..8bf4c0f8d --- /dev/null +++ b/tests/workloads/srt_slurm/test_installable.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import subprocess +import sys +from pathlib import Path +from unittest.mock import patch + +import pytest +import yaml + +from cloudai.core import BaseInstaller, GitRepo, InstallStatusResult, PythonExecutable +from cloudai.systems.slurm import SlurmSystem +from cloudai.workloads.srt_slurm.installable import SRTSlurmExecutable + + +@pytest.fixture +def runtime_installation(tmp_path: Path, slurm_system: SlurmSystem, monkeypatch): + executable = SRTSlurmExecutable(git_repo=GitRepo(url="https://github.com/NVIDIA/srt-slurm.git", commit="v2.2.1")) + repo = slurm_system.install_path / executable.git_repo.repo_name + repo.mkdir(parents=True) + executable.git_repo.installed_path = repo + executable.venv_path = slurm_system.install_path / executable.venv_name + srtctl = executable.venv_path / "bin" / "srtctl" + srtctl.parent.mkdir(parents=True) + srtctl.write_text("#!/bin/sh\n") + srtctl.chmod(0o755) + (repo / "Makefile").write_text( + f'# tachometer-scraper-download\nsetup:\n\ttest -f srtslurm.yaml\n\t"{sys.executable}" setup_fixture.py\n' + ) + tools = tmp_path / "compute tools" + tools.mkdir() + srun = tools / "srun" + srun.write_text(f"#!{sys.executable}\nimport os, sys\nos.execvp('make', sys.argv[sys.argv.index('make'):])\n") + srun.chmod(0o755) + monkeypatch.setenv("PATH", f"{tools}:{os.environ['PATH']}") + (repo / "setup_fixture.py").write_text( + "from pathlib import Path\n" + "for name in ['configs/nats-server', 'configs/etcd', 'bin/uv', 'bin/tachometer-scraper']:\n" + " path = Path(name)\n path.parent.mkdir(parents=True, exist_ok=True)\n" + " path.write_text('#!/bin/sh\\nexit 0\\n')\n path.chmod(0o755)\n" + "counter = Path('setup_count')\n" + "counter.write_text(str(int(counter.read_text()) + 1) if counter.exists() else '1')\n" + ) + with ( + patch.object(PythonExecutable, "install", return_value=InstallStatusResult(True)), + patch.object(GitRepo, "is_installed", return_value=InstallStatusResult(True)), + ): + yield executable, BaseInstaller(slurm_system), repo + + +def test_installer_sets_up_compute_runtime_and_reuses_it(runtime_installation): + executable, installer, repo = runtime_installation + assert installer.install([executable]).success + assert installer.is_installed([executable]).success + assert installer.install([executable]).success + assert (repo / "setup_count").read_text() == "1" + assert yaml.safe_load((repo / "srtslurm.yaml").read_text())["srtctl_root"] == str(repo.resolve()) + (repo / "bin" / "uv").unlink() + assert not executable.is_installed(installer).success + assert executable.install(installer).success + assert (repo / "setup_count").read_text() == "2" + + +def test_installer_preserves_existing_config(runtime_installation): + executable, installer, repo = runtime_installation + original = "default_partition: custom\nsrtctl_root: /shared/custom\n" + (repo / "srtslurm.yaml").write_text(original) + assert executable.install(installer).success + assert (repo / "srtslurm.yaml").read_text() == original + + +def test_installer_returns_setup_failure(runtime_installation): + executable, installer, _ = runtime_installation + with patch( + "cloudai.workloads.srt_slurm.installable.subprocess.run", + return_value=subprocess.CompletedProcess(["make"], 2, "setup output", "download failed"), + ): + result = executable.install(installer) + assert not result.success + assert "download failed" in result.message + + +def test_installable_identity_uses_revision(): + repo = GitRepo(url="https://github.com/NVIDIA/srt-slurm.git", commit="v2.2.1") + executable = SRTSlurmExecutable(git_repo=repo) + other = SRTSlurmExecutable(git_repo=GitRepo(url=repo.url, commit="other")) + assert len({executable, SRTSlurmExecutable(git_repo=repo), other}) == 2 diff --git a/tests/workloads/srt_slurm/test_report_generation_strategy.py b/tests/workloads/srt_slurm/test_report_generation_strategy.py new file mode 100644 index 000000000..a8fa8b2b2 --- /dev/null +++ b/tests/workloads/srt_slurm/test_report_generation_strategy.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +from pathlib import Path + +from cloudai.core import TestRun +from cloudai.systems.slurm import SlurmSystem +from cloudai.workloads.srt_slurm import SRTSlurmReportGenerationStrategy + +from .test_srt_slurm import make_tdef + + +def test_report_reads_best_run_directly_from_srt_rollup(slurm_system: SlurmSystem, tmp_path: Path): + tdef = make_tdef(tmp_path) + test_run = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path / "output") + test_run.output_path.mkdir(parents=True) + tdef.job_id_path(test_run).write_text("12345\n", encoding="utf-8") + rollup_path = test_run.output_path / "srt_slurm_outputs" / "12345" / "logs" / "benchmark-rollup.json" + rollup_path.parent.mkdir(parents=True) + rollup_path.write_text( + json.dumps( + { + "runs": [ + {"throughput_toks": 10.0, "ttft_mean_ms": 1.0}, + {"throughput_toks": 42.0, "ttft_mean_ms": 3.5}, + ] + } + ), + encoding="utf-8", + ) + report = SRTSlurmReportGenerationStrategy(slurm_system, test_run) + + assert report.can_handle_directory() + assert report.get_metric("default") == 42.0 + assert report.get_metric("ttft_mean_ms") == 3.5 + assert tdef.was_run_successful(test_run).is_successful diff --git a/tests/workloads/srt_slurm/test_results.py b/tests/workloads/srt_slurm/test_results.py new file mode 100644 index 000000000..c6698be29 --- /dev/null +++ b/tests/workloads/srt_slurm/test_results.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import csv +import json +from pathlib import Path + +import pytest + +from cloudai.core import METRIC_ERROR, TestRun +from cloudai.systems.slurm import SlurmSystem +from cloudai.workloads.srt_slurm import SRTSlurmReportGenerationStrategy +from cloudai.workloads.srt_slurm.results import agentx_metrics, load_results, sa_bench_metrics + +from .test_srt_slurm import make_tdef + + +def test_fixed_sequence_uses_serving_gpu_counts(tmp_path: Path): + tdef = make_tdef(tmp_path) + tdef.serving_gpus = 32 + tdef.prefill_gpus = 24 + tdef.decode_gpus = 8 + metrics = sa_bench_metrics( + {"total_token_throughput": 300583, "output_throughput": 32000, "median_tpot_ms": 44.59, "median_ttft_ms": 1000}, + tdef, + ) + assert metrics["tput_per_gpu"] == pytest.approx(9393.21875) + assert metrics["output_tput_per_gpu"] == 4000 + assert metrics["input_tput_per_gpu"] == pytest.approx((300583 - 32000) / 24) + assert metrics["median_intvty"] == pytest.approx(1000 / 44.59) + assert metrics["median_ttft"] == 1 + + +@pytest.mark.parametrize( + "runs", + [None, [], [{"throughput_toks": "bad"}], [{"throughput_toks": float("nan")}], [{"throughput_toks": {"bad": 1}}]], +) +def test_bad_rollup_is_unsuccessful(slurm_system: SlurmSystem, tmp_path: Path, runs): + tdef = make_tdef(tmp_path) + tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path) + tdef.job_id_path(tr).write_text("123") + rollup = tdef.rollup_path(tr) + assert rollup is not None + rollup.parent.mkdir(parents=True) + rollup.write_text(json.dumps({"runs": runs})) + assert not tdef.was_run_successful(tr).is_successful + assert SRTSlurmReportGenerationStrategy(slurm_system, tr).get_metric("default") == METRIC_ERROR + + +def test_agentx_requires_job_local_profiles_and_preserves_concurrencies(slurm_system: SlurmSystem, tmp_path: Path): + tdef = make_tdef(tmp_path) + tdef.cmd_args.result_format = "agentx" + tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path) + tdef.job_id_path(tr).write_text("123") + logs = tdef.logs_path(tr) + assert logs is not None + output = logs / "agentic" + output.mkdir(parents=True) + for concurrency, tput, interactivity in [(1, 5484, 518.1), (48, 88086, 42.8)]: + data = { + "conc": concurrency, + "request_metrics": { + "throughput": {"per_gpu": {"total_tput_tps": tput}, "output": {"tokens_per_second": tput * 2}}, + "latency": {"full_response_itl": {"p90": 1 / interactivity}, "itl": {"p90": 999}}, + }, + } + (output / f"result_{concurrency}.json").write_text(json.dumps(data)) + assert agentx_metrics(data)["p90_intvty"] == pytest.approx(interactivity) + assert not tdef.was_run_successful(tr).is_successful + for concurrency in (1, 48): + profile = output / f"conc_{concurrency}" / "aiperf_artifacts" / "profile_export.jsonl" + profile.parent.mkdir(parents=True) + profile.write_text('{"metrics": {}}\n') + assert tdef.was_run_successful(tr).is_successful + assert len(load_results(tdef, tr)) == 2 + report = SRTSlurmReportGenerationStrategy(slurm_system, tr) + assert report.get_metric("tput_per_gpu") == 88086 + assert report.get_metric("default") == 176172 + report.generate_report() + with (tmp_path / "srt_slurm_results.csv").open() as result: + rows = list(csv.DictReader(result)) + assert {float(row["concurrency"]) for row in rows} == {1, 48} diff --git a/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py b/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py new file mode 100644 index 000000000..1b9454546 --- /dev/null +++ b/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import subprocess +import sys +from pathlib import Path + +import pytest + +from cloudai.core import Registry, TestRun, TestScenario +from cloudai.systems.slurm import SlurmRunner, SlurmSystem +from cloudai.workloads.srt_slurm import SRTSlurmCommandGenStrategy + +from .test_srt_slurm import make_tdef + + +@pytest.mark.parametrize("schema", ["", "schema: 2\n"]) +def test_runner_tracks_the_recipe_job(slurm_system: SlurmSystem, tmp_path: Path, schema: str, monkeypatch): + tdef = make_tdef(tmp_path) + tdef.cmd_args.recipe_path.write_text(schema + "name: original\n") + tdef.srtctl.venv_path = tmp_path / "venv" + tdef.srtctl_path.parent.mkdir(parents=True) + repo = slurm_system.install_path / tdef.git_repos[0].repo_name + repo.mkdir(parents=True) + tdef.srtctl_path.write_text(f"#!{sys.executable}\nprint('Submitted batch job 111')\nprint('Job 222 submitted')\n") + tdef.srtctl_path.chmod(0o755) + monkeypatch.chdir(tmp_path) + tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=Path("relative output")) + runner = SlurmRunner("run", slurm_system, TestScenario(name="srt", test_runs=[tr]), tmp_path) + job = runner._submit_test(tr) + assert job.id == 222 + assert tdef.job_id_path(tr).read_text().strip() == "222" + runner.on_job_submit(tr) + assert (tr.output_path / "test-run.toml").is_file() + assert Registry().get_command_gen_strategy(SlurmSystem, type(tdef)) is SRTSlurmCommandGenStrategy + + +def test_failed_submission_does_not_publish_a_job_id(slurm_system: SlurmSystem, tmp_path: Path): + tdef = make_tdef(tmp_path) + tdef.srtctl.venv_path = tmp_path / "venv" + tdef.srtctl_path.parent.mkdir(parents=True) + repo = slurm_system.install_path / tdef.git_repos[0].repo_name + repo.mkdir(parents=True) + tdef.srtctl_path.write_text(f"#!{sys.executable}\nprint('Submitted batch job 111')\nraise SystemExit(2)\n") + tdef.srtctl_path.chmod(0o755) + tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path / "output") + strategy = SRTSlurmCommandGenStrategy(slurm_system, tr) + completed = subprocess.run(strategy.gen_exec_command(), shell=True, capture_output=True, text=True) + assert completed.returncode != 0 + assert "Submitted batch job" not in completed.stdout + assert not tdef.job_id_path(tr).exists() + + +def test_single_sbatch_rejected(slurm_system: SlurmSystem, tmp_path: Path): + tdef = make_tdef(tmp_path) + tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path / "output") + with pytest.raises(ValueError, match="single-sbatch"): + SRTSlurmCommandGenStrategy(slurm_system, tr).gen_srun_command() diff --git a/tests/workloads/srt_slurm/test_srt_slurm.py b/tests/workloads/srt_slurm/test_srt_slurm.py new file mode 100644 index 000000000..3a9656c98 --- /dev/null +++ b/tests/workloads/srt_slurm/test_srt_slurm.py @@ -0,0 +1,260 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest +import yaml +from pydantic import ValidationError + +from cloudai.core import GitRepo, TestRun +from cloudai.systems.slurm import SlurmSystem +from cloudai.workloads.srt_slurm import ( + SRTSlurmCommandGenStrategy, + SRTSlurmTestDefinition, +) + + +def make_tdef(tmp_path: Path) -> SRTSlurmTestDefinition: + recipe_path = tmp_path / "source_recipe.yaml" + recipe_path.write_text( + """\ +schema: 2 +name: demo +model: + path: model + container: image +resources: + gpu_type: h100 + gpus_per_node: 8 +frontend: + type: sglang +engine: sglang +roles: + agg: + nodes: 1 + workers: 1 + gpus: 1 +benchmark: + type: sa-bench + isl: 128 + osl: 128 + concurrencies: "4" +""", + encoding="utf-8", + ) + (tmp_path / "cluster.yaml").write_text("default_partition: main\n") + return SRTSlurmTestDefinition.model_validate( + { + "name": "srt", + "description": "srt-slurm test", + "test_template_name": "SRTSlurm", + "git_repos": [{"url": "git@github.com:NVIDIA/srt-slurm.git", "commit": "test-revision"}], + "cmd_args": { + "recipe_path": recipe_path, + "cluster_config": tmp_path / "cluster.yaml", + "recipe_overrides": {"roles": {"agg": {"gpus": 4}}}, + }, + "extra_cmd_args": {"no-preflight": ""}, + } + ) + + +def test_command_generation_returns_srt_slurm_job(slurm_system: SlurmSystem, tmp_path: Path): + tdef = make_tdef(tmp_path) + tdef.srtctl.venv_path = tmp_path / "venv" + test_run = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path / "output") + strategy = SRTSlurmCommandGenStrategy(slurm_system, test_run) + + command = strategy.gen_exec_command() + run_script = test_run.output_path / "cloudai_generated" / "run_srtctl.sh" + script = run_script.read_text(encoding="utf-8") + + subprocess.run(["bash", "-n", run_script], check=True) + assert command == f"bash {run_script}" + assert str(tdef.cmd_args.recipe_path) in script + assert "--set roles.agg.gpus=4" in script + assert not (run_script.parent / "recipe.yaml").exists() + assert "--no-preflight" in run_script.read_text(encoding="utf-8") + assert "srt_job_id.txt" in run_script.read_text(encoding="utf-8") + assert 'echo "Submitted batch job $SRT_JOB_ID"' in run_script.read_text(encoding="utf-8") + assert "--dependency" not in run_script.read_text(encoding="utf-8") + + +def test_run_script_exports_merged_environment(slurm_system: SlurmSystem, tmp_path: Path, monkeypatch): + monkeypatch.delenv("SRTSLURM_CONFIG", raising=False) + monkeypatch.setenv("VIRTUAL_ENV", str(tmp_path / "login-node-venv")) + slurm_system.global_env_vars = {"SRTSLURM_CONFIG": "system.yaml", "SRT_GLOBAL_ONLY": "global value"} + tdef = make_tdef(tmp_path) + repo_path = tmp_path / "installed srt-slurm" + repo_path.mkdir() + tdef.git_repos[0].installed_path = repo_path + caller_path = tmp_path / "cloudaix" + (caller_path / ".venv" / "bin").mkdir(parents=True) + (caller_path / ".venv" / "bin" / "python").write_bytes(b"invalid login-node interpreter") + (caller_path / ".venv" / "pyvenv.cfg").write_text("home = /login-node/python\n") + active_bin = tmp_path / "login-node-venv" / "bin" + active_bin.mkdir(parents=True) + (active_bin.parent / "pyvenv.cfg").write_text("home = /login-node/python\n") + retained_bin = tmp_path / "cluster tools" / "bin" + retained_bin.mkdir(parents=True) + monkeypatch.setenv("PATH", f"{caller_path / '.venv' / 'bin'}:{active_bin}:{retained_bin}:{os.environ['PATH']}") + tdef.extra_env_vars = { + "SRTSLURM_CONFIG": str(tmp_path / "cluster's config" / "srtslurm.yaml"), + "SRT_LITERAL": "$HOME; $(printf injected)", + "SRT_EMPTY": "", + } + expected_config = str(tdef.cmd_args.cluster_config.absolute()) + tdef.srtctl.venv_path = tmp_path / "venv" + tdef.srtctl_path.parent.mkdir(parents=True) + tdef.srtctl_path.write_text( + f"#!{sys.executable}\n" + "import json, os, sys\n" + "print(json.dumps({key: os.environ.get(key) for key in " + "['SRTSLURM_CONFIG', 'SRT_GLOBAL_ONLY', 'SRT_LITERAL', 'SRT_EMPTY', 'VIRTUAL_ENV']}))\n" + "print(os.getcwd())\n" + "print(os.environ['PATH'])\n" + "print(json.dumps(sys.argv[1:]))\n" + "print('Job 12345 submitted')\n", + encoding="utf-8", + ) + tdef.srtctl_path.chmod(0o755) + test_run = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path / "output") + strategy = SRTSlurmCommandGenStrategy(slurm_system, test_run) + strategy.gen_exec_command() + + subprocess.run( + ["bash", strategy._paths()["run_script"]], cwd=caller_path, check=True, capture_output=True, text=True + ) + received = json.loads(strategy._paths()["run_log"].read_text(encoding="utf-8").splitlines()[0]) + + assert received == { + "SRT_GLOBAL_ONLY": "global value", + **tdef.extra_env_vars, + "SRTSLURM_CONFIG": expected_config, + "VIRTUAL_ENV": None, + } + assert Path(strategy._paths()["run_log"].read_text().splitlines()[1]) == repo_path + submitted_path = strategy._paths()["run_log"].read_text().splitlines()[2].split(os.pathsep) + assert str(caller_path / ".venv" / "bin") not in submitted_path + assert str(active_bin) not in submitted_path + assert str(retained_bin) in submitted_path + submitted_args = json.loads(strategy._paths()["run_log"].read_text().splitlines()[3]) + assert submitted_args[submitted_args.index("-f") + 1] == str(tdef.cmd_args.recipe_path) + assert submitted_args[submitted_args.index("--set") + 1] == "roles.agg.gpus=4" + assert tdef.job_id_path(test_run).read_text(encoding="utf-8").strip() == "12345" + + +def test_existing_recipe_is_loaded_and_dse_values_are_overlaid(tmp_path: Path): + source = tmp_path / "existing.yaml" + original = """\ +schema: 2 +name: existing +roles: + agg: + gpus: 1 + args: + tensor-parallel-size: 1 +benchmark: + concurrencies: [4, 8] +""" + source.write_text(original, encoding="utf-8") + tdef = SRTSlurmTestDefinition.model_validate( + { + "name": "external", + "description": "existing recipe", + "test_template_name": "SRTSlurm", + "git_repos": [{"url": "git@github.com:NVIDIA/srt-slurm.git", "commit": "test-revision"}], + "cmd_args": { + "recipe_path": source, + "cluster_config": tmp_path / "cluster.yaml", + "recipe_overrides": { + "roles": {"agg": {"gpus": [1, 2, 4]}}, + "benchmark": {"concurrencies": [[32, 64], [64, 128]]}, + }, + }, + } + ) + + test_run = TestRun(name="dse", test=tdef, num_nodes=1, nodes=[]) + assert test_run.param_space == { + "recipe_overrides.roles.agg.gpus": [1, 2, 4], + "recipe_overrides.benchmark.concurrencies": [[32, 64], [64, 128]], + } + selected = test_run.apply_params_set( + { + "recipe_overrides.roles.agg.gpus": 4, + "recipe_overrides.benchmark.concurrencies": [32, 64], + } + ) + assert isinstance(selected.test, SRTSlurmTestDefinition) + assert SRTSlurmCommandGenStrategy._override_args(selected.test.cmd_args.recipe_overrides) == [ + "--set", + "roles.agg.gpus=4", + "--set", + "benchmark.concurrencies=[32, 64]", + ] + assert source.read_text(encoding="utf-8") == original + assert tdef.dse_excluded_args == ["cmd_args.recipe_path", "cmd_args.cluster_config", "cmd_args.result_format"] + assert "cmd_args.recipe_overrides" not in tdef.dse_excluded_args + + +def test_standard_git_repos_controls_srtctl_revision(tmp_path: Path): + config = make_tdef(tmp_path).model_dump() + config["extra_env_vars"] = {"SRTSLURM_CONFIG": str(tmp_path / "cluster.yaml")} + config["git_repos"][0]["commit"] = "updated-revision" + config["git_repos"].append({"url": "https://github.com/example/recipes.git", "commit": "main"}) + tdef = SRTSlurmTestDefinition.model_validate(config) + + assert tdef.srtctl.git_repo is tdef.git_repos[0] + assert tdef.srtctl.git_repo.commit == "updated-revision" + assert "updated-revision" in tdef.srtctl.venv_name + assert tdef.installables == [tdef.srtctl, tdef.git_repos[1]] + assert isinstance(tdef.installables[1], GitRepo) + + +@pytest.mark.parametrize("git_repos", [None, []]) +def test_srtctl_requires_a_configured_git_repo(tmp_path: Path, git_repos): + config = make_tdef(tmp_path).model_dump() + if git_repos is None: + config.pop("git_repos") + else: + config["git_repos"] = git_repos + + with pytest.raises(ValidationError, match="git_repos"): + SRTSlurmTestDefinition.model_validate(config) + + +@pytest.mark.parametrize("value", [4, 0.75, True, False, None, [32, 64], "true", "720", "", "$HOME; $(id)"]) +def test_set_arguments_preserve_yaml_value_types(value): + args = SRTSlurmCommandGenStrategy._override_args({"roles": {"agg": {"args": {"flag": value}}}}) + assert args[0] == "--set" + key, raw = args[1].split("=", 1) + assert key == "roles.agg.args.flag" + parsed = yaml.safe_load(raw) + assert parsed == value + assert type(parsed) is type(value) + + +def test_set_arguments_quote_literal_dotted_keys(): + assert SRTSlurmCommandGenStrategy._override_args({"container_mounts": {"/models/model.v1": "/model"}}) == [ + "--set", + 'container_mounts."/models/model.v1"="/model"', + ] From 0c779d369d0644f709fa183e11d60f7b904a6c1c Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Tue, 13 Oct 2026 07:37:52 -0600 Subject: [PATCH 02/20] use .json GPU counts to compute num_gpu normalized metrics --- conf/experimental/srt_slurm/scenario.toml | 1 - doc/workloads/srt_slurm.rst | 41 +++- src/cloudai/report_generator/dse_report.py | 2 +- src/cloudai/reporter.py | 2 +- .../slurm/slurm_command_gen_strategy.py | 2 +- .../srt_slurm/report_generation_strategy.py | 8 +- src/cloudai/workloads/srt_slurm/results.py | 115 +++++++-- .../srt_slurm/slurm_command_gen_strategy.py | 32 ++- src/cloudai/workloads/srt_slurm/srt_slurm.py | 36 ++- tests/test_reporter.py | 4 +- .../test_report_generation_strategy.py | 3 +- tests/workloads/srt_slurm/test_results.py | 226 +++++++++++++++++- .../test_slurm_command_gen_strategy.py | 56 ++++- tests/workloads/srt_slurm/test_srt_slurm.py | 23 +- 14 files changed, 486 insertions(+), 65 deletions(-) diff --git a/conf/experimental/srt_slurm/scenario.toml b/conf/experimental/srt_slurm/scenario.toml index 202e006e6..05dd9ce42 100644 --- a/conf/experimental/srt_slurm/scenario.toml +++ b/conf/experimental/srt_slurm/scenario.toml @@ -25,7 +25,6 @@ description = "SGLang through an existing srt-slurm recipe" test_template_name = "SRTSlurm" agent = "grid_search" agent_metrics = ["tput_per_gpu"] -serving_gpus = 1 [[Tests.git_repos]] url = "https://github.com/NVIDIA/srt-slurm.git" diff --git a/doc/workloads/srt_slurm.rst b/doc/workloads/srt_slurm.rst index 8d4b41ec8..7b47aa93a 100644 --- a/doc/workloads/srt_slurm.rst +++ b/doc/workloads/srt_slurm.rst @@ -21,7 +21,6 @@ and must be accessible on the shared filesystem. description = "SGLang with srt-slurm" test_template_name = "SRTSlurm" agent_metrics = ["tput_per_gpu"] - serving_gpus = 4 [[git_repos]] url = "https://github.com/NVIDIA/srt-slurm.git" @@ -37,6 +36,11 @@ and must be accessible on the shared filesystem. The first ``git_repos`` entry supplies srtctl; additional entries use ordinary CloudAI installation. Pin a revision supporting ``srtctl apply --set``. +For older revisions without ``--set``, set ``cmd_args.recipe_override_mode = "render"``. +This overlays existing mapping keys in a concrete recipe, preserves YAML value types, +and submits a trial-local ``cloudai_generated/recipe.yaml``. It rejects missing keys +and does not resolve recipe variants or interpret native override path expressions; +select the effective concrete recipe first. The default remains ``"native"``. Lists in ``recipe_overrides`` are DSE dimensions. Use nested candidate lists when the selected recipe value is itself a list. Recipe paths, cluster config paths, and result format are excluded from DSE. Additional apply flags use @@ -65,12 +69,27 @@ Results ``result_format = "sa-bench"`` reads job-local ``sa-bench_*/results_*.json`` files, falling back to ``benchmark-rollup.json`` when raw results are absent. -``default`` and ``output_throughput`` expose output tokens/s. Set -``serving_gpus`` to the actual serving GPU count to enable ``tput_per_gpu`` -from total token throughput, including input tokens. For disaggregated serving, -also set ``prefill_gpus`` and ``decode_gpus`` for the input/output per-GPU -divisors. Infrastructure and idle GPUs do not belong in these counts; shared -aggregate GPUs count once. A rollup lacking total throughput cannot provide +Both result formats require the submitted job's ``.json`` next to its +``logs`` directory. Its ``job_id`` must match the tracked job, and ``resources`` +must contain nonnegative integer worker counts for ``agg``, ``prefill`` and +``decode``, plus positive integer ``gpus_per_`` for each active role. +Missing, malformed, mismatched or incomplete metadata makes the trial unsuccessful; +there is no fallback to manual counts, filenames, node capacity or resource snapshots. +Use a pinned revision containing the resolved GPU metadata feature (introduced in +commit ``7858d3091d5e278ab56ebf0246df9ea913bc2477``, first released in ``v1.0.0``). + +CloudAI calculates each role's GPU total as workers times GPUs per worker and sums +those totals for the serving count. This automatically follows each trial's effective +DSE topology. ``serving_gpus``, ``prefill_gpus`` and ``decode_gpus`` are no longer +scenario fields; remove them from existing configs. Resolved counts appear in each +CSV result row as ``serving_gpus``, ``aggregate_gpus``, ``prefill_gpus`` and ``decode_gpus``. + +``default`` and ``output_throughput`` expose output tokens/s; +``total_token_throughput`` exposes input plus output tokens/s. +``tput_per_gpu`` divides total throughput by the metadata-derived serving count. +Disaggregated input/output per-GPU metrics use their respective prefill/decode pools; +aggregated metrics share the aggregate pool, counted once. Producer-normalized +rollup fields are ignored. A rollup lacking raw total throughput cannot provide ``tput_per_gpu``. Latency aliases such as ``median_ttft`` use seconds; ``ttft_median_ms`` uses milliseconds. ``median_intvty`` and ``p90_intvty`` use the reciprocal of the corresponding TPOT, in tokens/s/user. @@ -78,8 +97,12 @@ use the reciprocal of the corresponding TPOT, in tokens/s/user. ``result_format = "agentx"`` reads InferenceX aggregate JSON from the submitted job's ``logs/agentic`` directory. Each aggregate must have its corresponding ``conc_/aiperf_artifacts/profile_export.jsonl``. CloudAI reads the existing -aggregation rather than recomputing per-request statistics. Throughput uses -``request_metrics.throughput.per_gpu``; interactivity uses full-response ITL +aggregation rather than recomputing per-request statistics. CloudAI recomputes +per-GPU throughput from raw ``request_metrics.throughput`` total/input/output values +using the job metadata's total serving count for all three streams. Producer +``per_gpu`` values are ignored, since environment-derived divisors can become stale +when DSE changes topology. An aggregate without raw total throughput is unusable. +Interactivity uses full-response ITL when available, otherwise ITL. Latency aliases use the same units as sa-bench. An AgentX recipe uses ``benchmark.type: custom`` and its existing InferenceX diff --git a/src/cloudai/report_generator/dse_report.py b/src/cloudai/report_generator/dse_report.py index ddd0b8165..fa4cbed10 100644 --- a/src/cloudai/report_generator/dse_report.py +++ b/src/cloudai/report_generator/dse_report.py @@ -341,7 +341,7 @@ def _build_iteration_summary( observed_runtime=format_duration(total_observed_runtime_sec), efficiency_ratio=f"~{format_float(reduction_factor, 1)}x", efficiency_steps=f"{len(trajectory_steps):,} / {total_space:,} steps", - best_config_toml=toml.dumps(test_run_details.test_definition.model_dump()), + best_config_toml=toml.dumps(test_run_details.test_definition.model_dump(mode="json")), parameter_rows=_build_parameter_rows(test_case.param_space, best_step.action), reward_chart_data=_build_reward_chart_data(trajectory_steps), ) diff --git a/src/cloudai/reporter.py b/src/cloudai/reporter.py index 918fe26fb..86d8e7dca 100644 --- a/src/cloudai/reporter.py +++ b/src/cloudai/reporter.py @@ -293,7 +293,7 @@ def report_best_dse_config(self): best_config_path = tr_root / f"{tr.name}.toml" logging.info("Writing best config for %s to %s", tr.name, best_config_path) with best_config_path.open("w") as f: - toml.dump(trd.test_definition.model_dump(), f) + toml.dump(trd.test_definition.model_dump(mode="json"), f) class TarballReporter(Reporter): diff --git a/src/cloudai/systems/slurm/slurm_command_gen_strategy.py b/src/cloudai/systems/slurm/slurm_command_gen_strategy.py index 52343357c..15413bdcc 100644 --- a/src/cloudai/systems/slurm/slurm_command_gen_strategy.py +++ b/src/cloudai/systems/slurm/slurm_command_gen_strategy.py @@ -74,7 +74,7 @@ def store_test_run(self) -> None: test_cmd, srun_cmd = (" ".join(self.generate_test_command()), self.gen_srun_command()) with (self.test_run.output_path / self.TEST_RUN_DUMP_FILE_NAME).open("w") as f: trd = TestRunDetails.from_test_run(self.test_run, test_cmd=test_cmd, full_cmd=srun_cmd) - toml.dump(trd.model_dump(), f) + toml.dump(trd.model_dump(mode="json"), f) @final def container_mounts(self) -> list[str]: diff --git a/src/cloudai/workloads/srt_slurm/report_generation_strategy.py b/src/cloudai/workloads/srt_slurm/report_generation_strategy.py index 2ee62d92b..80309dc17 100644 --- a/src/cloudai/workloads/srt_slurm/report_generation_strategy.py +++ b/src/cloudai/workloads/srt_slurm/report_generation_strategy.py @@ -15,6 +15,7 @@ # limitations under the License. import csv +import logging from typing import ClassVar, cast from cloudai.core import METRIC_ERROR, MetricValue, ReportGenerationStrategy @@ -30,6 +31,7 @@ class SRTSlurmReportGenerationStrategy(ReportGenerationStrategy): "default", "output_throughput", "throughput_toks", + "total_token_throughput", "request_throughput", "tput_per_gpu", "input_tput_per_gpu", @@ -48,7 +50,11 @@ class SRTSlurmReportGenerationStrategy(ReportGenerationStrategy): ] def results(self) -> list[dict[str, float]]: - return load_results(cast(SRTSlurmTestDefinition, self.test_run.test), self.test_run) + try: + return load_results(cast(SRTSlurmTestDefinition, self.test_run.test), self.test_run) + except ValueError as error: + logging.warning("Unable to read srt-slurm results: %s", error) + return [] def can_handle_directory(self) -> bool: return bool(self.results()) diff --git a/src/cloudai/workloads/srt_slurm/results.py b/src/cloudai/workloads/srt_slurm/results.py index 2d9df819a..dd4fa0b4f 100644 --- a/src/cloudai/workloads/srt_slurm/results.py +++ b/src/cloudai/workloads/srt_slurm/results.py @@ -16,6 +16,7 @@ import json import math +from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Any @@ -43,7 +44,57 @@ def numeric_fields(data: dict[str, Any]) -> dict[str, float]: return {key: value for key, raw in data.items() if (value := number(raw)) is not None} -def sa_bench_metrics(data: dict[str, Any], tdef: "SRTSlurmTestDefinition") -> dict[str, float]: +@dataclass(frozen=True) +class JobGPUCounts: + """Backend GPU pools resolved by srt-slurm for one submitted job.""" + + aggregate: int + prefill: int + decode: int + + @property + def serving(self) -> int: + return self.aggregate + self.prefill + self.decode + + def as_metrics(self) -> dict[str, float]: + return { + "serving_gpus": float(self.serving), + "aggregate_gpus": float(self.aggregate), + "prefill_gpus": float(self.prefill), + "decode_gpus": float(self.decode), + } + + +def load_gpu_counts(path: Path) -> JobGPUCounts: + """Read required producer-resolved counts; never infer them from node capacity.""" + if not path.is_file(): + raise ValueError(f"Missing required srt-slurm job metadata: {path}") + metadata = read_json(path) + job_id = metadata.get("job_id") + if not isinstance(job_id, (str, int)) or isinstance(job_id, bool) or str(job_id) != path.stem: + raise ValueError(f"Invalid or mismatched job_id in srt-slurm job metadata: {path}") + resources = metadata.get("resources") + if not isinstance(resources, dict): + raise ValueError(f"Missing resources in srt-slurm job metadata: {path}") + counts = {} + for role in ("agg", "prefill", "decode"): + workers = resources.get(f"{role}_workers") + if type(workers) is not int or workers < 0: + raise ValueError(f"Invalid resources.{role}_workers in srt-slurm job metadata: {path}") + if workers == 0: + counts[role] = 0 + continue + per_worker = resources.get(f"gpus_per_{role}") + if type(per_worker) is not int or per_worker <= 0: + raise ValueError(f"Invalid resources.gpus_per_{role} in srt-slurm job metadata: {path}") + counts[role] = workers * per_worker + result = JobGPUCounts(aggregate=counts["agg"], prefill=counts["prefill"], decode=counts["decode"]) + if result.serving <= 0: + raise ValueError(f"No serving GPUs in srt-slurm job metadata: {path}") + return result + + +def sa_bench_metrics(data: dict[str, Any], gpus: JobGPUCounts) -> dict[str, float]: metrics = numeric_fields(data) aliases = {"max_concurrency": "concurrency", "output_throughput": "throughput_toks"} for old, new in aliases.items(): @@ -52,12 +103,16 @@ def sa_bench_metrics(data: dict[str, Any], tdef: "SRTSlurmTestDefinition") -> di if "throughput_toks" not in metrics: return {} metrics["output_throughput"] = metrics["throughput_toks"] - if tdef.serving_gpus and "total_token_throughput" in metrics: - metrics["tput_per_gpu"] = metrics["total_token_throughput"] / tdef.serving_gpus - output_gpus = tdef.decode_gpus or tdef.serving_gpus - metrics["output_tput_per_gpu"] = metrics["throughput_toks"] / output_gpus - input_gpus = tdef.prefill_gpus or tdef.serving_gpus - metrics["input_tput_per_gpu"] = (metrics["total_token_throughput"] - metrics["throughput_toks"]) / input_gpus + # Producer rollups may contain capacity-based normalization; use only raw measurements. + for key in ("tput_per_gpu", "output_tput_per_gpu", "input_tput_per_gpu"): + metrics.pop(key, None) + metrics.update(gpus.as_metrics()) + metrics["output_tput_per_gpu"] = metrics["throughput_toks"] / (gpus.decode or gpus.serving) + if "total_token_throughput" in metrics: + metrics["tput_per_gpu"] = metrics["total_token_throughput"] / gpus.serving + metrics["input_tput_per_gpu"] = (metrics["total_token_throughput"] - metrics["throughput_toks"]) / ( + gpus.prefill or gpus.serving + ) add_sa_latencies(data, metrics) return metrics @@ -86,29 +141,35 @@ def add_sa_latencies(data: dict[str, Any], metrics: dict[str, float]) -> None: metrics[f"{stat}_intvty"] = 1000 / value -def agentx_metrics(data: dict[str, Any]) -> dict[str, float]: +def agentx_metrics(data: dict[str, Any], gpus: JobGPUCounts) -> dict[str, float]: request = data.get("request_metrics") if not isinstance(request, dict): return {} - throughput = request.get("throughput", {}) - per_gpu = throughput.get("per_gpu", {}) if isinstance(throughput, dict) else {} - if not isinstance(per_gpu, dict) or number(per_gpu.get("total_tput_tps")) is None: + throughput = request.get("throughput") + if not isinstance(throughput, dict): return {} metrics = numeric_fields(data) if "conc" in metrics: metrics["concurrency"] = metrics["conc"] - for source, dest in ( - ("total_tput_tps", "tput_per_gpu"), - ("output_tput_tps", "output_tput_per_gpu"), - ("input_tput_tps", "input_tput_per_gpu"), + # AgentX normalizes every stream by total serving GPUs, including disaggregated runs. + # Its per_gpu values can reflect stale environment metadata after topology DSE. + for key in ("tput_per_gpu", "output_tput_per_gpu", "input_tput_per_gpu"): + metrics.pop(key, None) + metrics.update(gpus.as_metrics()) + for source, raw_key, normalized_key in ( + ("total", "total_tput_tps", "tput_per_gpu"), + ("input", "input_tput_tps", "input_tput_per_gpu"), + ("output", "output_tput_tps", "output_tput_per_gpu"), ): - if (value := number(per_gpu.get(source))) is not None: - metrics[dest] = value - metrics.update(numeric_fields(throughput)) - for source, dest in (("total", "total_tput_tps"), ("input", "input_tput_tps"), ("output", "output_tput_tps")): - block = throughput.get(source, {}) - if isinstance(block, dict) and (value := number(block.get("tokens_per_second"))) is not None: - metrics[dest] = value + block = throughput.get(source) + value = number(block.get("tokens_per_second")) if isinstance(block, dict) else None + if value is None: + value = number(throughput.get(raw_key)) + if value is not None: + metrics[raw_key] = value + metrics[normalized_key] = value / gpus.serving + if "tput_per_gpu" not in metrics: + return {} if "output_tput_tps" in metrics: metrics["output_throughput"] = metrics["throughput_toks"] = metrics["output_tput_tps"] add_agentx_latencies(request.get("latency", {}), metrics) @@ -136,10 +197,14 @@ def load_results(tdef: "SRTSlurmTestDefinition", tr: TestRun) -> list[dict[str, logs = tdef.logs_path(tr) if logs is None: return [] + metadata = tdef.metadata_path(tr) + if metadata is None: + raise ValueError("No submitted srt-slurm job ID for required job metadata") + gpus = load_gpu_counts(metadata) if tdef.cmd_args.result_format == "agentx": results = [] for path in sorted((logs / "agentic").glob("*.json")): - metrics = agentx_metrics(read_json(path)) + metrics = agentx_metrics(read_json(path), gpus) concurrency = metrics.get("concurrency") if concurrency is None or not concurrency.is_integer(): continue @@ -150,8 +215,8 @@ def load_results(tdef: "SRTSlurmTestDefinition", tr: TestRun) -> list[dict[str, return results raw_files = sorted(logs.glob("sa-bench_*/results_*.json")) if raw_files: - return [metrics for path in raw_files if (metrics := sa_bench_metrics(read_json(path), tdef))] + return [metrics for path in raw_files if (metrics := sa_bench_metrics(read_json(path), gpus))] runs = read_json(logs / "benchmark-rollup.json").get("runs", []) if not isinstance(runs, list): return [] - return [metrics for run in runs if isinstance(run, dict) and (metrics := sa_bench_metrics(run, tdef))] + return [metrics for run in runs if isinstance(run, dict) and (metrics := sa_bench_metrics(run, gpus))] diff --git a/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py b/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py index f736f04ce..61eb718f5 100644 --- a/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py +++ b/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py @@ -77,16 +77,24 @@ def write_run_script(self) -> None: repo = self.tdef.git_repos[0] repo_path = (repo.installed_path or self.system.install_path / repo.repo_name).absolute() cluster_config = self.tdef.cmd_args.cluster_config.expanduser().absolute() + recipe = paths["recipe"] + overrides = self._override_args(self.tdef.cmd_args.recipe_overrides) + if self.tdef.cmd_args.recipe_override_mode == "render": + data = self._load_recipe(recipe) + self._overlay_recipe(data, self.tdef.cmd_args.recipe_overrides) + recipe = paths["generated"] / "recipe.yaml" + recipe.write_text(yaml.safe_dump(data, sort_keys=False), encoding="utf-8") + overrides = [] apply_command = shlex.join( [ str(self.tdef.srtctl_path), "apply", "-f", - str(paths["recipe"]), + str(recipe), "-o", str(paths["nested_outputs"]), "-y", - *self._override_args(self.tdef.cmd_args.recipe_overrides), + *overrides, *self._extra_cmd_arg_parts(), ] ) @@ -128,6 +136,20 @@ def write_run_script(self) -> None: ) paths["run_script"].write_text(script, encoding="utf-8") + @staticmethod + def _overlay_recipe(data: dict[str, Any], overrides: dict[str, Any], prefix: str = "") -> None: + """Overlay existing keys of a concrete recipe for older CLIs without --set.""" + for key, value in overrides.items(): + path = f"{prefix}.{key}" if prefix else key + if key not in data: + raise ValueError(f"Rendered recipe override targets an absent key: {path}") + if isinstance(value, dict): + if not isinstance(data[key], dict): + raise ValueError(f"Rendered recipe override targets a non-mapping: {path}") + SRTSlurmCommandGenStrategy._overlay_recipe(data[key], value, path) + else: + data[key] = value + def _extra_cmd_arg_parts(self) -> list[str]: args = [] for name, value in self.tdef.extra_cmd_args.items(): @@ -139,6 +161,12 @@ def _extra_cmd_arg_parts(self) -> list[str]: def _container_mounts(self) -> list[str]: return [] + def store_test_run(self) -> None: + # SlurmRunner records command metadata before submitting the test. + # Prepare the wrapper first, without executing srtctl. + self.gen_exec_command() + super().store_test_run() + def gen_exec_command(self) -> str: self.test_run.output_path.mkdir(parents=True, exist_ok=True) self._load_recipe(self._paths()["recipe"]) diff --git a/src/cloudai/workloads/srt_slurm/srt_slurm.py b/src/cloudai/workloads/srt_slurm/srt_slurm.py index eb719a903..0cdb3d27c 100644 --- a/src/cloudai/workloads/srt_slurm/srt_slurm.py +++ b/src/cloudai/workloads/srt_slurm/srt_slurm.py @@ -31,19 +31,24 @@ class SRTSlurmCmdArgs(CmdArgs): recipe_path: Path cluster_config: Path result_format: Literal["sa-bench", "agentx"] = "sa-bench" + # native passes typed --set arguments to srtctl; render writes a per-trial + # recipe copy for older CLIs without --set, replacing existing keys only. + recipe_override_mode: Literal["native", "render"] = "native" recipe_overrides: dict[str, Any] = Field(default_factory=dict) def excluded_cmd_args(self) -> list[str]: - return ["cmd_args.recipe_path", "cmd_args.cluster_config", "cmd_args.result_format"] + return [ + "cmd_args.recipe_path", + "cmd_args.cluster_config", + "cmd_args.result_format", + "cmd_args.recipe_override_mode", + ] class SRTSlurmTestDefinition(TestDefinition): """Run one declarative srt-slurm recipe as one CloudAI trial.""" cmd_args: SRTSlurmCmdArgs - serving_gpus: int | None = Field(default=None, gt=0) - prefill_gpus: int | None = Field(default=None, gt=0) - decode_gpus: int | None = Field(default=None, gt=0) git_repos: list[GitRepo] = Field( default_factory=list, min_length=1, validate_default=True, description="The first repository supplies srtctl." ) @@ -74,22 +79,35 @@ def installables(self) -> list[Installable]: def job_id_path(tr: TestRun) -> Path: return tr.output_path / "srt_job_id.txt" - def rollup_path(self, tr: TestRun) -> Path | None: + def job_output_path(self, tr: TestRun) -> Path | None: + """Return the submitted job's output directory, without guessing its identity.""" try: job_id = self.job_id_path(tr).read_text(encoding="utf-8").strip() except OSError: return None if not job_id.isdigit(): return None - return tr.output_path / "srt_slurm_outputs" / job_id / "logs" / "benchmark-rollup.json" + return tr.output_path / "srt_slurm_outputs" / job_id + + def metadata_path(self, tr: TestRun) -> Path | None: + output = self.job_output_path(tr) + return output / f"{output.name}.json" if output else None + + def rollup_path(self, tr: TestRun) -> Path | None: + logs = self.logs_path(tr) + return logs / "benchmark-rollup.json" if logs else None def was_run_successful(self, tr: TestRun) -> JobStatusResult: from .results import load_results - if not load_results(self, tr): + try: + results = load_results(self, tr) + except ValueError as error: + return JobStatusResult(is_successful=False, error_message=str(error)) + if not results: return JobStatusResult(is_successful=False, error_message="No usable job-local srt-slurm results found.") return JobStatusResult(is_successful=True) def logs_path(self, tr: TestRun) -> Path | None: - rollup = self.rollup_path(tr) - return rollup.parent if rollup else None + output = self.job_output_path(tr) + return output / "logs" if output else None diff --git a/tests/test_reporter.py b/tests/test_reporter.py index 69bdcd33f..0788baa4e 100644 --- a/tests/test_reporter.py +++ b/tests/test_reporter.py @@ -658,7 +658,9 @@ def test_dse_reporter( "observed_runtime": "6m", "efficiency_ratio": "~1.3x", "efficiency_steps": "3 / 4 steps", - "best_config_toml": toml.dumps(TestRunDetails.from_test_run(best_tr, "", "").test_definition.model_dump()), + "best_config_toml": toml.dumps( + TestRunDetails.from_test_run(best_tr, "", "").test_definition.model_dump(mode="json") + ), "parameter_rows": [ { "name": "ngpus", diff --git a/tests/workloads/srt_slurm/test_report_generation_strategy.py b/tests/workloads/srt_slurm/test_report_generation_strategy.py index a8fa8b2b2..85e2f4ed2 100644 --- a/tests/workloads/srt_slurm/test_report_generation_strategy.py +++ b/tests/workloads/srt_slurm/test_report_generation_strategy.py @@ -21,7 +21,7 @@ from cloudai.systems.slurm import SlurmSystem from cloudai.workloads.srt_slurm import SRTSlurmReportGenerationStrategy -from .test_srt_slurm import make_tdef +from .test_srt_slurm import make_tdef, write_job_metadata def test_report_reads_best_run_directly_from_srt_rollup(slurm_system: SlurmSystem, tmp_path: Path): @@ -29,6 +29,7 @@ def test_report_reads_best_run_directly_from_srt_rollup(slurm_system: SlurmSyste test_run = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path / "output") test_run.output_path.mkdir(parents=True) tdef.job_id_path(test_run).write_text("12345\n", encoding="utf-8") + write_job_metadata(tdef, test_run) rollup_path = test_run.output_path / "srt_slurm_outputs" / "12345" / "logs" / "benchmark-rollup.json" rollup_path.parent.mkdir(parents=True) rollup_path.write_text( diff --git a/tests/workloads/srt_slurm/test_results.py b/tests/workloads/srt_slurm/test_results.py index c6698be29..65eeb4ce6 100644 --- a/tests/workloads/srt_slurm/test_results.py +++ b/tests/workloads/srt_slurm/test_results.py @@ -22,20 +22,23 @@ from cloudai.core import METRIC_ERROR, TestRun from cloudai.systems.slurm import SlurmSystem -from cloudai.workloads.srt_slurm import SRTSlurmReportGenerationStrategy -from cloudai.workloads.srt_slurm.results import agentx_metrics, load_results, sa_bench_metrics +from cloudai.workloads.srt_slurm import SRTSlurmReportGenerationStrategy, SRTSlurmTestDefinition +from cloudai.workloads.srt_slurm.results import ( + JobGPUCounts, + agentx_metrics, + load_gpu_counts, + load_results, + sa_bench_metrics, +) -from .test_srt_slurm import make_tdef +from .test_srt_slurm import make_tdef, write_job_metadata def test_fixed_sequence_uses_serving_gpu_counts(tmp_path: Path): - tdef = make_tdef(tmp_path) - tdef.serving_gpus = 32 - tdef.prefill_gpus = 24 - tdef.decode_gpus = 8 + gpus = JobGPUCounts(aggregate=0, prefill=24, decode=8) metrics = sa_bench_metrics( {"total_token_throughput": 300583, "output_throughput": 32000, "median_tpot_ms": 44.59, "median_ttft_ms": 1000}, - tdef, + gpus, ) assert metrics["tput_per_gpu"] == pytest.approx(9393.21875) assert metrics["output_tput_per_gpu"] == 4000 @@ -44,6 +47,26 @@ def test_fixed_sequence_uses_serving_gpu_counts(tmp_path: Path): assert metrics["median_ttft"] == 1 +def test_agentx_recomputes_normalization_from_job_counts(): + data = { + "conc": 1, + "request_metrics": { + "throughput": { + "per_gpu": {}, + "total": {"tokens_per_second": 14407}, + "output": {"tokens_per_second": 112}, + }, + "latency": {"full_response_itl": {"p90": 0.00554}, "tpot": {"p90": 0.006}}, + }, + } + metrics = agentx_metrics(data, JobGPUCounts(aggregate=8, prefill=0, decode=0)) + assert metrics["tput_per_gpu"] == 14407 / 8 + assert metrics["output_tput_per_gpu"] == 14 + assert metrics["p90_intvty"] == pytest.approx(1 / 0.00554) + data["request_metrics"]["throughput"]["per_gpu"] = {"total_tput_tps": 1801} + assert agentx_metrics(data, JobGPUCounts(aggregate=16, prefill=0, decode=0))["tput_per_gpu"] == 14407 / 16 + + @pytest.mark.parametrize( "runs", [None, [], [{"throughput_toks": "bad"}], [{"throughput_toks": float("nan")}], [{"throughput_toks": {"bad": 1}}]], @@ -52,6 +75,7 @@ def test_bad_rollup_is_unsuccessful(slurm_system: SlurmSystem, tmp_path: Path, r tdef = make_tdef(tmp_path) tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path) tdef.job_id_path(tr).write_text("123") + write_job_metadata(tdef, tr) rollup = tdef.rollup_path(tr) assert rollup is not None rollup.parent.mkdir(parents=True) @@ -65,6 +89,7 @@ def test_agentx_requires_job_local_profiles_and_preserves_concurrencies(slurm_sy tdef.cmd_args.result_format = "agentx" tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path) tdef.job_id_path(tr).write_text("123") + write_job_metadata(tdef, tr) logs = tdef.logs_path(tr) assert logs is not None output = logs / "agentic" @@ -73,12 +98,18 @@ def test_agentx_requires_job_local_profiles_and_preserves_concurrencies(slurm_sy data = { "conc": concurrency, "request_metrics": { - "throughput": {"per_gpu": {"total_tput_tps": tput}, "output": {"tokens_per_second": tput * 2}}, + "throughput": { + "per_gpu": {"total_tput_tps": 999999}, + "total": {"tokens_per_second": tput * 4}, + "output": {"tokens_per_second": tput * 2}, + }, "latency": {"full_response_itl": {"p90": 1 / interactivity}, "itl": {"p90": 999}}, }, } (output / f"result_{concurrency}.json").write_text(json.dumps(data)) - assert agentx_metrics(data)["p90_intvty"] == pytest.approx(interactivity) + assert agentx_metrics(data, JobGPUCounts(aggregate=4, prefill=0, decode=0))["p90_intvty"] == pytest.approx( + interactivity + ) assert not tdef.was_run_successful(tr).is_successful for concurrency in (1, 48): profile = output / f"conc_{concurrency}" / "aiperf_artifacts" / "profile_export.jsonl" @@ -93,3 +124,178 @@ def test_agentx_requires_job_local_profiles_and_preserves_concurrencies(slurm_sy with (tmp_path / "srt_slurm_results.csv").open() as result: rows = list(csv.DictReader(result)) assert {float(row["concurrency"]) for row in rows} == {1, 48} + + +@pytest.mark.parametrize( + "resources, expected", + [ + ({"agg_workers": 1, "gpus_per_agg": 1}, JobGPUCounts(1, 0, 0)), + ({"agg_workers": 2, "gpus_per_agg": 2}, JobGPUCounts(4, 0, 0)), + ({"agg_workers": 1, "gpus_per_agg": 16, "agg_nodes": 2}, JobGPUCounts(16, 0, 0)), + ( + {"agg_workers": 0, "prefill_workers": 3, "gpus_per_prefill": 8, "decode_workers": 1, "gpus_per_decode": 8}, + JobGPUCounts(0, 24, 8), + ), + ( + { + "agg_workers": 0, + "prefill_workers": 1, + "gpus_per_prefill": 4, + "decode_workers": 1, + "gpus_per_decode": 4, + "decode_nodes": 0, + }, + JobGPUCounts(0, 4, 4), + ), + ], +) +def test_job_metadata_resolves_worker_gpu_totals(tmp_path: Path, resources: dict, expected: JobGPUCounts): + tdef = make_tdef(tmp_path) + tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path) + tdef.job_id_path(tr).write_text("123") + metadata = write_job_metadata(tdef, tr, **resources) + assert load_gpu_counts(metadata) == expected + + +@pytest.mark.parametrize("agentic", [False, True]) +def test_missing_metadata_fails_even_with_results(slurm_system: SlurmSystem, tmp_path: Path, agentic: bool): + tdef = make_tdef(tmp_path) + tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path) + tdef.job_id_path(tr).write_text("123") + rollup = tdef.rollup_path(tr) + assert rollup is not None + rollup.parent.mkdir(parents=True) + rollup.write_text(json.dumps({"runs": [{"throughput_toks": 100, "tput_per_gpu": 50}]})) + if agentic: + tdef.cmd_args.result_format = "agentx" + output = rollup.parent / "agentic" + output.mkdir() + (output / "result.json").write_text( + json.dumps( + { + "conc": 1, + "request_metrics": { + "throughput": {"total": {"tokens_per_second": 100}, "per_gpu": {"total_tput_tps": 50}} + }, + } + ) + ) + profile = output / "conc_1/aiperf_artifacts/profile_export.jsonl" + profile.parent.mkdir(parents=True) + profile.write_text("{}\n") + with pytest.raises(ValueError, match="Missing required srt-slurm job metadata"): + load_results(tdef, tr) + status = tdef.was_run_successful(tr) + assert not status.is_successful + assert "123.json" in status.error_message + report = SRTSlurmReportGenerationStrategy(slurm_system, tr) + assert not report.can_handle_directory() + assert report.get_metric("tput_per_gpu") == METRIC_ERROR + + +@pytest.mark.parametrize( + "invalid", + [ + None, + [], + {"job_id": "999", "resources": {}}, + {"job_id": "123"}, + {"job_id": "123", "resources": {"agg_workers": None}}, + {"job_id": "123", "resources": {"gpus_per_agg": None}}, + {"job_id": "123", "resources": {"gpus_per_agg": True}}, + {"job_id": "123", "resources": {"agg_workers": True}}, + {"job_id": "123", "resources": {"agg_workers": -1}}, + {"job_id": "123", "resources": {"agg_workers": 1, "gpus_per_agg": 0}}, + {"job_id": "123", "resources": {"agg_workers": 1, "gpus_per_agg": 1.5}}, + {"job_id": "123", "resources": {"agg_workers": 1, "gpus_per_agg": "4"}}, + {"job_id": "123", "resources": {"agg_workers": 0, "prefill_workers": 0, "decode_workers": 0}}, + ], +) +def test_invalid_metadata_is_unsuccessful(slurm_system: SlurmSystem, tmp_path: Path, invalid): + tdef = make_tdef(tmp_path) + tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path) + tdef.job_id_path(tr).write_text("123") + path = write_job_metadata(tdef, tr) + if isinstance(invalid, dict) and isinstance(invalid.get("resources"), dict): + complete = json.loads(path.read_text()) + complete["resources"].update(invalid["resources"]) + invalid = {**invalid, "resources": complete["resources"]} + path.write_text(json.dumps(invalid)) + with pytest.raises(ValueError): + load_results(tdef, tr) + assert not tdef.was_run_successful(tr).is_successful + assert not SRTSlurmReportGenerationStrategy(slurm_system, tr).can_handle_directory() + + +def test_malformed_metadata_is_unsuccessful(tmp_path: Path): + tdef = make_tdef(tmp_path) + tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path) + tdef.job_id_path(tr).write_text("123") + path = write_job_metadata(tdef, tr) + path.write_text('{"job_id":') + status = tdef.was_run_successful(tr) + assert not status.is_successful + assert "job_id" in status.error_message + + +@pytest.mark.parametrize("workers, per_worker", [(1, 1), (2, 2), (1, 16)]) +def test_dse_reporting_uses_each_jobs_metadata(tmp_path: Path, workers: int, per_worker: int): + tdef = make_tdef(tmp_path) + base = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path) + tr = base.apply_params_set( + {"recipe_overrides.roles.agg.workers": workers, "recipe_overrides.roles.agg.gpus": per_worker} + ) + assert isinstance(tr.test, SRTSlurmTestDefinition) + tr.output_path = tmp_path / "selected" + tr.output_path.mkdir() + tr.test.job_id_path(tr).write_text("123") + write_job_metadata(tr.test, tr, agg_workers=workers, gpus_per_agg=per_worker) + logs = tr.test.logs_path(tr) + assert logs is not None + raw = logs / "sa-bench_one" / "results_concurrency_1_gpus_999.json" + raw.parent.mkdir(parents=True) + raw.write_text(json.dumps({"total_token_throughput": 1600, "output_throughput": 800, "max_concurrency": 1})) + result = load_results(tr.test, tr)[0] + assert result["serving_gpus"] == workers * per_worker + assert result["tput_per_gpu"] == 1600 / (workers * per_worker) + assert tdef.cmd_args.recipe_overrides == {"roles": {"agg": {"gpus": 4}}} + + +def test_agentx_disaggregated_metrics_use_total_pool(): + data = { + "conc": 1, + "request_metrics": { + "throughput": { + "total": {"tokens_per_second": 1000}, + "input": {"tokens_per_second": 800}, + "output": {"tokens_per_second": 200}, + "per_gpu": {"total_tput_tps": 999, "input_tput_tps": 999, "output_tput_tps": 999}, + } + }, + } + metrics = agentx_metrics(data, JobGPUCounts(aggregate=0, prefill=8, decode=4)) + assert metrics["tput_per_gpu"] == 1000 / 12 + assert metrics["input_tput_per_gpu"] == 800 / 12 + assert metrics["output_tput_per_gpu"] == 200 / 12 + + +def test_agentx_without_raw_total_is_unusable(): + data = {"conc": 1, "request_metrics": {"throughput": {"per_gpu": {"total_tput_tps": 100}}}} + assert not agentx_metrics(data, JobGPUCounts(aggregate=4, prefill=0, decode=0)) + + +def test_rollup_normalization_is_replaced_and_missing_total_stays_unavailable(): + gpus = JobGPUCounts(aggregate=2, prefill=0, decode=0) + data = { + "throughput_toks": 200, + "total_token_throughput": 1000, + "tput_per_gpu": 125, + "input_tput_per_gpu": 100, + "output_tput_per_gpu": 25, + } + metrics = sa_bench_metrics(data, gpus) + assert metrics["tput_per_gpu"] == 500 + assert metrics["input_tput_per_gpu"] == 400 + assert metrics["output_tput_per_gpu"] == 100 + data.pop("total_token_throughput") + assert "tput_per_gpu" not in sa_bench_metrics(data, gpus) diff --git a/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py b/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py index 1b9454546..5e9f1a906 100644 --- a/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py +++ b/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py @@ -19,8 +19,12 @@ from pathlib import Path import pytest +import toml +import yaml from cloudai.core import Registry, TestRun, TestScenario +from cloudai.models.scenario import ReportConfig, TestRunDetails +from cloudai.reporter import DSEReporter from cloudai.systems.slurm import SlurmRunner, SlurmSystem from cloudai.workloads.srt_slurm import SRTSlurmCommandGenStrategy @@ -35,19 +39,49 @@ def test_runner_tracks_the_recipe_job(slurm_system: SlurmSystem, tmp_path: Path, tdef.srtctl_path.parent.mkdir(parents=True) repo = slurm_system.install_path / tdef.git_repos[0].repo_name repo.mkdir(parents=True) + tdef.git_repos[0].installed_path = repo tdef.srtctl_path.write_text(f"#!{sys.executable}\nprint('Submitted batch job 111')\nprint('Job 222 submitted')\n") tdef.srtctl_path.chmod(0o755) monkeypatch.chdir(tmp_path) tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=Path("relative output")) runner = SlurmRunner("run", slurm_system, TestScenario(name="srt", test_runs=[tr]), tmp_path) - job = runner._submit_test(tr) + runner.submit_test(tr) + job = runner.jobs[0] assert job.id == 222 assert tdef.job_id_path(tr).read_text().strip() == "222" - runner.on_job_submit(tr) assert (tr.output_path / "test-run.toml").is_file() + details = TestRunDetails.model_validate(toml.load(tr.output_path / "test-run.toml")) + assert details.test_definition.cmd_args.recipe_path == tdef.cmd_args.recipe_path + assert details.test_definition.cmd_args.cluster_config == tdef.cmd_args.cluster_config + assert details.test_definition.git_repos[0].installed_path == repo assert Registry().get_command_gen_strategy(SlurmSystem, type(tdef)) is SRTSlurmCommandGenStrategy +def test_best_dse_config_preserves_paths(slurm_system: SlurmSystem, tmp_path: Path): + tdef = make_tdef(tmp_path) + tdef.srtctl.venv_path = tmp_path / "venv" + tdef.cmd_args.recipe_overrides = {"roles": {"agg": {"args": {"tensor-parallel-size": [1, 2]}}}} + tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path / "results") + root = tmp_path / "results" + iteration = root / "srt" / "0" + iteration.mkdir(parents=True) + key = "recipe_overrides.roles.agg.args.tensor-parallel-size" + (iteration / "trajectory.csv").write_text( + f"step,action,reward,observation\n1,\"{{'{key}': 1}}\",10,[10]\n2,\"{{'{key}': 2}}\",20,[20]\n" + ) + for step in (1, 2): + selected = tr.apply_params_set({key: step}) + selected.output_path = iteration / str(step) + SRTSlurmCommandGenStrategy(slurm_system, selected).store_test_run() + reporter = DSEReporter(slurm_system, TestScenario(name="srt", test_runs=[tr]), root, ReportConfig()) + reporter.report_best_dse_config() + exported = toml.load(iteration / "srt.toml") + reloaded = type(tdef).model_validate(exported) + assert reloaded.cmd_args.recipe_path == tdef.cmd_args.recipe_path + assert reloaded.cmd_args.cluster_config == tdef.cmd_args.cluster_config + assert reloaded.cmd_args.recipe_overrides["roles"]["agg"]["args"]["tensor-parallel-size"] == 2 + + def test_failed_submission_does_not_publish_a_job_id(slurm_system: SlurmSystem, tmp_path: Path): tdef = make_tdef(tmp_path) tdef.srtctl.venv_path = tmp_path / "venv" @@ -69,3 +103,21 @@ def test_single_sbatch_rejected(slurm_system: SlurmSystem, tmp_path: Path): tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path / "output") with pytest.raises(ValueError, match="single-sbatch"): SRTSlurmCommandGenStrategy(slurm_system, tr).gen_srun_command() + + +def test_rendered_overrides_support_old_cli(slurm_system: SlurmSystem, tmp_path: Path): + tdef = make_tdef(tmp_path) + original = tdef.cmd_args.recipe_path.read_text() + tdef.cmd_args.recipe_override_mode = "render" + tdef.srtctl.venv_path = tmp_path / "venv" + tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path / "results") + strategy = SRTSlurmCommandGenStrategy(slurm_system, tr) + strategy.gen_exec_command() + rendered = tr.output_path / "cloudai_generated" / "recipe.yaml" + assert yaml.safe_load(rendered.read_text())["roles"]["agg"]["gpus"] == 4 + assert tdef.cmd_args.recipe_path.read_text() == original + assert str(rendered) in strategy._paths()["run_script"].read_text() + assert "--set" not in strategy._paths()["run_script"].read_text() + tdef.cmd_args.recipe_overrides = {"roles": {"agg": {"typo": 1}}} + with pytest.raises(ValueError, match=r"absent key: roles\.agg\.typo"): + strategy.gen_exec_command() diff --git a/tests/workloads/srt_slurm/test_srt_slurm.py b/tests/workloads/srt_slurm/test_srt_slurm.py index 3a9656c98..b40e8f72f 100644 --- a/tests/workloads/srt_slurm/test_srt_slurm.py +++ b/tests/workloads/srt_slurm/test_srt_slurm.py @@ -77,6 +77,22 @@ def make_tdef(tmp_path: Path) -> SRTSlurmTestDefinition: ) +def write_job_metadata(tdef: SRTSlurmTestDefinition, tr: TestRun, **resources) -> Path: + path = tdef.metadata_path(tr) + assert path is not None + path.parent.mkdir(parents=True, exist_ok=True) + defaults = { + "gpus_per_node": 8, + "agg_nodes": 1, + "agg_workers": 1, + "gpus_per_agg": 4, + "prefill_workers": 0, + "decode_workers": 0, + } + path.write_text(json.dumps({"job_id": path.stem, "resources": defaults | resources})) + return path + + def test_command_generation_returns_srt_slurm_job(slurm_system: SlurmSystem, tmp_path: Path): tdef = make_tdef(tmp_path) tdef.srtctl.venv_path = tmp_path / "venv" @@ -212,7 +228,12 @@ def test_existing_recipe_is_loaded_and_dse_values_are_overlaid(tmp_path: Path): "benchmark.concurrencies=[32, 64]", ] assert source.read_text(encoding="utf-8") == original - assert tdef.dse_excluded_args == ["cmd_args.recipe_path", "cmd_args.cluster_config", "cmd_args.result_format"] + assert tdef.dse_excluded_args == [ + "cmd_args.recipe_path", + "cmd_args.cluster_config", + "cmd_args.result_format", + "cmd_args.recipe_override_mode", + ] assert "cmd_args.recipe_overrides" not in tdef.dse_excluded_args From e0d0f62ccbfad000078e219afb4a394e32856bde Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Wed, 14 Oct 2026 12:18:28 -0600 Subject: [PATCH 03/20] write generate srt-slurm artifacts into result directory without nested cloudai_generated/ --- doc/workloads/srt_slurm.rst | 5 ++++- src/cloudai/workloads/srt_slurm/installable.py | 4 ++-- .../workloads/srt_slurm/slurm_command_gen_strategy.py | 8 +++----- tests/test_acceptance.py | 4 +++- .../srt_slurm/test_slurm_command_gen_strategy.py | 4 +++- tests/workloads/srt_slurm/test_srt_slurm.py | 2 +- 6 files changed, 16 insertions(+), 11 deletions(-) diff --git a/doc/workloads/srt_slurm.rst b/doc/workloads/srt_slurm.rst index 7b47aa93a..21dc30b65 100644 --- a/doc/workloads/srt_slurm.rst +++ b/doc/workloads/srt_slurm.rst @@ -38,9 +38,12 @@ The first ``git_repos`` entry supplies srtctl; additional entries use ordinary CloudAI installation. Pin a revision supporting ``srtctl apply --set``. For older revisions without ``--set``, set ``cmd_args.recipe_override_mode = "render"``. This overlays existing mapping keys in a concrete recipe, preserves YAML value types, -and submits a trial-local ``cloudai_generated/recipe.yaml``. It rejects missing keys +and submits a trial-local ``srt_slurm_recipe.yaml``. It rejects missing keys and does not resolve recipe variants or interpret native override path expressions; select the effective concrete recipe first. The default remains ``"native"``. +The launch script ``run_srtctl.sh`` and rendered recipe are stored directly in +each trial's output directory, alongside ``run_srtctl.log`` and ``srt_job_id.txt``. +SRT Slurm job outputs are stored under ``srt_slurm_outputs/`` in that directory. Lists in ``recipe_overrides`` are DSE dimensions. Use nested candidate lists when the selected recipe value is itself a list. Recipe paths, cluster config paths, and result format are excluded from DSE. Additional apply flags use diff --git a/src/cloudai/workloads/srt_slurm/installable.py b/src/cloudai/workloads/srt_slurm/installable.py index f39391943..7fee779bf 100644 --- a/src/cloudai/workloads/srt_slurm/installable.py +++ b/src/cloudai/workloads/srt_slurm/installable.py @@ -59,8 +59,8 @@ def is_installed(self, installer: BaseInstaller) -> InstallStatusResult: return self._runtime_status(installer) def _prepare_cluster_config(self, installer: BaseInstaller, repo_path: Path) -> None: - # make setup prompts interactively without this file. This bootstrap - # config is checkout-local; trial submission always uses the supplied config. + # We create a placeholder config (srtslurm.yaml) file to avoid the interactive prompt during `make setup`. + # This is purely a placeholder - the actual config file used is specified via Tests.cmd_args.cluster_config in the CloudAI config config = repo_path / "srtslurm.yaml" if not config.is_file(): self._write_cluster_config(installer, config, repo_path) diff --git a/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py b/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py index 61eb718f5..42e29fb23 100644 --- a/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py +++ b/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py @@ -61,19 +61,17 @@ def tdef(self) -> SRTSlurmTestDefinition: def _paths(self) -> dict[str, Path]: output = self.test_run.output_path.absolute() - generated = output / "cloudai_generated" return { - "generated": generated, "recipe": self.tdef.cmd_args.recipe_path.expanduser().absolute(), + "rendered_recipe": output / "srt_slurm_recipe.yaml", "nested_outputs": output / "srt_slurm_outputs", "run_log": output / "run_srtctl.log", - "run_script": generated / "run_srtctl.sh", + "run_script": output / "run_srtctl.sh", "job_id": self.tdef.job_id_path(self.test_run).absolute(), } def write_run_script(self) -> None: paths = self._paths() - paths["generated"].mkdir(parents=True, exist_ok=True) repo = self.tdef.git_repos[0] repo_path = (repo.installed_path or self.system.install_path / repo.repo_name).absolute() cluster_config = self.tdef.cmd_args.cluster_config.expanduser().absolute() @@ -82,7 +80,7 @@ def write_run_script(self) -> None: if self.tdef.cmd_args.recipe_override_mode == "render": data = self._load_recipe(recipe) self._overlay_recipe(data, self.tdef.cmd_args.recipe_overrides) - recipe = paths["generated"] / "recipe.yaml" + recipe = paths["rendered_recipe"] recipe.write_text(yaml.safe_dump(data, sort_keys=False), encoding="utf-8") overrides = [] apply_command = shlex.join( diff --git a/tests/test_acceptance.py b/tests/test_acceptance.py index 4a2a2bea3..6d8bec878 100644 --- a/tests/test_acceptance.py +++ b/tests/test_acceptance.py @@ -967,7 +967,9 @@ def test_srt_slurm_example_dry_run(slurm_system: SlurmSystem, tmp_path: Path): job = runner._submit_test(tr) runner.on_job_submit(tr) assert job.id == 0 - script = tr.output_path / "cloudai_generated" / "run_srtctl.sh" + script = tr.output_path / "run_srtctl.sh" assert script.is_file() + assert not (tr.output_path / "cloudai_generated").exists() + assert not (tr.output_path / "srt_slurm_recipe.yaml").exists() assert "--set" in script.read_text() assert (tr.output_path / "test-run.toml").is_file() diff --git a/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py b/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py index 5e9f1a906..df7a4205e 100644 --- a/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py +++ b/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py @@ -113,7 +113,9 @@ def test_rendered_overrides_support_old_cli(slurm_system: SlurmSystem, tmp_path: tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path / "results") strategy = SRTSlurmCommandGenStrategy(slurm_system, tr) strategy.gen_exec_command() - rendered = tr.output_path / "cloudai_generated" / "recipe.yaml" + rendered = tr.output_path / "srt_slurm_recipe.yaml" + assert (tr.output_path / "run_srtctl.sh").is_file() + assert not (tr.output_path / "cloudai_generated").exists() assert yaml.safe_load(rendered.read_text())["roles"]["agg"]["gpus"] == 4 assert tdef.cmd_args.recipe_path.read_text() == original assert str(rendered) in strategy._paths()["run_script"].read_text() diff --git a/tests/workloads/srt_slurm/test_srt_slurm.py b/tests/workloads/srt_slurm/test_srt_slurm.py index b40e8f72f..aaa12911f 100644 --- a/tests/workloads/srt_slurm/test_srt_slurm.py +++ b/tests/workloads/srt_slurm/test_srt_slurm.py @@ -100,7 +100,7 @@ def test_command_generation_returns_srt_slurm_job(slurm_system: SlurmSystem, tmp strategy = SRTSlurmCommandGenStrategy(slurm_system, test_run) command = strategy.gen_exec_command() - run_script = test_run.output_path / "cloudai_generated" / "run_srtctl.sh" + run_script = test_run.output_path / "run_srtctl.sh" script = run_script.read_text(encoding="utf-8") subprocess.run(["bash", "-n", run_script], check=True) From bf983f44f597604565f285dbc066c3238fdf98c7 Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Tue, 27 Oct 2026 14:19:36 -0600 Subject: [PATCH 04/20] harden srt-slurm submission wrapper --- .../srt_slurm/slurm_command_gen_strategy.py | 46 ++++++++++++++----- src/cloudai/workloads/srt_slurm/srt_slurm.py | 2 + tests/test_init.py | 5 +- tests/test_test_scenario.py | 4 +- .../test_slurm_command_gen_strategy.py | 30 ++++++++++++ tests/workloads/srt_slurm/test_srt_slurm.py | 17 ++++++- 6 files changed, 89 insertions(+), 15 deletions(-) diff --git a/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py b/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py index 42e29fb23..71bfe7b82 100644 --- a/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py +++ b/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py @@ -70,19 +70,26 @@ def _paths(self) -> dict[str, Path]: "job_id": self.tdef.job_id_path(self.test_run).absolute(), } - def write_run_script(self) -> None: - paths = self._paths() - repo = self.tdef.git_repos[0] - repo_path = (repo.installed_path or self.system.install_path / repo.repo_name).absolute() - cluster_config = self.tdef.cmd_args.cluster_config.expanduser().absolute() + def _prepare_recipe(self, paths: dict[str, Path]) -> tuple[Path, list[str]]: + """Select the original recipe with CLI overrides or write a rendered copy.""" recipe = paths["recipe"] overrides = self._override_args(self.tdef.cmd_args.recipe_overrides) + # "render" mode is for older CLIs that don't support passing `--set` to srtctl. + # Apply the overrides to a generated recipe instead of passing them on the CLI. if self.tdef.cmd_args.recipe_override_mode == "render": data = self._load_recipe(recipe) self._overlay_recipe(data, self.tdef.cmd_args.recipe_overrides) recipe = paths["rendered_recipe"] recipe.write_text(yaml.safe_dump(data, sort_keys=False), encoding="utf-8") overrides = [] + return recipe, overrides + + def write_run_script(self) -> None: + paths = self._paths() + repo = self.tdef.git_repos[0] + repo_path = (repo.installed_path or self.system.install_path / repo.repo_name).absolute() + cluster_config = self.tdef.cmd_args.cluster_config.expanduser().absolute() + recipe, overrides = self._prepare_recipe(paths) apply_command = shlex.join( [ str(self.tdef.srtctl_path), @@ -100,11 +107,15 @@ def write_run_script(self) -> None: [ "#!/usr/bin/env bash", "set -euo pipefail", + "", + "# Configure submission environment", *[f"export {key}={shlex.quote(str(value))}" for key, value in self.final_env_vars.items()], f"export SRTSLURM_CONFIG={shlex.quote(str(cluster_config))}", *(["export AGENTIC_OUTPUT_DIR=/logs/agentic"] if self.tdef.cmd_args.result_format == "agentx" else []), - # uv also discovers Python through PATH after VIRTUAL_ENV is - # unset. Strip login-node venvs before Slurm exports the PATH. + "", + "# Prevent login-node Python environments from leaking into compute jobs.", + "# A: uv searches PATH after VIRTUAL_ENV. Remove every virtualenv-backed", + "# PATH entry so it cannot select an incompatible login-node python/python3.", 'IFS=: read -r -a SRT_PATH_PARTS <<< "$PATH"', "SRT_CLEAN_PATH=()", 'for SRT_PATH_PART in "${SRT_PATH_PARTS[@]}"; do', @@ -113,20 +124,31 @@ def write_run_script(self) -> None: " fi", "done", 'export PATH="$(IFS=:; printf \'%s\' "${SRT_CLEAN_PATH[*]}")"', + "# B: uv prioritizes the active environment. Unset VIRTUAL_ENV so it cannot", + "# inspect or execute CloudAI's login-node interpreter on a compute node.", "unset VIRTUAL_ENV", + "", + "# Submit the recipe", f"mkdir -p {shlex.quote(str(paths['nested_outputs']))}", f"rm -f {shlex.quote(str(paths['job_id']))}", - # Slurm inherits this directory. The early Dynamo prefetch must - # not discover the login-node CloudAI .venv in the caller's cwd. + "# C: uv discovers .venv from the current directory even with --no-project.", + "# Submit from srt-slurm so it cannot discover CloudAI's checkout .venv.", f"cd {shlex.quote(str(repo_path))}", f"if ! {apply_command} > {shlex.quote(str(paths['run_log']))} 2>&1; then", f" cat {shlex.quote(str(paths['run_log']))} >&2", " exit 1", "fi", f"cat {shlex.quote(str(paths['run_log']))} >&2", - "SRT_JOB_ID=\"$(sed -nE 's/.*Job ([0-9]+) submitted.*/\\1/p; s/.*Submitted batch job ([0-9]+).*/\\1/p' " - f'{shlex.quote(str(paths["run_log"]))} | tail -1)"', - '[[ "$SRT_JOB_ID" =~ ^[0-9]+$ ]] || { echo "Unable to find submitted srt-slurm job ID"; exit 1; }', + "", + "# Record the submitted job ID for CloudAI", + "if ! SRT_JOB_ID=\"$(sed -nE 's/.*Job ([0-9]+) submitted.*/\\1/p; " + "s/.*Submitted batch job ([0-9]+).*/\\1/p' " + f'{shlex.quote(str(paths["run_log"]))} | tail -1)" || [[ ! "$SRT_JOB_ID" =~ ^[0-9]+$ ]]; then', + ' echo "Unable to find submitted srt-slurm job ID" >&2', + " exit 1", + "fi", + "# Persist the result lookup ID now; slurm-job.toml is written at completion", + "# only when accounting is available. Stdout separately supplies the runner ID.", f'printf "%s\\n" "$SRT_JOB_ID" > {shlex.quote(str(paths["job_id"]))}', 'echo "Submitted batch job $SRT_JOB_ID"', "", diff --git a/src/cloudai/workloads/srt_slurm/srt_slurm.py b/src/cloudai/workloads/srt_slurm/srt_slurm.py index 0cdb3d27c..023b622c6 100644 --- a/src/cloudai/workloads/srt_slurm/srt_slurm.py +++ b/src/cloudai/workloads/srt_slurm/srt_slurm.py @@ -77,10 +77,12 @@ def installables(self) -> list[Installable]: @staticmethod def job_id_path(tr: TestRun) -> Path: + """Locate the saved job ID used for result lookup and later report regeneration.""" return tr.output_path / "srt_job_id.txt" def job_output_path(self, tr: TestRun) -> Path | None: """Return the submitted job's output directory, without guessing its identity.""" + # TestRun has no runtime SlurmJob ID; avoid guessing among old submission directories. try: job_id = self.job_id_path(tr).read_text(encoding="utf-8").strip() except OSError: diff --git a/tests/test_init.py b/tests/test_init.py index 868a8144f..40f14f7e6 100644 --- a/tests/test_init.py +++ b/tests/test_init.py @@ -108,6 +108,7 @@ SleepTestDefinition, ) from cloudai.workloads.slurm_container import SlurmContainerCommandGenStrategy, SlurmContainerTestDefinition +from cloudai.workloads.srt_slurm import SRTSlurmCommandGenStrategy, SRTSlurmTestDefinition from cloudai.workloads.triton_inference import ( TritonInferenceSlurmCommandGenStrategy, TritonInferenceTestDefinition, @@ -171,6 +172,7 @@ def test_runners(): (SlurmSystem, OSUBenchTestDefinition): OSUBenchSlurmCommandGenStrategy, (SlurmSystem, FioTestDefinition): FioSlurmCommandGenStrategy, (SlurmSystem, SglangTestDefinition): SglangSlurmCommandGenStrategy, + (SlurmSystem, SRTSlurmTestDefinition): SRTSlurmCommandGenStrategy, (SlurmSystem, VllmTestDefinition): VllmSlurmCommandGenStrategy, } JSON_GEN_STRATEGIES = { @@ -245,7 +247,7 @@ def test_installers(): def test_definitions(): test_defs = Registry().test_definitions_map - assert len(test_defs) == 28 + assert len(test_defs) == 29 for tdef in [ ("UCCTest", UCCTestDefinition), ("DDLBTest", DDLBTestDefinition), @@ -275,6 +277,7 @@ def test_definitions(): ("sglang", SglangTestDefinition), ("vllm", VllmTestDefinition), ("DynamoMocker", DynamoMockerTestDefinition), + ("SRTSlurm", SRTSlurmTestDefinition), ]: assert test_defs[tdef[0]] == tdef[1] diff --git a/tests/test_test_scenario.py b/tests/test_test_scenario.py index 1e2d92158..e51db7b3b 100644 --- a/tests/test_test_scenario.py +++ b/tests/test_test_scenario.py @@ -76,6 +76,7 @@ from cloudai.workloads.nixl_perftest import NIXLKVBenchDummyReport, NixlPerftestTestDefinition from cloudai.workloads.osu_bench import OSUBenchReportGenerationStrategy, OSUBenchTestDefinition from cloudai.workloads.sglang import SGLangBenchReportGenerationStrategy, SglangTestDefinition +from cloudai.workloads.srt_slurm import SRTSlurmReportGenerationStrategy, SRTSlurmTestDefinition from cloudai.workloads.triton_inference import TritonInferenceReportGenerationStrategy, TritonInferenceTestDefinition from cloudai.workloads.ucc_test import UCCTestDefinition, UCCTestReportGenerationStrategy @@ -757,7 +758,7 @@ def test_default(self): assert len(reporters) == 0 def test_default_reporters_size(self): - assert len(Registry().reports_map) == 22 + assert len(Registry().reports_map) == 23 @pytest.mark.parametrize( "tdef,expected_reporters", @@ -798,6 +799,7 @@ def test_default_reporters_size(self): (FioTestDefinition, {FioReportGenerationStrategy}), (NixlPerftestTestDefinition, {NIXLKVBenchDummyReport}), (AiconfiguratorTestDefinition, {AiconfiguratorReportGenerationStrategy}), + (SRTSlurmTestDefinition, {SRTSlurmReportGenerationStrategy}), ], ) def test_custom_reporters(self, tdef: Type[TestDefinition], expected_reporters: Set[ReportGenerationStrategy]): diff --git a/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py b/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py index df7a4205e..fd6736575 100644 --- a/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py +++ b/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py @@ -14,6 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os import subprocess import sys from pathlib import Path @@ -105,6 +106,35 @@ def test_single_sbatch_rejected(slurm_system: SlurmSystem, tmp_path: Path): SRTSlurmCommandGenStrategy(slurm_system, tr).gen_srun_command() +@pytest.mark.parametrize("extraction_fails", [False, True]) +def test_job_id_errors_are_reported_to_stderr(slurm_system: SlurmSystem, tmp_path: Path, monkeypatch, extraction_fails): + tdef = make_tdef(tmp_path) + tdef.srtctl.venv_path = tmp_path / "venv" + tdef.srtctl_path.parent.mkdir(parents=True) + repo = slurm_system.install_path / tdef.git_repos[0].repo_name + repo.mkdir(parents=True) + response = "Submitted batch job 111" if extraction_fails else "Submission completed without a job ID" + tdef.srtctl_path.write_text(f"#!{sys.executable}\nprint({response!r})\n") + tdef.srtctl_path.chmod(0o755) + if extraction_fails: + tools = tmp_path / "tools" + tools.mkdir() + sed = tools / "sed" + sed.write_text("#!/usr/bin/env bash\nexit 1\n") + sed.chmod(0o755) + monkeypatch.setenv("PATH", f"{tools}:{os.environ['PATH']}") + tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path / "output") + strategy = SRTSlurmCommandGenStrategy(slurm_system, tr) + strategy.gen_exec_command() + tdef.job_id_path(tr).write_text("999\n") + completed = subprocess.run(["bash", strategy._paths()["run_script"]], capture_output=True, text=True) + assert completed.returncode != 0 + assert completed.stdout == "" + assert response in completed.stderr + assert "Unable to find submitted srt-slurm job ID" in completed.stderr + assert not tdef.job_id_path(tr).exists() + + def test_rendered_overrides_support_old_cli(slurm_system: SlurmSystem, tmp_path: Path): tdef = make_tdef(tmp_path) original = tdef.cmd_args.recipe_path.read_text() diff --git a/tests/workloads/srt_slurm/test_srt_slurm.py b/tests/workloads/srt_slurm/test_srt_slurm.py index aaa12911f..e7deb1efb 100644 --- a/tests/workloads/srt_slurm/test_srt_slurm.py +++ b/tests/workloads/srt_slurm/test_srt_slurm.py @@ -131,7 +131,21 @@ def test_run_script_exports_merged_environment(slurm_system: SlurmSystem, tmp_pa (active_bin.parent / "pyvenv.cfg").write_text("home = /login-node/python\n") retained_bin = tmp_path / "cluster tools" / "bin" retained_bin.mkdir(parents=True) - monkeypatch.setenv("PATH", f"{caller_path / '.venv' / 'bin'}:{active_bin}:{retained_bin}:{os.environ['PATH']}") + other_venv_bin = tmp_path / "other login-node venv" / "bin" + other_venv_bin.mkdir(parents=True) + (other_venv_bin.parent / "pyvenv.cfg").write_text("home = /other/login-node/python\n") + monkeypatch.setenv( + "PATH", + os.pathsep.join( + [ + str(caller_path / ".venv" / "bin"), + str(active_bin), + str(other_venv_bin), + str(retained_bin), + os.environ["PATH"], + ] + ), + ) tdef.extra_env_vars = { "SRTSLURM_CONFIG": str(tmp_path / "cluster's config" / "srtslurm.yaml"), "SRT_LITERAL": "$HOME; $(printf injected)", @@ -171,6 +185,7 @@ def test_run_script_exports_merged_environment(slurm_system: SlurmSystem, tmp_pa submitted_path = strategy._paths()["run_log"].read_text().splitlines()[2].split(os.pathsep) assert str(caller_path / ".venv" / "bin") not in submitted_path assert str(active_bin) not in submitted_path + assert str(other_venv_bin) not in submitted_path assert str(retained_bin) in submitted_path submitted_args = json.loads(strategy._paths()["run_log"].read_text().splitlines()[3]) assert submitted_args[submitted_args.index("-f") + 1] == str(tdef.cmd_args.recipe_path) From 040699718cd8080fa7abd28533846cc8f08a28b5 Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Tue, 27 Oct 2026 14:19:46 -0600 Subject: [PATCH 05/20] clarify srt-slurm scheduling ownership --- conf/experimental/srt_slurm/scenario.toml | 2 + conf/experimental/srt_slurm/srtslurm.yaml | 2 + doc/workloads/srt_slurm.rst | 53 ++++++++++++++++++++--- 3 files changed, 51 insertions(+), 6 deletions(-) diff --git a/conf/experimental/srt_slurm/scenario.toml b/conf/experimental/srt_slurm/scenario.toml index 05dd9ce42..abb3e31b6 100644 --- a/conf/experimental/srt_slurm/scenario.toml +++ b/conf/experimental/srt_slurm/scenario.toml @@ -18,6 +18,7 @@ name = "srt_slurm" [[Tests]] id = "srt.sglang" +# Benchmark node allocation is defined by recipe.yaml, not this value. num_nodes = 1 name = "srt_sglang" @@ -32,6 +33,7 @@ agent_metrics = ["tput_per_gpu"] [Tests.cmd_args] recipe_path = "conf/experimental/srt_slurm/recipe.yaml" + # Owns benchmark scheduling defaults; CloudAI allocation settings are not merged. cluster_config = "conf/experimental/srt_slurm/srtslurm.yaml" result_format = "sa-bench" diff --git a/conf/experimental/srt_slurm/srtslurm.yaml b/conf/experimental/srt_slurm/srtslurm.yaml index 6441472a8..3768b2736 100644 --- a/conf/experimental/srt_slurm/srtslurm.yaml +++ b/conf/experimental/srt_slurm/srtslurm.yaml @@ -1,3 +1,5 @@ +# Benchmark scheduling defaults belong here and in the recipe. +# CloudAI system/scenario allocation settings are not forwarded to srt-slurm. # Adapt these paths and Slurm defaults to your shared filesystem. default_partition: gpu gpus_per_node: 8 diff --git a/doc/workloads/srt_slurm.rst b/doc/workloads/srt_slurm.rst index 21dc30b65..f4673b904 100644 --- a/doc/workloads/srt_slurm.rst +++ b/doc/workloads/srt_slurm.rst @@ -2,7 +2,9 @@ SRT-Slurm ========= The ``SRTSlurm`` workload submits an existing recipe with ``srtctl apply`` and -monitors the resulting Slurm job directly. The pinned srtctl version validates +monitors the resulting Slurm job directly. CloudAI is a pass-through layer for +benchmark submission: srt-slurm owns scheduling and resource orchestration. +The pinned srtctl version validates the recipe, including recipes without a ``schema`` field. Each CloudAI trial passes selected ``recipe_overrides`` as typed ``--set KEY=VALUE`` arguments; the original YAML remains unchanged. @@ -52,8 +54,13 @@ paths, and result format are excluded from DSE. Additional apply flags use ``cluster_config`` is required and takes precedence over ``SRTSLURM_CONFIG`` in the submitting environment. Supply a complete config with the recipe's aliases, mounts, and cluster settings. CloudAI exports merged system and test -environment variables, removes login-node virtual environments from ``PATH``, -unsets ``VIRTUAL_ENV``, and submits from the installed srt-slurm checkout. +environment variables, removes virtual-environment directories from ``PATH``, +unsets ``VIRTUAL_ENV``, and submits from the installed srt-slurm checkout. These +steps prevent compute-node ``uv`` from selecting login-node Python through its +PATH search, active-environment lookup, or current-directory ``.venv`` lookup. +Tools supplied through virtual-environment PATH entries are removed as part of +this compatibility cleanup; supply required compute-node tools outside a Python +virtual environment. Installation runs ``make setup`` in a single-node, 15-minute ``srun`` using the CloudAI system's partition, account, and extra srun arguments. This downloads @@ -61,9 +68,43 @@ runtime tools for the compute-node architecture. A checkout-local bootstrap config prevents interactive setup prompts; existing checkout configs are preserved. Trial submission always uses the explicit ``cluster_config``. -The recipe and cluster config own job resources and orchestration. Configure -nodes, time limits, container mounts, profiling, and setup/cleanup commands -there. CloudAI's scenario allocation directives do not rewrite the recipe. +Scheduling ownership +~~~~~~~~~~~~~~~~~~~~ + +Configure benchmark scheduling in the supplied ``cluster_config`` +(``srtslurm.yaml``) and recipe. CloudAI does not translate or merge its system +or scenario allocation settings into these files or into ``srtctl`` scheduling +arguments. Recipe overrides explicitly selected through ``recipe_overrides`` +are still passed to srt-slurm as described above. + +* ``srtslurm.yaml`` supplies cluster defaults such as account, partition, time + limit, GPU capacity, and supported SBATCH directives. +* The recipe specifies serving roles, nodes, workers, GPU requirements, + container mounts, profiling, and setup/cleanup commands, together with any + recipe-specific scheduling overrides. srt-slurm resolves these inputs. +* CloudAI's ``system.toml`` selects the Slurm runner and supplies installation, + cache, and output locations, shared environment variables, and monitoring + settings. Its partition, account, and extra srun arguments are used for the + installation setup allocation described above, not forwarded to benchmark + submission. + +For example, changing ``default_partition`` or ``account`` in CloudAI's system +config does not change the benchmark's partition or account. Likewise, +``extra_sbatch_args``, ``ntasks_per_node``, and scenario allocation settings +such as ``num_nodes``, ``nodes``, ``exclude_nodes``, and ``time_limit`` are not +translated into the benchmark allocation. Set the corresponding values in +srt-slurm's inputs. GPU capacity for benchmark scheduling comes from those +inputs, not CloudAI's ``gpus_per_node``. + +Keep ``scheduler = "slurm"`` in the CloudAI system config: the Slurm runner +tracks the job ID returned by srt-slurm, monitors completion, handles +cancellation, and collects accounting information. Using this runner does not +make CloudAI the owner of the benchmark allocation. + +When changing clusters, select both the target CloudAI system config and its +matching ``cluster_config`` in the scenario, and check the recipe for explicit +cluster-specific paths or scheduling overrides. + Use normal per-trial Slurm CLI execution; single-sbatch, Slurm REST submission, and CloudAI pre/post-test hooks are unsupported for this workload. From f9389db07e3506d3bc0a717042a0bde8bf5b1cf8 Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Tue, 27 Oct 2026 14:21:31 -0600 Subject: [PATCH 06/20] fix srt-slurm installer lint --- src/cloudai/workloads/srt_slurm/installable.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/cloudai/workloads/srt_slurm/installable.py b/src/cloudai/workloads/srt_slurm/installable.py index 7fee779bf..755b49bda 100644 --- a/src/cloudai/workloads/srt_slurm/installable.py +++ b/src/cloudai/workloads/srt_slurm/installable.py @@ -59,8 +59,8 @@ def is_installed(self, installer: BaseInstaller) -> InstallStatusResult: return self._runtime_status(installer) def _prepare_cluster_config(self, installer: BaseInstaller, repo_path: Path) -> None: - # We create a placeholder config (srtslurm.yaml) file to avoid the interactive prompt during `make setup`. - # This is purely a placeholder - the actual config file used is specified via Tests.cmd_args.cluster_config in the CloudAI config + # Create a placeholder srtslurm.yaml to avoid the interactive prompt during `make setup`. + # The actual config is supplied by Tests.cmd_args.cluster_config in the CloudAI scenario. config = repo_path / "srtslurm.yaml" if not config.is_file(): self._write_cluster_config(installer, config, repo_path) From ac675556279338f8ff983aca3754a3e2485e80ed Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Tue, 27 Oct 2026 14:27:06 -0600 Subject: [PATCH 07/20] fix srt-slurm copyright headers --- conf/experimental/srt_slurm/scenario.toml | 2 +- src/cloudai/workloads/srt_slurm/__init__.py | 2 +- src/cloudai/workloads/srt_slurm/installable.py | 2 +- src/cloudai/workloads/srt_slurm/report_generation_strategy.py | 2 +- src/cloudai/workloads/srt_slurm/results.py | 2 +- src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py | 2 +- src/cloudai/workloads/srt_slurm/srt_slurm.py | 2 +- tests/workloads/srt_slurm/__init__.py | 2 +- tests/workloads/srt_slurm/test_installable.py | 2 +- tests/workloads/srt_slurm/test_report_generation_strategy.py | 2 +- tests/workloads/srt_slurm/test_results.py | 2 +- tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py | 2 +- tests/workloads/srt_slurm/test_srt_slurm.py | 2 +- 13 files changed, 13 insertions(+), 13 deletions(-) diff --git a/conf/experimental/srt_slurm/scenario.toml b/conf/experimental/srt_slurm/scenario.toml index abb3e31b6..b0e6afe7c 100644 --- a/conf/experimental/srt_slurm/scenario.toml +++ b/conf/experimental/srt_slurm/scenario.toml @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/src/cloudai/workloads/srt_slurm/__init__.py b/src/cloudai/workloads/srt_slurm/__init__.py index 202cda6e9..1119c007a 100644 --- a/src/cloudai/workloads/srt_slurm/__init__.py +++ b/src/cloudai/workloads/srt_slurm/__init__.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/src/cloudai/workloads/srt_slurm/installable.py b/src/cloudai/workloads/srt_slurm/installable.py index 755b49bda..c77441259 100644 --- a/src/cloudai/workloads/srt_slurm/installable.py +++ b/src/cloudai/workloads/srt_slurm/installable.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/src/cloudai/workloads/srt_slurm/report_generation_strategy.py b/src/cloudai/workloads/srt_slurm/report_generation_strategy.py index 80309dc17..7c1f49eaf 100644 --- a/src/cloudai/workloads/srt_slurm/report_generation_strategy.py +++ b/src/cloudai/workloads/srt_slurm/report_generation_strategy.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/src/cloudai/workloads/srt_slurm/results.py b/src/cloudai/workloads/srt_slurm/results.py index dd4fa0b4f..1f9b6b9f7 100644 --- a/src/cloudai/workloads/srt_slurm/results.py +++ b/src/cloudai/workloads/srt_slurm/results.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py b/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py index 71bfe7b82..d2bf2ff1a 100644 --- a/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py +++ b/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/src/cloudai/workloads/srt_slurm/srt_slurm.py b/src/cloudai/workloads/srt_slurm/srt_slurm.py index 023b622c6..8145cacf0 100644 --- a/src/cloudai/workloads/srt_slurm/srt_slurm.py +++ b/src/cloudai/workloads/srt_slurm/srt_slurm.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/tests/workloads/srt_slurm/__init__.py b/tests/workloads/srt_slurm/__init__.py index 19b239dd6..3ac11851d 100644 --- a/tests/workloads/srt_slurm/__init__.py +++ b/tests/workloads/srt_slurm/__init__.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/tests/workloads/srt_slurm/test_installable.py b/tests/workloads/srt_slurm/test_installable.py index 8bf4c0f8d..e36b2f9e3 100644 --- a/tests/workloads/srt_slurm/test_installable.py +++ b/tests/workloads/srt_slurm/test_installable.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/tests/workloads/srt_slurm/test_report_generation_strategy.py b/tests/workloads/srt_slurm/test_report_generation_strategy.py index 85e2f4ed2..095512a8f 100644 --- a/tests/workloads/srt_slurm/test_report_generation_strategy.py +++ b/tests/workloads/srt_slurm/test_report_generation_strategy.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/tests/workloads/srt_slurm/test_results.py b/tests/workloads/srt_slurm/test_results.py index 65eeb4ce6..87e97c33c 100644 --- a/tests/workloads/srt_slurm/test_results.py +++ b/tests/workloads/srt_slurm/test_results.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py b/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py index fd6736575..d155049f1 100644 --- a/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py +++ b/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/tests/workloads/srt_slurm/test_srt_slurm.py b/tests/workloads/srt_slurm/test_srt_slurm.py index e7deb1efb..e4b19fda0 100644 --- a/tests/workloads/srt_slurm/test_srt_slurm.py +++ b/tests/workloads/srt_slurm/test_srt_slurm.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); From 5634f2971796ecaffe6eb3f916326216e3b8973f Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Tue, 27 Oct 2026 14:44:55 -0600 Subject: [PATCH 08/20] address srt-slurm review feedback --- src/cloudai/workloads/srt_slurm/installable.py | 8 ++++++++ tests/test_test_scenario.py | 5 ++++- tests/workloads/srt_slurm/test_installable.py | 14 +++++++++++++- .../srt_slurm/test_slurm_command_gen_strategy.py | 3 ++- 4 files changed, 27 insertions(+), 3 deletions(-) diff --git a/src/cloudai/workloads/srt_slurm/installable.py b/src/cloudai/workloads/srt_slurm/installable.py index c77441259..ca8d5a6ce 100644 --- a/src/cloudai/workloads/srt_slurm/installable.py +++ b/src/cloudai/workloads/srt_slurm/installable.py @@ -29,6 +29,8 @@ class SRTSlurmExecutable(PythonExecutable): """Install srtctl and the binaries required on the compute nodes.""" + _SETUP_TIMEOUT_SECONDS = 30 * 60 + def _runtime_status(self, installer: BaseInstaller) -> InstallStatusResult: repo_path = self.git_repo.installed_path or installer.system.install_path / self.git_repo.repo_name binaries = ["configs/nats-server", "configs/etcd", "bin/uv"] @@ -99,6 +101,12 @@ def install(self, installer: BaseInstaller) -> InstallStatusResult: stdin=subprocess.DEVNULL, capture_output=True, text=True, + timeout=self._SETUP_TIMEOUT_SECONDS, + ) + except subprocess.TimeoutExpired: + return InstallStatusResult( + False, + f"srt-slurm setup timed out after {self._SETUP_TIMEOUT_SECONDS} seconds", ) except OSError as error: return InstallStatusResult(False, f"srt-slurm setup failed: {error}") diff --git a/tests/test_test_scenario.py b/tests/test_test_scenario.py index e51db7b3b..9b260c3eb 100644 --- a/tests/test_test_scenario.py +++ b/tests/test_test_scenario.py @@ -76,7 +76,10 @@ from cloudai.workloads.nixl_perftest import NIXLKVBenchDummyReport, NixlPerftestTestDefinition from cloudai.workloads.osu_bench import OSUBenchReportGenerationStrategy, OSUBenchTestDefinition from cloudai.workloads.sglang import SGLangBenchReportGenerationStrategy, SglangTestDefinition -from cloudai.workloads.srt_slurm import SRTSlurmReportGenerationStrategy, SRTSlurmTestDefinition +from cloudai.workloads.srt_slurm import ( + SRTSlurmReportGenerationStrategy, + SRTSlurmTestDefinition, +) from cloudai.workloads.triton_inference import TritonInferenceReportGenerationStrategy, TritonInferenceTestDefinition from cloudai.workloads.ucc_test import UCCTestDefinition, UCCTestReportGenerationStrategy diff --git a/tests/workloads/srt_slurm/test_installable.py b/tests/workloads/srt_slurm/test_installable.py index e36b2f9e3..2e3beb1ff 100644 --- a/tests/workloads/srt_slurm/test_installable.py +++ b/tests/workloads/srt_slurm/test_installable.py @@ -89,10 +89,22 @@ def test_installer_returns_setup_failure(runtime_installation): with patch( "cloudai.workloads.srt_slurm.installable.subprocess.run", return_value=subprocess.CompletedProcess(["make"], 2, "setup output", "download failed"), - ): + ) as run: result = executable.install(installer) assert not result.success assert "download failed" in result.message + assert run.call_args.kwargs["timeout"] == executable._SETUP_TIMEOUT_SECONDS + + +def test_installer_returns_setup_timeout(runtime_installation): + executable, installer, _ = runtime_installation + with patch( + "cloudai.workloads.srt_slurm.installable.subprocess.run", + side_effect=subprocess.TimeoutExpired(["srun", "make", "setup"], executable._SETUP_TIMEOUT_SECONDS), + ): + result = executable.install(installer) + assert not result.success + assert f"timed out after {executable._SETUP_TIMEOUT_SECONDS} seconds" in result.message def test_installable_identity_uses_revision(): diff --git a/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py b/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py index d155049f1..73a66c1db 100644 --- a/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py +++ b/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py @@ -93,7 +93,8 @@ def test_failed_submission_does_not_publish_a_job_id(slurm_system: SlurmSystem, tdef.srtctl_path.chmod(0o755) tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path / "output") strategy = SRTSlurmCommandGenStrategy(slurm_system, tr) - completed = subprocess.run(strategy.gen_exec_command(), shell=True, capture_output=True, text=True) + strategy.gen_exec_command() + completed = subprocess.run(["bash", strategy._paths()["run_script"]], capture_output=True, text=True) assert completed.returncode != 0 assert "Submitted batch job" not in completed.stdout assert not tdef.job_id_path(tr).exists() From ee99eff1cd89f8d454ffb68ce43d359eb975c288 Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Wed, 28 Oct 2026 06:39:39 -0600 Subject: [PATCH 09/20] remove hardcoded srt-slurm checkout path --- conf/experimental/srt_slurm/srtslurm.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/conf/experimental/srt_slurm/srtslurm.yaml b/conf/experimental/srt_slurm/srtslurm.yaml index 3768b2736..16927cf15 100644 --- a/conf/experimental/srt_slurm/srtslurm.yaml +++ b/conf/experimental/srt_slurm/srtslurm.yaml @@ -1,6 +1,6 @@ # Benchmark scheduling defaults belong here and in the recipe. # CloudAI system/scenario allocation settings are not forwarded to srt-slurm. -# Adapt these paths and Slurm defaults to your shared filesystem. +# Adapt these Slurm defaults to your cluster. srtctl_root is omitted so the +# pinned srtctl installation derives its own checkout path. default_partition: gpu gpus_per_node: 8 -srtctl_root: /shared/cloudai/install/srt-slurm__v2.2.1 From afa6bf161112dd55d1e435d3c9c119ea2bfc22d2 Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Thu, 29 Oct 2026 07:41:14 -0600 Subject: [PATCH 10/20] address comments --- conf/experimental/srt_slurm/recipe.yaml | 2 +- doc/workloads/srt_slurm.rst | 35 +++---- src/cloudai/report_generator/dse_report.py | 2 +- src/cloudai/reporter.py | 2 +- .../slurm/slurm_command_gen_strategy.py | 2 +- tests/ref_data/srt-slurm.sbatch | 43 +++++++++ tests/test_acceptance.py | 52 +++++----- tests/test_reporter.py | 4 +- tests/workloads/srt_slurm/__init__.py | 15 --- tests/workloads/srt_slurm/conftest.py | 95 +++++++++++++++++++ .../test_report_generation_strategy.py | 6 +- tests/workloads/srt_slurm/test_results.py | 24 +++-- .../test_slurm_command_gen_strategy.py | 49 ++-------- tests/workloads/srt_slurm/test_srt_slurm.py | 88 +---------------- 14 files changed, 220 insertions(+), 199 deletions(-) create mode 100644 tests/ref_data/srt-slurm.sbatch delete mode 100644 tests/workloads/srt_slurm/__init__.py create mode 100644 tests/workloads/srt_slurm/conftest.py diff --git a/conf/experimental/srt_slurm/recipe.yaml b/conf/experimental/srt_slurm/recipe.yaml index 32f2980f0..e79286e31 100644 --- a/conf/experimental/srt_slurm/recipe.yaml +++ b/conf/experimental/srt_slurm/recipe.yaml @@ -2,7 +2,7 @@ schema: 2 name: sglang-example model: path: Qwen/Qwen3-0.6B - container: docker.io/lmsysorg/sglang:latest + container: docker.io/lmsysorg/sglang:v0.5.9 resources: gpu_type: h100 gpus_per_node: 8 diff --git a/doc/workloads/srt_slurm.rst b/doc/workloads/srt_slurm.rst index f4673b904..42dacd0be 100644 --- a/doc/workloads/srt_slurm.rst +++ b/doc/workloads/srt_slurm.rst @@ -62,9 +62,11 @@ Tools supplied through virtual-environment PATH entries are removed as part of this compatibility cleanup; supply required compute-node tools outside a Python virtual environment. -Installation runs ``make setup`` in a single-node, 15-minute ``srun`` using the -CloudAI system's partition, account, and extra srun arguments. This downloads -runtime tools for the compute-node architecture. A checkout-local bootstrap +Installation runs ``make setup`` in a single-node, 15-minute ``srun`` using +``system.toml``'s ``default_partition``, ``account``, and ``extra_srun_args``. +These settings apply only to the installation allocation; benchmark scheduling +uses the supplied srt-slurm cluster config and recipe. Setup downloads runtime +tools for the compute-node architecture. A checkout-local bootstrap config prevents interactive setup prompts; existing checkout configs are preserved. Trial submission always uses the explicit ``cluster_config``. @@ -124,19 +126,20 @@ commit ``7858d3091d5e278ab56ebf0246df9ea913bc2477``, first released in ``v1.0.0` CloudAI calculates each role's GPU total as workers times GPUs per worker and sums those totals for the serving count. This automatically follows each trial's effective -DSE topology. ``serving_gpus``, ``prefill_gpus`` and ``decode_gpus`` are no longer -scenario fields; remove them from existing configs. Resolved counts appear in each -CSV result row as ``serving_gpus``, ``aggregate_gpus``, ``prefill_gpus`` and ``decode_gpus``. - -``default`` and ``output_throughput`` expose output tokens/s; -``total_token_throughput`` exposes input plus output tokens/s. -``tput_per_gpu`` divides total throughput by the metadata-derived serving count. -Disaggregated input/output per-GPU metrics use their respective prefill/decode pools; -aggregated metrics share the aggregate pool, counted once. Producer-normalized -rollup fields are ignored. A rollup lacking raw total throughput cannot provide -``tput_per_gpu``. Latency aliases such as ``median_ttft`` use seconds; -``ttft_median_ms`` uses milliseconds. ``median_intvty`` and ``p90_intvty`` -use the reciprocal of the corresponding TPOT, in tokens/s/user. +DSE topology. CSV result rows include ``serving_gpus``, ``aggregate_gpus``, +``prefill_gpus``, and ``decode_gpus``. + +* ``default`` and ``output_throughput`` expose output tokens/s. +* ``total_token_throughput`` exposes input plus output tokens/s. +* ``tput_per_gpu`` divides total throughput by the metadata-derived serving count. + Disaggregated input/output per-GPU metrics use their respective prefill/decode + pools; aggregated metrics share the aggregate pool, counted once. +* Producer-normalized rollup fields are ignored. A rollup lacking raw total + throughput cannot provide ``tput_per_gpu``. +* Latency aliases such as ``median_ttft`` use seconds; ``ttft_median_ms`` uses + milliseconds. +* ``median_intvty`` and ``p90_intvty`` use the reciprocal of the corresponding + TPOT, in tokens/s/user. ``result_format = "agentx"`` reads InferenceX aggregate JSON from the submitted job's ``logs/agentic`` directory. Each aggregate must have its corresponding diff --git a/src/cloudai/report_generator/dse_report.py b/src/cloudai/report_generator/dse_report.py index fa4cbed10..ddd0b8165 100644 --- a/src/cloudai/report_generator/dse_report.py +++ b/src/cloudai/report_generator/dse_report.py @@ -341,7 +341,7 @@ def _build_iteration_summary( observed_runtime=format_duration(total_observed_runtime_sec), efficiency_ratio=f"~{format_float(reduction_factor, 1)}x", efficiency_steps=f"{len(trajectory_steps):,} / {total_space:,} steps", - best_config_toml=toml.dumps(test_run_details.test_definition.model_dump(mode="json")), + best_config_toml=toml.dumps(test_run_details.test_definition.model_dump()), parameter_rows=_build_parameter_rows(test_case.param_space, best_step.action), reward_chart_data=_build_reward_chart_data(trajectory_steps), ) diff --git a/src/cloudai/reporter.py b/src/cloudai/reporter.py index 86d8e7dca..918fe26fb 100644 --- a/src/cloudai/reporter.py +++ b/src/cloudai/reporter.py @@ -293,7 +293,7 @@ def report_best_dse_config(self): best_config_path = tr_root / f"{tr.name}.toml" logging.info("Writing best config for %s to %s", tr.name, best_config_path) with best_config_path.open("w") as f: - toml.dump(trd.test_definition.model_dump(mode="json"), f) + toml.dump(trd.test_definition.model_dump(), f) class TarballReporter(Reporter): diff --git a/src/cloudai/systems/slurm/slurm_command_gen_strategy.py b/src/cloudai/systems/slurm/slurm_command_gen_strategy.py index 15413bdcc..52343357c 100644 --- a/src/cloudai/systems/slurm/slurm_command_gen_strategy.py +++ b/src/cloudai/systems/slurm/slurm_command_gen_strategy.py @@ -74,7 +74,7 @@ def store_test_run(self) -> None: test_cmd, srun_cmd = (" ".join(self.generate_test_command()), self.gen_srun_command()) with (self.test_run.output_path / self.TEST_RUN_DUMP_FILE_NAME).open("w") as f: trd = TestRunDetails.from_test_run(self.test_run, test_cmd=test_cmd, full_cmd=srun_cmd) - toml.dump(trd.model_dump(mode="json"), f) + toml.dump(trd.model_dump(), f) @final def container_mounts(self) -> list[str]: diff --git a/tests/ref_data/srt-slurm.sbatch b/tests/ref_data/srt-slurm.sbatch new file mode 100644 index 000000000..95ec4730a --- /dev/null +++ b/tests/ref_data/srt-slurm.sbatch @@ -0,0 +1,43 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Configure submission environment +export TEST_VAR='value with spaces' +export SRTSLURM_CONFIG=__CLOUDAI_DIR__/conf/experimental/srt_slurm/srtslurm.yaml + +# Prevent login-node Python environments from leaking into compute jobs. +# A: uv searches PATH after VIRTUAL_ENV. Remove every virtualenv-backed +# PATH entry so it cannot select an incompatible login-node python/python3. +IFS=: read -r -a SRT_PATH_PARTS <<< "$PATH" +SRT_CLEAN_PATH=() +for SRT_PATH_PART in "${SRT_PATH_PARTS[@]}"; do + if [[ ! -f "$SRT_PATH_PART/../pyvenv.cfg" ]]; then + SRT_CLEAN_PATH+=("$SRT_PATH_PART") + fi +done +export PATH="$(IFS=:; printf '%s' "${SRT_CLEAN_PATH[*]}")" +# B: uv prioritizes the active environment. Unset VIRTUAL_ENV so it cannot +# inspect or execute CloudAI's login-node interpreter on a compute node. +unset VIRTUAL_ENV + +# Submit the recipe +mkdir -p __OUTPUT_DIR__/output/srt_slurm_outputs +rm -f __OUTPUT_DIR__/output/srt_job_id.txt +# C: uv discovers .venv from the current directory even with --no-project. +# Submit from srt-slurm so it cannot discover CloudAI's checkout .venv. +cd __INSTALL_DIR__/srt-slurm__v2.2.1 +if ! __INSTALL_DIR__/srt-slurm__v2.2.1-venv/bin/srtctl apply -f __CLOUDAI_DIR__/conf/experimental/srt_slurm/recipe.yaml -o __OUTPUT_DIR__/output/srt_slurm_outputs -y --set 'benchmark.concurrencies="4"' --no-preflight > __OUTPUT_DIR__/output/run_srtctl.log 2>&1; then + cat __OUTPUT_DIR__/output/run_srtctl.log >&2 + exit 1 +fi +cat __OUTPUT_DIR__/output/run_srtctl.log >&2 + +# Record the submitted job ID for CloudAI +if ! SRT_JOB_ID="$(sed -nE 's/.*Job ([0-9]+) submitted.*/\1/p; s/.*Submitted batch job ([0-9]+).*/\1/p' __OUTPUT_DIR__/output/run_srtctl.log | tail -1)" || [[ ! "$SRT_JOB_ID" =~ ^[0-9]+$ ]]; then + echo "Unable to find submitted srt-slurm job ID" >&2 + exit 1 +fi +# Persist the result lookup ID now; slurm-job.toml is written at completion +# only when accounting is available. Stdout separately supplies the runner ID. +printf "%s\n" "$SRT_JOB_ID" > __OUTPUT_DIR__/output/srt_job_id.txt +echo "Submitted batch job $SRT_JOB_ID" diff --git a/tests/test_acceptance.py b/tests/test_acceptance.py index 6d8bec878..efd8b1572 100644 --- a/tests/test_acceptance.py +++ b/tests/test_acceptance.py @@ -29,7 +29,7 @@ import cloudai.cli.cli import cloudai.models.output from cloudai.cli import setup_logging -from cloudai.core import CommandGenStrategy, GitRepo, TestDefinition, TestRun, TestScenario +from cloudai.core import BaseInstaller, CommandGenStrategy, GitRepo, TestDefinition, TestRun, TestScenario from cloudai.models.scenario import TestRunDetails from cloudai.systems.slurm import SlurmCommandGenStrategy, SlurmRunner, SlurmSystem from cloudai.workloads.ai_dynamo import ( @@ -83,6 +83,7 @@ SlurmContainerCmdArgs, SlurmContainerTestDefinition, ) +from cloudai.workloads.srt_slurm import SRTSlurmCmdArgs, SRTSlurmTestDefinition from cloudai.workloads.triton_inference import ( TritonInferenceCmdArgs, TritonInferenceTestDefinition, @@ -239,6 +240,12 @@ def with_installed_dynamo_repo(test_definition: AIDynamoTestDefinition, installe return test_definition +def with_installed_srt_repo(test_definition: SRTSlurmTestDefinition, system: SlurmSystem) -> SRTSlurmTestDefinition: + """Set the installation paths used by acceptance command generation.""" + BaseInstaller(system).mark_as_installed(test_definition.installables) + return test_definition + + def build_special_test_run( partial_tr: partial[TestRun], param: str, test_mapping: Dict[str, Callable[[], TestRun]] ) -> Tuple[TestRun, str, Optional[str]]: @@ -312,6 +319,7 @@ def build_special_test_run( @pytest.fixture( params=[ + "srt-slurm", "ucc", "ddlb", "nccl", @@ -351,6 +359,26 @@ def build_special_test_run( ) def test_req(request, slurm_system: SlurmSystem, partial_tr: partial[TestRun]) -> Tuple[TestRun, str, Optional[str]]: test_mapping: Dict[str, Callable[[], TestRun]] = { + "srt-slurm": lambda: create_test_run( + partial_tr, + "srt-slurm", + with_installed_srt_repo( + SRTSlurmTestDefinition( + name="srt-slurm", + description="SRT Slurm recipe submission", + test_template_name="SRTSlurm", + git_repos=[GitRepo(url="https://github.com/NVIDIA/srt-slurm.git", commit="v2.2.1")], + cmd_args=SRTSlurmCmdArgs( + recipe_path=Path("conf/experimental/srt_slurm/recipe.yaml"), + cluster_config=Path("conf/experimental/srt_slurm/srtslurm.yaml"), + recipe_overrides={"benchmark": {"concurrencies": "4"}}, + ), + extra_env_vars={"TEST_VAR": "value with spaces"}, + extra_cmd_args={"no-preflight": ""}, + ), + slurm_system, + ), + ), "ucc": lambda: create_test_run( partial_tr, "ucc", @@ -951,25 +979,3 @@ def test_sbatch_generation(slurm_system: SlurmSystem, test_req: tuple[TestRun, s curr_aiperf = aiperf_script.read_text().strip() ref_aiperf = (Path(__file__).parent / "ref_data" / "ai-dynamo-aiperf.sh").read_text().strip() assert curr_aiperf == ref_aiperf, "aiperf.sh does not match reference" - - -def test_srt_slurm_example_dry_run(slurm_system: SlurmSystem, tmp_path: Path): - from cloudai.core import BaseInstaller - from cloudai.test_scenario_parser import TestScenarioParser - from cloudai.workloads.srt_slurm import SRTSlurmTestDefinition - - scenario = TestScenarioParser(Path("conf/experimental/srt_slurm/scenario.toml"), slurm_system, {}, {}).parse() - tr = scenario.test_runs[0].apply_params_set({"recipe_overrides.benchmark.concurrencies": "4"}) - assert isinstance(tr.test, SRTSlurmTestDefinition) - tr.output_path = tmp_path / "trial" - BaseInstaller(slurm_system).mark_as_installed(tr.test.installables) - runner = SlurmRunner("dry-run", slurm_system, TestScenario(name="srt", test_runs=[tr]), tmp_path) - job = runner._submit_test(tr) - runner.on_job_submit(tr) - assert job.id == 0 - script = tr.output_path / "run_srtctl.sh" - assert script.is_file() - assert not (tr.output_path / "cloudai_generated").exists() - assert not (tr.output_path / "srt_slurm_recipe.yaml").exists() - assert "--set" in script.read_text() - assert (tr.output_path / "test-run.toml").is_file() diff --git a/tests/test_reporter.py b/tests/test_reporter.py index 0788baa4e..69bdcd33f 100644 --- a/tests/test_reporter.py +++ b/tests/test_reporter.py @@ -658,9 +658,7 @@ def test_dse_reporter( "observed_runtime": "6m", "efficiency_ratio": "~1.3x", "efficiency_steps": "3 / 4 steps", - "best_config_toml": toml.dumps( - TestRunDetails.from_test_run(best_tr, "", "").test_definition.model_dump(mode="json") - ), + "best_config_toml": toml.dumps(TestRunDetails.from_test_run(best_tr, "", "").test_definition.model_dump()), "parameter_rows": [ { "name": "ngpus", diff --git a/tests/workloads/srt_slurm/__init__.py b/tests/workloads/srt_slurm/__init__.py deleted file mode 100644 index 3ac11851d..000000000 --- a/tests/workloads/srt_slurm/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. diff --git a/tests/workloads/srt_slurm/conftest.py b/tests/workloads/srt_slurm/conftest.py new file mode 100644 index 000000000..2463f77c8 --- /dev/null +++ b/tests/workloads/srt_slurm/conftest.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +from collections.abc import Callable +from pathlib import Path + +import pytest + +from cloudai.core import TestRun +from cloudai.workloads.srt_slurm import SRTSlurmTestDefinition + + +def _make_tdef(tmp_path: Path) -> SRTSlurmTestDefinition: + recipe_path = tmp_path / "source_recipe.yaml" + recipe_path.write_text( + """\ +schema: 2 +name: demo +model: + path: model + container: image +resources: + gpu_type: h100 + gpus_per_node: 8 +frontend: + type: sglang +engine: sglang +roles: + agg: + nodes: 1 + workers: 1 + gpus: 1 +benchmark: + type: sa-bench + isl: 128 + osl: 128 + concurrencies: "4" +""", + encoding="utf-8", + ) + (tmp_path / "cluster.yaml").write_text("default_partition: main\n") + return SRTSlurmTestDefinition.model_validate( + { + "name": "srt", + "description": "srt-slurm test", + "test_template_name": "SRTSlurm", + "git_repos": [{"url": "git@github.com:NVIDIA/srt-slurm.git", "commit": "test-revision"}], + "cmd_args": { + "recipe_path": recipe_path, + "cluster_config": tmp_path / "cluster.yaml", + "recipe_overrides": {"roles": {"agg": {"gpus": 4}}}, + }, + "extra_cmd_args": {"no-preflight": ""}, + } + ) + + +def _write_job_metadata(tdef: SRTSlurmTestDefinition, tr: TestRun, **resources) -> Path: + path = tdef.metadata_path(tr) + assert path is not None + path.parent.mkdir(parents=True, exist_ok=True) + defaults = { + "gpus_per_node": 8, + "agg_nodes": 1, + "agg_workers": 1, + "gpus_per_agg": 4, + "prefill_workers": 0, + "decode_workers": 0, + } + path.write_text(json.dumps({"job_id": path.stem, "resources": defaults | resources})) + return path + + +@pytest.fixture(name="make_tdef") +def fixture_make_tdef() -> Callable[[Path], SRTSlurmTestDefinition]: + return _make_tdef + + +@pytest.fixture(name="write_job_metadata") +def fixture_write_job_metadata() -> Callable[..., Path]: + return _write_job_metadata diff --git a/tests/workloads/srt_slurm/test_report_generation_strategy.py b/tests/workloads/srt_slurm/test_report_generation_strategy.py index 095512a8f..4bdec5018 100644 --- a/tests/workloads/srt_slurm/test_report_generation_strategy.py +++ b/tests/workloads/srt_slurm/test_report_generation_strategy.py @@ -21,10 +21,10 @@ from cloudai.systems.slurm import SlurmSystem from cloudai.workloads.srt_slurm import SRTSlurmReportGenerationStrategy -from .test_srt_slurm import make_tdef, write_job_metadata - -def test_report_reads_best_run_directly_from_srt_rollup(slurm_system: SlurmSystem, tmp_path: Path): +def test_report_reads_best_run_directly_from_srt_rollup( + slurm_system: SlurmSystem, tmp_path: Path, make_tdef, write_job_metadata +): tdef = make_tdef(tmp_path) test_run = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path / "output") test_run.output_path.mkdir(parents=True) diff --git a/tests/workloads/srt_slurm/test_results.py b/tests/workloads/srt_slurm/test_results.py index 87e97c33c..adae617f1 100644 --- a/tests/workloads/srt_slurm/test_results.py +++ b/tests/workloads/srt_slurm/test_results.py @@ -31,8 +31,6 @@ sa_bench_metrics, ) -from .test_srt_slurm import make_tdef, write_job_metadata - def test_fixed_sequence_uses_serving_gpu_counts(tmp_path: Path): gpus = JobGPUCounts(aggregate=0, prefill=24, decode=8) @@ -71,7 +69,7 @@ def test_agentx_recomputes_normalization_from_job_counts(): "runs", [None, [], [{"throughput_toks": "bad"}], [{"throughput_toks": float("nan")}], [{"throughput_toks": {"bad": 1}}]], ) -def test_bad_rollup_is_unsuccessful(slurm_system: SlurmSystem, tmp_path: Path, runs): +def test_bad_rollup_is_unsuccessful(slurm_system: SlurmSystem, tmp_path: Path, runs, make_tdef, write_job_metadata): tdef = make_tdef(tmp_path) tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path) tdef.job_id_path(tr).write_text("123") @@ -84,7 +82,9 @@ def test_bad_rollup_is_unsuccessful(slurm_system: SlurmSystem, tmp_path: Path, r assert SRTSlurmReportGenerationStrategy(slurm_system, tr).get_metric("default") == METRIC_ERROR -def test_agentx_requires_job_local_profiles_and_preserves_concurrencies(slurm_system: SlurmSystem, tmp_path: Path): +def test_agentx_requires_job_local_profiles_and_preserves_concurrencies( + slurm_system: SlurmSystem, tmp_path: Path, make_tdef, write_job_metadata +): tdef = make_tdef(tmp_path) tdef.cmd_args.result_format = "agentx" tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path) @@ -149,7 +149,9 @@ def test_agentx_requires_job_local_profiles_and_preserves_concurrencies(slurm_sy ), ], ) -def test_job_metadata_resolves_worker_gpu_totals(tmp_path: Path, resources: dict, expected: JobGPUCounts): +def test_job_metadata_resolves_worker_gpu_totals( + tmp_path: Path, resources: dict, expected: JobGPUCounts, make_tdef, write_job_metadata +): tdef = make_tdef(tmp_path) tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path) tdef.job_id_path(tr).write_text("123") @@ -158,7 +160,7 @@ def test_job_metadata_resolves_worker_gpu_totals(tmp_path: Path, resources: dict @pytest.mark.parametrize("agentic", [False, True]) -def test_missing_metadata_fails_even_with_results(slurm_system: SlurmSystem, tmp_path: Path, agentic: bool): +def test_missing_metadata_fails_even_with_results(slurm_system: SlurmSystem, tmp_path: Path, agentic: bool, make_tdef): tdef = make_tdef(tmp_path) tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path) tdef.job_id_path(tr).write_text("123") @@ -211,7 +213,9 @@ def test_missing_metadata_fails_even_with_results(slurm_system: SlurmSystem, tmp {"job_id": "123", "resources": {"agg_workers": 0, "prefill_workers": 0, "decode_workers": 0}}, ], ) -def test_invalid_metadata_is_unsuccessful(slurm_system: SlurmSystem, tmp_path: Path, invalid): +def test_invalid_metadata_is_unsuccessful( + slurm_system: SlurmSystem, tmp_path: Path, invalid, make_tdef, write_job_metadata +): tdef = make_tdef(tmp_path) tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path) tdef.job_id_path(tr).write_text("123") @@ -227,7 +231,7 @@ def test_invalid_metadata_is_unsuccessful(slurm_system: SlurmSystem, tmp_path: P assert not SRTSlurmReportGenerationStrategy(slurm_system, tr).can_handle_directory() -def test_malformed_metadata_is_unsuccessful(tmp_path: Path): +def test_malformed_metadata_is_unsuccessful(tmp_path: Path, make_tdef, write_job_metadata): tdef = make_tdef(tmp_path) tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path) tdef.job_id_path(tr).write_text("123") @@ -239,7 +243,9 @@ def test_malformed_metadata_is_unsuccessful(tmp_path: Path): @pytest.mark.parametrize("workers, per_worker", [(1, 1), (2, 2), (1, 16)]) -def test_dse_reporting_uses_each_jobs_metadata(tmp_path: Path, workers: int, per_worker: int): +def test_dse_reporting_uses_each_jobs_metadata( + tmp_path: Path, workers: int, per_worker: int, make_tdef, write_job_metadata +): tdef = make_tdef(tmp_path) base = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path) tr = base.apply_params_set( diff --git a/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py b/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py index 73a66c1db..cfd7a4dbe 100644 --- a/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py +++ b/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py @@ -20,20 +20,15 @@ from pathlib import Path import pytest -import toml import yaml -from cloudai.core import Registry, TestRun, TestScenario -from cloudai.models.scenario import ReportConfig, TestRunDetails -from cloudai.reporter import DSEReporter +from cloudai.core import TestRun, TestScenario from cloudai.systems.slurm import SlurmRunner, SlurmSystem from cloudai.workloads.srt_slurm import SRTSlurmCommandGenStrategy -from .test_srt_slurm import make_tdef - @pytest.mark.parametrize("schema", ["", "schema: 2\n"]) -def test_runner_tracks_the_recipe_job(slurm_system: SlurmSystem, tmp_path: Path, schema: str, monkeypatch): +def test_runner_tracks_the_recipe_job(slurm_system: SlurmSystem, tmp_path: Path, schema: str, monkeypatch, make_tdef): tdef = make_tdef(tmp_path) tdef.cmd_args.recipe_path.write_text(schema + "name: original\n") tdef.srtctl.venv_path = tmp_path / "venv" @@ -51,39 +46,9 @@ def test_runner_tracks_the_recipe_job(slurm_system: SlurmSystem, tmp_path: Path, assert job.id == 222 assert tdef.job_id_path(tr).read_text().strip() == "222" assert (tr.output_path / "test-run.toml").is_file() - details = TestRunDetails.model_validate(toml.load(tr.output_path / "test-run.toml")) - assert details.test_definition.cmd_args.recipe_path == tdef.cmd_args.recipe_path - assert details.test_definition.cmd_args.cluster_config == tdef.cmd_args.cluster_config - assert details.test_definition.git_repos[0].installed_path == repo - assert Registry().get_command_gen_strategy(SlurmSystem, type(tdef)) is SRTSlurmCommandGenStrategy - - -def test_best_dse_config_preserves_paths(slurm_system: SlurmSystem, tmp_path: Path): - tdef = make_tdef(tmp_path) - tdef.srtctl.venv_path = tmp_path / "venv" - tdef.cmd_args.recipe_overrides = {"roles": {"agg": {"args": {"tensor-parallel-size": [1, 2]}}}} - tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path / "results") - root = tmp_path / "results" - iteration = root / "srt" / "0" - iteration.mkdir(parents=True) - key = "recipe_overrides.roles.agg.args.tensor-parallel-size" - (iteration / "trajectory.csv").write_text( - f"step,action,reward,observation\n1,\"{{'{key}': 1}}\",10,[10]\n2,\"{{'{key}': 2}}\",20,[20]\n" - ) - for step in (1, 2): - selected = tr.apply_params_set({key: step}) - selected.output_path = iteration / str(step) - SRTSlurmCommandGenStrategy(slurm_system, selected).store_test_run() - reporter = DSEReporter(slurm_system, TestScenario(name="srt", test_runs=[tr]), root, ReportConfig()) - reporter.report_best_dse_config() - exported = toml.load(iteration / "srt.toml") - reloaded = type(tdef).model_validate(exported) - assert reloaded.cmd_args.recipe_path == tdef.cmd_args.recipe_path - assert reloaded.cmd_args.cluster_config == tdef.cmd_args.cluster_config - assert reloaded.cmd_args.recipe_overrides["roles"]["agg"]["args"]["tensor-parallel-size"] == 2 -def test_failed_submission_does_not_publish_a_job_id(slurm_system: SlurmSystem, tmp_path: Path): +def test_failed_submission_does_not_publish_a_job_id(slurm_system: SlurmSystem, tmp_path: Path, make_tdef): tdef = make_tdef(tmp_path) tdef.srtctl.venv_path = tmp_path / "venv" tdef.srtctl_path.parent.mkdir(parents=True) @@ -100,7 +65,7 @@ def test_failed_submission_does_not_publish_a_job_id(slurm_system: SlurmSystem, assert not tdef.job_id_path(tr).exists() -def test_single_sbatch_rejected(slurm_system: SlurmSystem, tmp_path: Path): +def test_single_sbatch_rejected(slurm_system: SlurmSystem, tmp_path: Path, make_tdef): tdef = make_tdef(tmp_path) tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path / "output") with pytest.raises(ValueError, match="single-sbatch"): @@ -108,7 +73,9 @@ def test_single_sbatch_rejected(slurm_system: SlurmSystem, tmp_path: Path): @pytest.mark.parametrize("extraction_fails", [False, True]) -def test_job_id_errors_are_reported_to_stderr(slurm_system: SlurmSystem, tmp_path: Path, monkeypatch, extraction_fails): +def test_job_id_errors_are_reported_to_stderr( + slurm_system: SlurmSystem, tmp_path: Path, monkeypatch, extraction_fails, make_tdef +): tdef = make_tdef(tmp_path) tdef.srtctl.venv_path = tmp_path / "venv" tdef.srtctl_path.parent.mkdir(parents=True) @@ -136,7 +103,7 @@ def test_job_id_errors_are_reported_to_stderr(slurm_system: SlurmSystem, tmp_pat assert not tdef.job_id_path(tr).exists() -def test_rendered_overrides_support_old_cli(slurm_system: SlurmSystem, tmp_path: Path): +def test_rendered_overrides_support_old_cli(slurm_system: SlurmSystem, tmp_path: Path, make_tdef): tdef = make_tdef(tmp_path) original = tdef.cmd_args.recipe_path.read_text() tdef.cmd_args.recipe_override_mode = "render" diff --git a/tests/workloads/srt_slurm/test_srt_slurm.py b/tests/workloads/srt_slurm/test_srt_slurm.py index e4b19fda0..c696ed77a 100644 --- a/tests/workloads/srt_slurm/test_srt_slurm.py +++ b/tests/workloads/srt_slurm/test_srt_slurm.py @@ -32,89 +32,7 @@ ) -def make_tdef(tmp_path: Path) -> SRTSlurmTestDefinition: - recipe_path = tmp_path / "source_recipe.yaml" - recipe_path.write_text( - """\ -schema: 2 -name: demo -model: - path: model - container: image -resources: - gpu_type: h100 - gpus_per_node: 8 -frontend: - type: sglang -engine: sglang -roles: - agg: - nodes: 1 - workers: 1 - gpus: 1 -benchmark: - type: sa-bench - isl: 128 - osl: 128 - concurrencies: "4" -""", - encoding="utf-8", - ) - (tmp_path / "cluster.yaml").write_text("default_partition: main\n") - return SRTSlurmTestDefinition.model_validate( - { - "name": "srt", - "description": "srt-slurm test", - "test_template_name": "SRTSlurm", - "git_repos": [{"url": "git@github.com:NVIDIA/srt-slurm.git", "commit": "test-revision"}], - "cmd_args": { - "recipe_path": recipe_path, - "cluster_config": tmp_path / "cluster.yaml", - "recipe_overrides": {"roles": {"agg": {"gpus": 4}}}, - }, - "extra_cmd_args": {"no-preflight": ""}, - } - ) - - -def write_job_metadata(tdef: SRTSlurmTestDefinition, tr: TestRun, **resources) -> Path: - path = tdef.metadata_path(tr) - assert path is not None - path.parent.mkdir(parents=True, exist_ok=True) - defaults = { - "gpus_per_node": 8, - "agg_nodes": 1, - "agg_workers": 1, - "gpus_per_agg": 4, - "prefill_workers": 0, - "decode_workers": 0, - } - path.write_text(json.dumps({"job_id": path.stem, "resources": defaults | resources})) - return path - - -def test_command_generation_returns_srt_slurm_job(slurm_system: SlurmSystem, tmp_path: Path): - tdef = make_tdef(tmp_path) - tdef.srtctl.venv_path = tmp_path / "venv" - test_run = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path / "output") - strategy = SRTSlurmCommandGenStrategy(slurm_system, test_run) - - command = strategy.gen_exec_command() - run_script = test_run.output_path / "run_srtctl.sh" - script = run_script.read_text(encoding="utf-8") - - subprocess.run(["bash", "-n", run_script], check=True) - assert command == f"bash {run_script}" - assert str(tdef.cmd_args.recipe_path) in script - assert "--set roles.agg.gpus=4" in script - assert not (run_script.parent / "recipe.yaml").exists() - assert "--no-preflight" in run_script.read_text(encoding="utf-8") - assert "srt_job_id.txt" in run_script.read_text(encoding="utf-8") - assert 'echo "Submitted batch job $SRT_JOB_ID"' in run_script.read_text(encoding="utf-8") - assert "--dependency" not in run_script.read_text(encoding="utf-8") - - -def test_run_script_exports_merged_environment(slurm_system: SlurmSystem, tmp_path: Path, monkeypatch): +def test_run_script_exports_merged_environment(slurm_system: SlurmSystem, tmp_path: Path, monkeypatch, make_tdef): monkeypatch.delenv("SRTSLURM_CONFIG", raising=False) monkeypatch.setenv("VIRTUAL_ENV", str(tmp_path / "login-node-venv")) slurm_system.global_env_vars = {"SRTSLURM_CONFIG": "system.yaml", "SRT_GLOBAL_ONLY": "global value"} @@ -252,7 +170,7 @@ def test_existing_recipe_is_loaded_and_dse_values_are_overlaid(tmp_path: Path): assert "cmd_args.recipe_overrides" not in tdef.dse_excluded_args -def test_standard_git_repos_controls_srtctl_revision(tmp_path: Path): +def test_standard_git_repos_controls_srtctl_revision(tmp_path: Path, make_tdef): config = make_tdef(tmp_path).model_dump() config["extra_env_vars"] = {"SRTSLURM_CONFIG": str(tmp_path / "cluster.yaml")} config["git_repos"][0]["commit"] = "updated-revision" @@ -267,7 +185,7 @@ def test_standard_git_repos_controls_srtctl_revision(tmp_path: Path): @pytest.mark.parametrize("git_repos", [None, []]) -def test_srtctl_requires_a_configured_git_repo(tmp_path: Path, git_repos): +def test_srtctl_requires_a_configured_git_repo(tmp_path: Path, git_repos, make_tdef): config = make_tdef(tmp_path).model_dump() if git_repos is None: config.pop("git_repos") From 3f5fd5307a18b4c21521bd56f6c0c21a9183732a Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Thu, 29 Oct 2026 07:56:51 -0600 Subject: [PATCH 11/20] Add portable AgentX DSE example Signed-off-by: Jakob Johnson --- .../srt_slurm/agentx_dse/README.md | 45 ++++++++++++++++ .../srt_slurm/agentx_dse/agentx_client.sh | 49 ++++++++++++++++++ .../srt_slurm/agentx_dse/recipe.yaml | 26 ++++++++++ .../srt_slurm/agentx_dse/scenario.toml | 51 +++++++++++++++++++ .../srt_slurm/agentx_dse/srtslurm.yaml | 13 +++++ .../srt_slurm/agentx_dse/system.toml | 26 ++++++++++ 6 files changed, 210 insertions(+) create mode 100644 conf/experimental/srt_slurm/agentx_dse/README.md create mode 100644 conf/experimental/srt_slurm/agentx_dse/agentx_client.sh create mode 100644 conf/experimental/srt_slurm/agentx_dse/recipe.yaml create mode 100644 conf/experimental/srt_slurm/agentx_dse/scenario.toml create mode 100644 conf/experimental/srt_slurm/agentx_dse/srtslurm.yaml create mode 100644 conf/experimental/srt_slurm/agentx_dse/system.toml diff --git a/conf/experimental/srt_slurm/agentx_dse/README.md b/conf/experimental/srt_slurm/agentx_dse/README.md new file mode 100644 index 000000000..22eab6131 --- /dev/null +++ b/conf/experimental/srt_slurm/agentx_dse/README.md @@ -0,0 +1,45 @@ +# AgentX DSE smoke test + +This example provisions pinned SRT Slurm and InferenceX repositories, then runs +SGLang with Qwen3-0.6B on one serving GPU of a four-GPU GB200 node. Grid search +tries context lengths 2048 and 4096 and maximizes `p90_intvty`. + +The client replays two public fixtures from InferenceX's pinned AIPerf submodule +for 60 seconds. It uses `--unsafe-override`; this is a correctness smoke test, +not a canonical SemiAnalysis submission or a performance comparison. + +## Configure + +All `/shared/...` paths and the `gpu` partition are placeholders: + +- Put these five input files in `/shared/cloudai/example`, or update the + scenario paths and the corresponding `/agentx-test` mount together. +- Set the installation and results locations in `system.toml`. If you change + the installation root, update `srtctl_root` and both InferenceX paths in + `srtslurm.yaml` and `scenario.toml` as well. +- Provide a Qwen3-0.6B model directory at `/shared/models/Qwen3-0.6B` and a + compatible SGLang container at `/shared/images/sglang.sqsh`. The client uses + the model's mounted tokenizer without downloading it from Hugging Face. +- Choose your partition in both system and SRT cluster configs. CloudAI uses + its partition for the setup allocation; SRT Slurm owns benchmark scheduling. +- Configure the Slurm CLI environment and any site-required container settings. + Dependencies require network access. uv caches and managed Python downloads + use node-local `/tmp`; other application caches still need writable storage. + +The first repository entry installs SRT Slurm. The second installs InferenceX +with its pinned submodules. `INFMAX_WORKSPACE` and the `/infmax-workspace` mount +connect that checkout to the custom benchmark client. + +## Run + +After adapting the paths, run from the directory containing the input files: + +```bash +mkdir -p /shared/cloudai/install /shared/cloudai/results +cloudai install --system-config system.toml --test-scenario scenario.toml +cloudai run --system-config system.toml --test-scenario scenario.toml +``` + +Each trial writes its SRT job artifacts and `srt_slurm_results.csv`. The DSE +iteration writes `trajectory.csv` and the selected `agentx.toml` configuration; +scenario and DSE HTML reports summarize the results. diff --git a/conf/experimental/srt_slurm/agentx_dse/agentx_client.sh b/conf/experimental/srt_slurm/agentx_dse/agentx_client.sh new file mode 100644 index 000000000..9ed840f29 --- /dev/null +++ b/conf/experimental/srt_slurm/agentx_dse/agentx_client.sh @@ -0,0 +1,49 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -euo pipefail +export INFMAX_CONTAINER_WORKSPACE=/infmax-workspace +source "$INFMAX_CONTAINER_WORKSPACE/benchmarks/benchmark_lib.sh" +install_agentic_deps + +# Use the mounted model tokenizer without fetching it from the Hub. +export HF_HUB_CACHE="$(mktemp -d)" +export HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 +tokenizer_cache="$HF_HUB_CACHE/models--Qwen--Qwen3-0.6B" +revision=c1899de289a04d12100db370d81485cdf75e47ca +mkdir -p "$tokenizer_cache/refs" "$tokenizer_cache/snapshots" +ln -s /model "$tokenizer_cache/snapshots/$revision" +printf '%s' "$revision" > "$tokenizer_cache/refs/main" + +# Replay the same two pinned public test fixtures in each trial. +traces="$(mktemp -d)" +cp "$AIPERF_DIR/tests/fixtures/weka_traces/simple.json" "$traces/" +cp "$AIPERF_DIR/tests/fixtures/weka_traces/one_subagent.json" "$traces/" +export MODEL=Qwen/Qwen3-0.6B MODEL_PREFIX=qwen3 FRAMEWORK=sglang PRECISION=bf16 +export DURATION=60 CONC=1 PORT=8000 RESULT_FILENAME=qwen_c1 +export IS_MULTINODE=false DISAGG=false KV_OFFLOADING=none KV_OFFLOAD_BACKEND='' +export AIPERF_WARMUP_REQUESTS_PER_LANE=1 AIPERF_EXPERIMENTAL_FAST=0 +export AIPERF_EXTRA_INPUTS=ignore_eos:true +export AGENTIC_OUTPUT_DIR=/logs/agentic +result_dir="$AGENTIC_OUTPUT_DIR/conc_1" +mkdir -p "$result_dir" +TRACE_SOURCE_FLAG="--input-file $traces --custom-dataset-type weka_trace" +build_replay_cmd "$result_dir" +REPLAY_CMD="${REPLAY_CMD/--trajectory-start-min-ratio 0.25/--trajectory-start-min-ratio 0.0}" +REPLAY_CMD="${REPLAY_CMD/--trajectory-start-max-ratio 0.75/--trajectory-start-max-ratio 0.0}" +REPLAY_CMD+=' --tokenizer Qwen/Qwen3-0.6B' +run_agentic_replay_and_write_outputs "$result_dir" diff --git a/conf/experimental/srt_slurm/agentx_dse/recipe.yaml b/conf/experimental/srt_slurm/agentx_dse/recipe.yaml new file mode 100644 index 000000000..beebd402d --- /dev/null +++ b/conf/experimental/srt_slurm/agentx_dse/recipe.yaml @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +schema: 2 +name: qwen-agentx +model: {path: qwen, container: sglang, precision: bf16} +resources: {gpu_type: gb200, gpus_per_node: 4} +slurm: {time_limit: '00:20:00'} +frontend: {type: sglang, enable_multiple_frontends: false} +engine: sglang +roles: + agg: + nodes: 1 + workers: 1 + gpus: 1 + args: + tensor-parallel-size: 1 + context-length: 2048 + mem-fraction-static: 0.2 + disable-cuda-graph: true + served-model-name: Qwen/Qwen3-0.6B +health_check: {max_attempts: 60, interval_seconds: 5} +benchmark: + type: custom + command: bash /agentx-test/agentx_client.sh diff --git a/conf/experimental/srt_slurm/agentx_dse/scenario.toml b/conf/experimental/srt_slurm/agentx_dse/scenario.toml new file mode 100644 index 000000000..0892063ed --- /dev/null +++ b/conf/experimental/srt_slurm/agentx_dse/scenario.toml @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +name = "agentx-dse" + +[[Tests]] +id = "agentx" +name = "qwen_agentx" +description = "Two-step AgentX DSE correctness smoke test" +test_template_name = "SRTSlurm" +agent = "grid_search" +agent_steps = 2 +agent_metrics = ["p90_intvty"] +agent_reward_function = "identity" + + [[Tests.git_repos]] + url = "https://github.com/NVIDIA/srt-slurm.git" + commit = "098e15aca8504d4747c6f62c7a845e178b7f4b10" + + [[Tests.git_repos]] + url = "https://github.com/SemiAnalysisAI/InferenceX.git" + commit = "f8dff19a6898e3ee0dd95232f867151d4d2d6392" + init_submodules = true + + [Tests.cmd_args] + recipe_path = "/shared/cloudai/example/recipe.yaml" + cluster_config = "/shared/cloudai/example/srtslurm.yaml" + result_format = "agentx" + + [Tests.cmd_args.recipe_overrides.roles.agg.args] + context-length = [2048, 4096] + + [Tests.extra_env_vars] + INFMAX_WORKSPACE = "/shared/cloudai/install/InferenceX__f8dff19a6898e3ee0dd95232f867151d4d2d6392" + + # Keep host-side uv work off shared home storage. + UV_CACHE_DIR = "/tmp/cloudai-agentx-uv-cache" + UV_PYTHON_INSTALL_DIR = "/tmp/cloudai-agentx-python" diff --git a/conf/experimental/srt_slurm/agentx_dse/srtslurm.yaml b/conf/experimental/srt_slurm/agentx_dse/srtslurm.yaml new file mode 100644 index 000000000..1e4eb5d8d --- /dev/null +++ b/conf/experimental/srt_slurm/agentx_dse/srtslurm.yaml @@ -0,0 +1,13 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +default_partition: gpu +gpus_per_node: 4 +srtctl_root: /shared/cloudai/install/srt-slurm__098e15aca8504d4747c6f62c7a845e178b7f4b10 +model_paths: {qwen: /shared/models/Qwen3-0.6B} +containers: {sglang: /shared/images/sglang.sqsh} +default_mounts: + /shared/cloudai/install/InferenceX__f8dff19a6898e3ee0dd95232f867151d4d2d6392: /infmax-workspace + /shared/cloudai/example: /agentx-test + /shared/models/Qwen3-0.6B: /model diff --git a/conf/experimental/srt_slurm/agentx_dse/system.toml b/conf/experimental/srt_slurm/agentx_dse/system.toml new file mode 100644 index 000000000..1d3d67df6 --- /dev/null +++ b/conf/experimental/srt_slurm/agentx_dse/system.toml @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +name = "example" +scheduler = "slurm" +install_path = "/shared/cloudai/install" +output_path = "/shared/cloudai/results" +default_partition = "gpu" +gpus_per_node = 4 +monitor_interval = 15 + +[[partitions]] +name = "gpu" From 8511b60be6f80a83b9c9f8ce223d4d6a1cee1e5e Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Fri, 30 Oct 2026 12:16:21 -0600 Subject: [PATCH 12/20] remove unneeded system.toml --- .../srt_slurm/agentx_dse/README.md | 18 ++++++++----- .../srt_slurm/agentx_dse/system.toml | 26 ------------------- 2 files changed, 12 insertions(+), 32 deletions(-) delete mode 100644 conf/experimental/srt_slurm/agentx_dse/system.toml diff --git a/conf/experimental/srt_slurm/agentx_dse/README.md b/conf/experimental/srt_slurm/agentx_dse/README.md index 22eab6131..251cc263a 100644 --- a/conf/experimental/srt_slurm/agentx_dse/README.md +++ b/conf/experimental/srt_slurm/agentx_dse/README.md @@ -12,9 +12,14 @@ not a canonical SemiAnalysis submission or a performance comparison. All `/shared/...` paths and the `gpu` partition are placeholders: -- Put these five input files in `/shared/cloudai/example`, or update the - scenario paths and the corresponding `/agentx-test` mount together. -- Set the installation and results locations in `system.toml`. If you change +- Copy the common [Slurm system template](../../../common/system/example_slurm_cluster.toml) + to your own `system.toml`. Adapt its partition, node groups, GPU counts, and + environment variables to your cluster. +- Put `scenario.toml`, `srtslurm.yaml`, `recipe.yaml`, and `agentx_client.sh` in + `/shared/cloudai/example`, or update the scenario paths and the corresponding + `/agentx-test` mount together. +- Set `install_path = "/shared/cloudai/install"` and + `output_path = "/shared/cloudai/results"` in your system config. If you change the installation root, update `srtctl_root` and both InferenceX paths in `srtslurm.yaml` and `scenario.toml` as well. - Provide a Qwen3-0.6B model directory at `/shared/models/Qwen3-0.6B` and a @@ -32,12 +37,13 @@ connect that checkout to the custom benchmark client. ## Run -After adapting the paths, run from the directory containing the input files: +After adapting the paths, run from the directory containing the input files. +Pass the path to your adapted copy of the common system template: ```bash mkdir -p /shared/cloudai/install /shared/cloudai/results -cloudai install --system-config system.toml --test-scenario scenario.toml -cloudai run --system-config system.toml --test-scenario scenario.toml +cloudai install --system-config /path/to/system.toml --test-scenario scenario.toml +cloudai run --system-config /path/to/system.toml --test-scenario scenario.toml ``` Each trial writes its SRT job artifacts and `srt_slurm_results.csv`. The DSE diff --git a/conf/experimental/srt_slurm/agentx_dse/system.toml b/conf/experimental/srt_slurm/agentx_dse/system.toml deleted file mode 100644 index 1d3d67df6..000000000 --- a/conf/experimental/srt_slurm/agentx_dse/system.toml +++ /dev/null @@ -1,26 +0,0 @@ -# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -name = "example" -scheduler = "slurm" -install_path = "/shared/cloudai/install" -output_path = "/shared/cloudai/results" -default_partition = "gpu" -gpus_per_node = 4 -monitor_interval = 15 - -[[partitions]] -name = "gpu" From a13d89f534ec689f171663c84b39f4f2f12efdbe Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Fri, 30 Oct 2026 12:23:50 -0600 Subject: [PATCH 13/20] remove git_repo on subclass and rely on repo name instead of order --- .../srt_slurm/agentx_dse/README.md | 4 ++-- doc/workloads/srt_slurm.rst | 5 +++-- .../srt_slurm/slurm_command_gen_strategy.py | 2 +- src/cloudai/workloads/srt_slurm/srt_slurm.py | 20 ++++++++++++----- tests/workloads/srt_slurm/test_srt_slurm.py | 22 +++++++++++++------ 5 files changed, 35 insertions(+), 18 deletions(-) diff --git a/conf/experimental/srt_slurm/agentx_dse/README.md b/conf/experimental/srt_slurm/agentx_dse/README.md index 251cc263a..b3bba31cd 100644 --- a/conf/experimental/srt_slurm/agentx_dse/README.md +++ b/conf/experimental/srt_slurm/agentx_dse/README.md @@ -31,8 +31,8 @@ All `/shared/...` paths and the `gpu` partition are placeholders: Dependencies require network access. uv caches and managed Python downloads use node-local `/tmp`; other application caches still need writable storage. -The first repository entry installs SRT Slurm. The second installs InferenceX -with its pinned submodules. `INFMAX_WORKSPACE` and the `/infmax-workspace` mount +The `srt-slurm` repository entry supplies srtctl, regardless of entry order. +The InferenceX entry installs its pinned submodules. `INFMAX_WORKSPACE` and the `/infmax-workspace` mount connect that checkout to the custom benchmark client. ## Run diff --git a/doc/workloads/srt_slurm.rst b/doc/workloads/srt_slurm.rst index 42dacd0be..104951091 100644 --- a/doc/workloads/srt_slurm.rst +++ b/doc/workloads/srt_slurm.rst @@ -36,8 +36,9 @@ and must be accessible on the shared filesystem. [cmd_args.recipe_overrides.benchmark] concurrencies = [[1, 4], [8, 16]] -The first ``git_repos`` entry supplies srtctl; additional entries use ordinary -CloudAI installation. Pin a revision supporting ``srtctl apply --set``. +The repository named ``srt-slurm`` in ``git_repos`` supplies srtctl, regardless +of entry order. Forks with that repository name are supported; other entries use +ordinary CloudAI installation. Pin a revision supporting ``srtctl apply --set``. For older revisions without ``--set``, set ``cmd_args.recipe_override_mode = "render"``. This overlays existing mapping keys in a concrete recipe, preserves YAML value types, and submits a trial-local ``srt_slurm_recipe.yaml``. It rejects missing keys diff --git a/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py b/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py index d2bf2ff1a..29aa9df63 100644 --- a/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py +++ b/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py @@ -86,7 +86,7 @@ def _prepare_recipe(self, paths: dict[str, Path]) -> tuple[Path, list[str]]: def write_run_script(self) -> None: paths = self._paths() - repo = self.tdef.git_repos[0] + repo = self.tdef.srtctl.git_repo repo_path = (repo.installed_path or self.system.install_path / repo.repo_name).absolute() cluster_config = self.tdef.cmd_args.cluster_config.expanduser().absolute() recipe, overrides = self._prepare_recipe(paths) diff --git a/src/cloudai/workloads/srt_slurm/srt_slurm.py b/src/cloudai/workloads/srt_slurm/srt_slurm.py index 8145cacf0..b928296e8 100644 --- a/src/cloudai/workloads/srt_slurm/srt_slurm.py +++ b/src/cloudai/workloads/srt_slurm/srt_slurm.py @@ -19,7 +19,7 @@ from pydantic import Field, PrivateAttr, model_validator -from cloudai.core import GitRepo, Installable, JobStatusResult, TestRun +from cloudai.core import Installable, JobStatusResult, TestRun from cloudai.models.workload import CmdArgs, TestDefinition from .installable import SRTSlurmExecutable @@ -49,9 +49,6 @@ class SRTSlurmTestDefinition(TestDefinition): """Run one declarative srt-slurm recipe as one CloudAI trial.""" cmd_args: SRTSlurmCmdArgs - git_repos: list[GitRepo] = Field( - default_factory=list, min_length=1, validate_default=True, description="The first repository supplies srtctl." - ) _srtctl: SRTSlurmExecutable | None = PrivateAttr(default=None) @model_validator(mode="after") @@ -62,7 +59,17 @@ def exclude_recipe_structure_from_dse(self) -> "SRTSlurmTestDefinition": @property def srtctl(self) -> SRTSlurmExecutable: if self._srtctl is None: - self._srtctl = SRTSlurmExecutable(git_repo=self.git_repos[0]) + repo = next( + ( + repo + for repo in self.git_repos + if repo.url.rstrip("/").rsplit("/", 1)[-1].removesuffix(".git") == "srt-slurm" + ), + None, + ) + if repo is None: + raise ValueError("SRTSlurm requires a repository named 'srt-slurm' in git_repos.") + self._srtctl = SRTSlurmExecutable(git_repo=repo) return self._srtctl @property @@ -73,7 +80,8 @@ def srtctl_path(self) -> Path: @property def installables(self) -> list[Installable]: - return [self.srtctl, *self.git_repos[1:]] + srtctl = self.srtctl + return [srtctl, *(repo for repo in self.git_repos if repo is not srtctl.git_repo)] @staticmethod def job_id_path(tr: TestRun) -> Path: diff --git a/tests/workloads/srt_slurm/test_srt_slurm.py b/tests/workloads/srt_slurm/test_srt_slurm.py index c696ed77a..b80002466 100644 --- a/tests/workloads/srt_slurm/test_srt_slurm.py +++ b/tests/workloads/srt_slurm/test_srt_slurm.py @@ -22,7 +22,6 @@ import pytest import yaml -from pydantic import ValidationError from cloudai.core import GitRepo, TestRun from cloudai.systems.slurm import SlurmSystem @@ -40,6 +39,7 @@ def test_run_script_exports_merged_environment(slurm_system: SlurmSystem, tmp_pa repo_path = tmp_path / "installed srt-slurm" repo_path.mkdir() tdef.git_repos[0].installed_path = repo_path + tdef.git_repos.insert(0, GitRepo(url="https://github.com/example/recipes.git", commit="main")) caller_path = tmp_path / "cloudaix" (caller_path / ".venv" / "bin").mkdir(parents=True) (caller_path / ".venv" / "bin" / "python").write_bytes(b"invalid login-node interpreter") @@ -170,21 +170,28 @@ def test_existing_recipe_is_loaded_and_dse_values_are_overlaid(tmp_path: Path): assert "cmd_args.recipe_overrides" not in tdef.dse_excluded_args -def test_standard_git_repos_controls_srtctl_revision(tmp_path: Path, make_tdef): +@pytest.mark.parametrize("srt_first", [True, False]) +@pytest.mark.parametrize("url", ["https://github.com/example/srt-slurm.git", "git@github.com:example/srt-slurm"]) +def test_standard_git_repos_controls_srtctl_revision(tmp_path: Path, make_tdef, srt_first: bool, url: str): config = make_tdef(tmp_path).model_dump() config["extra_env_vars"] = {"SRTSLURM_CONFIG": str(tmp_path / "cluster.yaml")} config["git_repos"][0]["commit"] = "updated-revision" + config["git_repos"][0]["url"] = url config["git_repos"].append({"url": "https://github.com/example/recipes.git", "commit": "main"}) + if not srt_first: + config["git_repos"].reverse() tdef = SRTSlurmTestDefinition.model_validate(config) + srt_repo = tdef.git_repos[0 if srt_first else 1] + other_repo = tdef.git_repos[1 if srt_first else 0] - assert tdef.srtctl.git_repo is tdef.git_repos[0] + assert tdef.srtctl.git_repo is srt_repo assert tdef.srtctl.git_repo.commit == "updated-revision" assert "updated-revision" in tdef.srtctl.venv_name - assert tdef.installables == [tdef.srtctl, tdef.git_repos[1]] + assert tdef.installables == [tdef.srtctl, other_repo] assert isinstance(tdef.installables[1], GitRepo) -@pytest.mark.parametrize("git_repos", [None, []]) +@pytest.mark.parametrize("git_repos", [None, [], [{"url": "https://github.com/example/recipes.git", "commit": "main"}]]) def test_srtctl_requires_a_configured_git_repo(tmp_path: Path, git_repos, make_tdef): config = make_tdef(tmp_path).model_dump() if git_repos is None: @@ -192,8 +199,9 @@ def test_srtctl_requires_a_configured_git_repo(tmp_path: Path, git_repos, make_t else: config["git_repos"] = git_repos - with pytest.raises(ValidationError, match="git_repos"): - SRTSlurmTestDefinition.model_validate(config) + tdef = SRTSlurmTestDefinition.model_validate(config) + with pytest.raises(ValueError, match="requires a repository named 'srt-slurm' in git_repos"): + _ = tdef.installables @pytest.mark.parametrize("value", [4, 0.75, True, False, None, [32, 64], "true", "720", "", "$HOME; $(id)"]) From b1fc587e0a61e26306a10d654feffa5caf6cc4c6 Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Fri, 30 Oct 2026 12:30:38 -0600 Subject: [PATCH 14/20] remove bash functionality tests --- tests/test_acceptance.py | 5 +- .../test_slurm_command_gen_strategy.py | 76 +---------------- tests/workloads/srt_slurm/test_srt_slurm.py | 85 ------------------- 3 files changed, 6 insertions(+), 160 deletions(-) diff --git a/tests/test_acceptance.py b/tests/test_acceptance.py index efd8b1572..a6852cf36 100644 --- a/tests/test_acceptance.py +++ b/tests/test_acceptance.py @@ -367,7 +367,10 @@ def test_req(request, slurm_system: SlurmSystem, partial_tr: partial[TestRun]) - name="srt-slurm", description="SRT Slurm recipe submission", test_template_name="SRTSlurm", - git_repos=[GitRepo(url="https://github.com/NVIDIA/srt-slurm.git", commit="v2.2.1")], + git_repos=[ + GitRepo(url="https://github.com/example/recipes.git", commit="main"), + GitRepo(url="https://github.com/NVIDIA/srt-slurm.git", commit="v2.2.1"), + ], cmd_args=SRTSlurmCmdArgs( recipe_path=Path("conf/experimental/srt_slurm/recipe.yaml"), cluster_config=Path("conf/experimental/srt_slurm/srtslurm.yaml"), diff --git a/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py b/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py index cfd7a4dbe..1048a8b99 100644 --- a/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py +++ b/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py @@ -14,57 +14,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os -import subprocess -import sys from pathlib import Path import pytest import yaml -from cloudai.core import TestRun, TestScenario -from cloudai.systems.slurm import SlurmRunner, SlurmSystem +from cloudai.core import TestRun +from cloudai.systems.slurm import SlurmSystem from cloudai.workloads.srt_slurm import SRTSlurmCommandGenStrategy -@pytest.mark.parametrize("schema", ["", "schema: 2\n"]) -def test_runner_tracks_the_recipe_job(slurm_system: SlurmSystem, tmp_path: Path, schema: str, monkeypatch, make_tdef): - tdef = make_tdef(tmp_path) - tdef.cmd_args.recipe_path.write_text(schema + "name: original\n") - tdef.srtctl.venv_path = tmp_path / "venv" - tdef.srtctl_path.parent.mkdir(parents=True) - repo = slurm_system.install_path / tdef.git_repos[0].repo_name - repo.mkdir(parents=True) - tdef.git_repos[0].installed_path = repo - tdef.srtctl_path.write_text(f"#!{sys.executable}\nprint('Submitted batch job 111')\nprint('Job 222 submitted')\n") - tdef.srtctl_path.chmod(0o755) - monkeypatch.chdir(tmp_path) - tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=Path("relative output")) - runner = SlurmRunner("run", slurm_system, TestScenario(name="srt", test_runs=[tr]), tmp_path) - runner.submit_test(tr) - job = runner.jobs[0] - assert job.id == 222 - assert tdef.job_id_path(tr).read_text().strip() == "222" - assert (tr.output_path / "test-run.toml").is_file() - - -def test_failed_submission_does_not_publish_a_job_id(slurm_system: SlurmSystem, tmp_path: Path, make_tdef): - tdef = make_tdef(tmp_path) - tdef.srtctl.venv_path = tmp_path / "venv" - tdef.srtctl_path.parent.mkdir(parents=True) - repo = slurm_system.install_path / tdef.git_repos[0].repo_name - repo.mkdir(parents=True) - tdef.srtctl_path.write_text(f"#!{sys.executable}\nprint('Submitted batch job 111')\nraise SystemExit(2)\n") - tdef.srtctl_path.chmod(0o755) - tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path / "output") - strategy = SRTSlurmCommandGenStrategy(slurm_system, tr) - strategy.gen_exec_command() - completed = subprocess.run(["bash", strategy._paths()["run_script"]], capture_output=True, text=True) - assert completed.returncode != 0 - assert "Submitted batch job" not in completed.stdout - assert not tdef.job_id_path(tr).exists() - - def test_single_sbatch_rejected(slurm_system: SlurmSystem, tmp_path: Path, make_tdef): tdef = make_tdef(tmp_path) tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path / "output") @@ -72,37 +31,6 @@ def test_single_sbatch_rejected(slurm_system: SlurmSystem, tmp_path: Path, make_ SRTSlurmCommandGenStrategy(slurm_system, tr).gen_srun_command() -@pytest.mark.parametrize("extraction_fails", [False, True]) -def test_job_id_errors_are_reported_to_stderr( - slurm_system: SlurmSystem, tmp_path: Path, monkeypatch, extraction_fails, make_tdef -): - tdef = make_tdef(tmp_path) - tdef.srtctl.venv_path = tmp_path / "venv" - tdef.srtctl_path.parent.mkdir(parents=True) - repo = slurm_system.install_path / tdef.git_repos[0].repo_name - repo.mkdir(parents=True) - response = "Submitted batch job 111" if extraction_fails else "Submission completed without a job ID" - tdef.srtctl_path.write_text(f"#!{sys.executable}\nprint({response!r})\n") - tdef.srtctl_path.chmod(0o755) - if extraction_fails: - tools = tmp_path / "tools" - tools.mkdir() - sed = tools / "sed" - sed.write_text("#!/usr/bin/env bash\nexit 1\n") - sed.chmod(0o755) - monkeypatch.setenv("PATH", f"{tools}:{os.environ['PATH']}") - tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path / "output") - strategy = SRTSlurmCommandGenStrategy(slurm_system, tr) - strategy.gen_exec_command() - tdef.job_id_path(tr).write_text("999\n") - completed = subprocess.run(["bash", strategy._paths()["run_script"]], capture_output=True, text=True) - assert completed.returncode != 0 - assert completed.stdout == "" - assert response in completed.stderr - assert "Unable to find submitted srt-slurm job ID" in completed.stderr - assert not tdef.job_id_path(tr).exists() - - def test_rendered_overrides_support_old_cli(slurm_system: SlurmSystem, tmp_path: Path, make_tdef): tdef = make_tdef(tmp_path) original = tdef.cmd_args.recipe_path.read_text() diff --git a/tests/workloads/srt_slurm/test_srt_slurm.py b/tests/workloads/srt_slurm/test_srt_slurm.py index b80002466..4f00b828b 100644 --- a/tests/workloads/srt_slurm/test_srt_slurm.py +++ b/tests/workloads/srt_slurm/test_srt_slurm.py @@ -14,103 +14,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json -import os -import subprocess -import sys from pathlib import Path import pytest import yaml from cloudai.core import GitRepo, TestRun -from cloudai.systems.slurm import SlurmSystem from cloudai.workloads.srt_slurm import ( SRTSlurmCommandGenStrategy, SRTSlurmTestDefinition, ) -def test_run_script_exports_merged_environment(slurm_system: SlurmSystem, tmp_path: Path, monkeypatch, make_tdef): - monkeypatch.delenv("SRTSLURM_CONFIG", raising=False) - monkeypatch.setenv("VIRTUAL_ENV", str(tmp_path / "login-node-venv")) - slurm_system.global_env_vars = {"SRTSLURM_CONFIG": "system.yaml", "SRT_GLOBAL_ONLY": "global value"} - tdef = make_tdef(tmp_path) - repo_path = tmp_path / "installed srt-slurm" - repo_path.mkdir() - tdef.git_repos[0].installed_path = repo_path - tdef.git_repos.insert(0, GitRepo(url="https://github.com/example/recipes.git", commit="main")) - caller_path = tmp_path / "cloudaix" - (caller_path / ".venv" / "bin").mkdir(parents=True) - (caller_path / ".venv" / "bin" / "python").write_bytes(b"invalid login-node interpreter") - (caller_path / ".venv" / "pyvenv.cfg").write_text("home = /login-node/python\n") - active_bin = tmp_path / "login-node-venv" / "bin" - active_bin.mkdir(parents=True) - (active_bin.parent / "pyvenv.cfg").write_text("home = /login-node/python\n") - retained_bin = tmp_path / "cluster tools" / "bin" - retained_bin.mkdir(parents=True) - other_venv_bin = tmp_path / "other login-node venv" / "bin" - other_venv_bin.mkdir(parents=True) - (other_venv_bin.parent / "pyvenv.cfg").write_text("home = /other/login-node/python\n") - monkeypatch.setenv( - "PATH", - os.pathsep.join( - [ - str(caller_path / ".venv" / "bin"), - str(active_bin), - str(other_venv_bin), - str(retained_bin), - os.environ["PATH"], - ] - ), - ) - tdef.extra_env_vars = { - "SRTSLURM_CONFIG": str(tmp_path / "cluster's config" / "srtslurm.yaml"), - "SRT_LITERAL": "$HOME; $(printf injected)", - "SRT_EMPTY": "", - } - expected_config = str(tdef.cmd_args.cluster_config.absolute()) - tdef.srtctl.venv_path = tmp_path / "venv" - tdef.srtctl_path.parent.mkdir(parents=True) - tdef.srtctl_path.write_text( - f"#!{sys.executable}\n" - "import json, os, sys\n" - "print(json.dumps({key: os.environ.get(key) for key in " - "['SRTSLURM_CONFIG', 'SRT_GLOBAL_ONLY', 'SRT_LITERAL', 'SRT_EMPTY', 'VIRTUAL_ENV']}))\n" - "print(os.getcwd())\n" - "print(os.environ['PATH'])\n" - "print(json.dumps(sys.argv[1:]))\n" - "print('Job 12345 submitted')\n", - encoding="utf-8", - ) - tdef.srtctl_path.chmod(0o755) - test_run = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path / "output") - strategy = SRTSlurmCommandGenStrategy(slurm_system, test_run) - strategy.gen_exec_command() - - subprocess.run( - ["bash", strategy._paths()["run_script"]], cwd=caller_path, check=True, capture_output=True, text=True - ) - received = json.loads(strategy._paths()["run_log"].read_text(encoding="utf-8").splitlines()[0]) - - assert received == { - "SRT_GLOBAL_ONLY": "global value", - **tdef.extra_env_vars, - "SRTSLURM_CONFIG": expected_config, - "VIRTUAL_ENV": None, - } - assert Path(strategy._paths()["run_log"].read_text().splitlines()[1]) == repo_path - submitted_path = strategy._paths()["run_log"].read_text().splitlines()[2].split(os.pathsep) - assert str(caller_path / ".venv" / "bin") not in submitted_path - assert str(active_bin) not in submitted_path - assert str(other_venv_bin) not in submitted_path - assert str(retained_bin) in submitted_path - submitted_args = json.loads(strategy._paths()["run_log"].read_text().splitlines()[3]) - assert submitted_args[submitted_args.index("-f") + 1] == str(tdef.cmd_args.recipe_path) - assert submitted_args[submitted_args.index("--set") + 1] == "roles.agg.gpus=4" - assert tdef.job_id_path(test_run).read_text(encoding="utf-8").strip() == "12345" - - def test_existing_recipe_is_loaded_and_dse_values_are_overlaid(tmp_path: Path): source = tmp_path / "existing.yaml" original = """\ From 6f985244b85c97c5533c8fb1df8879566e7fee8f Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Fri, 30 Oct 2026 12:42:22 -0600 Subject: [PATCH 15/20] add clarifying comments --- .../workloads/srt_slurm/slurm_command_gen_strategy.py | 5 +++-- tests/ref_data/srt-slurm.sbatch | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py b/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py index 29aa9df63..415c87933 100644 --- a/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py +++ b/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py @@ -130,9 +130,10 @@ def write_run_script(self) -> None: "", "# Submit the recipe", f"mkdir -p {shlex.quote(str(paths['nested_outputs']))}", + "# Clear the old result lookup ID in case this wrapper is manually rerun and submission fails.", f"rm -f {shlex.quote(str(paths['job_id']))}", - "# C: uv discovers .venv from the current directory even with --no-project.", - "# Submit from srt-slurm so it cannot discover CloudAI's checkout .venv.", + "# C: Slurm inherits this working directory. Leave CloudAI's checkout so", + "# compute-side uv does not discover its login-node .venv, even with --no-project.", f"cd {shlex.quote(str(repo_path))}", f"if ! {apply_command} > {shlex.quote(str(paths['run_log']))} 2>&1; then", f" cat {shlex.quote(str(paths['run_log']))} >&2", diff --git a/tests/ref_data/srt-slurm.sbatch b/tests/ref_data/srt-slurm.sbatch index 95ec4730a..591bbd893 100644 --- a/tests/ref_data/srt-slurm.sbatch +++ b/tests/ref_data/srt-slurm.sbatch @@ -22,9 +22,10 @@ unset VIRTUAL_ENV # Submit the recipe mkdir -p __OUTPUT_DIR__/output/srt_slurm_outputs +# Clear the old result lookup ID in case this wrapper is manually rerun and submission fails. rm -f __OUTPUT_DIR__/output/srt_job_id.txt -# C: uv discovers .venv from the current directory even with --no-project. -# Submit from srt-slurm so it cannot discover CloudAI's checkout .venv. +# C: Slurm inherits this working directory. Leave CloudAI's checkout so +# compute-side uv does not discover its login-node .venv, even with --no-project. cd __INSTALL_DIR__/srt-slurm__v2.2.1 if ! __INSTALL_DIR__/srt-slurm__v2.2.1-venv/bin/srtctl apply -f __CLOUDAI_DIR__/conf/experimental/srt_slurm/recipe.yaml -o __OUTPUT_DIR__/output/srt_slurm_outputs -y --set 'benchmark.concurrencies="4"' --no-preflight > __OUTPUT_DIR__/output/run_srtctl.log 2>&1; then cat __OUTPUT_DIR__/output/run_srtctl.log >&2 From 83f7fd45dd274224f6d50888ca33e6c737690689 Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Fri, 30 Oct 2026 12:44:01 -0600 Subject: [PATCH 16/20] module level import + rename load_srt_slurm_results --- .../workloads/srt_slurm/report_generation_strategy.py | 4 ++-- src/cloudai/workloads/srt_slurm/results.py | 2 +- src/cloudai/workloads/srt_slurm/srt_slurm.py | 5 ++--- tests/workloads/srt_slurm/test_results.py | 10 +++++----- 4 files changed, 10 insertions(+), 11 deletions(-) diff --git a/src/cloudai/workloads/srt_slurm/report_generation_strategy.py b/src/cloudai/workloads/srt_slurm/report_generation_strategy.py index 7c1f49eaf..d9dc23f69 100644 --- a/src/cloudai/workloads/srt_slurm/report_generation_strategy.py +++ b/src/cloudai/workloads/srt_slurm/report_generation_strategy.py @@ -20,7 +20,7 @@ from cloudai.core import METRIC_ERROR, MetricValue, ReportGenerationStrategy -from .results import load_results +from .results import load_srt_slurm_results from .srt_slurm import SRTSlurmTestDefinition @@ -51,7 +51,7 @@ class SRTSlurmReportGenerationStrategy(ReportGenerationStrategy): def results(self) -> list[dict[str, float]]: try: - return load_results(cast(SRTSlurmTestDefinition, self.test_run.test), self.test_run) + return load_srt_slurm_results(cast(SRTSlurmTestDefinition, self.test_run.test), self.test_run) except ValueError as error: logging.warning("Unable to read srt-slurm results: %s", error) return [] diff --git a/src/cloudai/workloads/srt_slurm/results.py b/src/cloudai/workloads/srt_slurm/results.py index 1f9b6b9f7..cbaf30214 100644 --- a/src/cloudai/workloads/srt_slurm/results.py +++ b/src/cloudai/workloads/srt_slurm/results.py @@ -192,7 +192,7 @@ def add_agentx_latencies(latency: Any, metrics: dict[str, float]) -> None: metrics[f"{stat}_intvty"] = 1 / value -def load_results(tdef: "SRTSlurmTestDefinition", tr: TestRun) -> list[dict[str, float]]: +def load_srt_slurm_results(tdef: "SRTSlurmTestDefinition", tr: TestRun) -> list[dict[str, float]]: """Read only the current job's outputs, retaining each concurrency.""" logs = tdef.logs_path(tr) if logs is None: diff --git a/src/cloudai/workloads/srt_slurm/srt_slurm.py b/src/cloudai/workloads/srt_slurm/srt_slurm.py index b928296e8..facb69c0f 100644 --- a/src/cloudai/workloads/srt_slurm/srt_slurm.py +++ b/src/cloudai/workloads/srt_slurm/srt_slurm.py @@ -23,6 +23,7 @@ from cloudai.models.workload import CmdArgs, TestDefinition from .installable import SRTSlurmExecutable +from .results import load_srt_slurm_results class SRTSlurmCmdArgs(CmdArgs): @@ -108,10 +109,8 @@ def rollup_path(self, tr: TestRun) -> Path | None: return logs / "benchmark-rollup.json" if logs else None def was_run_successful(self, tr: TestRun) -> JobStatusResult: - from .results import load_results - try: - results = load_results(self, tr) + results = load_srt_slurm_results(self, tr) except ValueError as error: return JobStatusResult(is_successful=False, error_message=str(error)) if not results: diff --git a/tests/workloads/srt_slurm/test_results.py b/tests/workloads/srt_slurm/test_results.py index adae617f1..5b9ab5d0e 100644 --- a/tests/workloads/srt_slurm/test_results.py +++ b/tests/workloads/srt_slurm/test_results.py @@ -27,7 +27,7 @@ JobGPUCounts, agentx_metrics, load_gpu_counts, - load_results, + load_srt_slurm_results, sa_bench_metrics, ) @@ -116,7 +116,7 @@ def test_agentx_requires_job_local_profiles_and_preserves_concurrencies( profile.parent.mkdir(parents=True) profile.write_text('{"metrics": {}}\n') assert tdef.was_run_successful(tr).is_successful - assert len(load_results(tdef, tr)) == 2 + assert len(load_srt_slurm_results(tdef, tr)) == 2 report = SRTSlurmReportGenerationStrategy(slurm_system, tr) assert report.get_metric("tput_per_gpu") == 88086 assert report.get_metric("default") == 176172 @@ -186,7 +186,7 @@ def test_missing_metadata_fails_even_with_results(slurm_system: SlurmSystem, tmp profile.parent.mkdir(parents=True) profile.write_text("{}\n") with pytest.raises(ValueError, match="Missing required srt-slurm job metadata"): - load_results(tdef, tr) + load_srt_slurm_results(tdef, tr) status = tdef.was_run_successful(tr) assert not status.is_successful assert "123.json" in status.error_message @@ -226,7 +226,7 @@ def test_invalid_metadata_is_unsuccessful( invalid = {**invalid, "resources": complete["resources"]} path.write_text(json.dumps(invalid)) with pytest.raises(ValueError): - load_results(tdef, tr) + load_srt_slurm_results(tdef, tr) assert not tdef.was_run_successful(tr).is_successful assert not SRTSlurmReportGenerationStrategy(slurm_system, tr).can_handle_directory() @@ -261,7 +261,7 @@ def test_dse_reporting_uses_each_jobs_metadata( raw = logs / "sa-bench_one" / "results_concurrency_1_gpus_999.json" raw.parent.mkdir(parents=True) raw.write_text(json.dumps({"total_token_throughput": 1600, "output_throughput": 800, "max_concurrency": 1})) - result = load_results(tr.test, tr)[0] + result = load_srt_slurm_results(tr.test, tr)[0] assert result["serving_gpus"] == workers * per_worker assert result["tput_per_gpu"] == 1600 / (workers * per_worker) assert tdef.cmd_args.recipe_overrides == {"roles": {"agg": {"gpus": 4}}} From 3026ec619828c84dffe427340476e18649420985 Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Fri, 30 Oct 2026 12:54:20 -0600 Subject: [PATCH 17/20] add comment to agentx script --- conf/experimental/srt_slurm/agentx_dse/agentx_client.sh | 3 +++ 1 file changed, 3 insertions(+) diff --git a/conf/experimental/srt_slurm/agentx_dse/agentx_client.sh b/conf/experimental/srt_slurm/agentx_dse/agentx_client.sh index 9ed840f29..8c0cd439d 100644 --- a/conf/experimental/srt_slurm/agentx_dse/agentx_client.sh +++ b/conf/experimental/srt_slurm/agentx_dse/agentx_client.sh @@ -15,6 +15,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +# AgentX/DSE smoke-test client: replay two pinned AIPerf fixtures for 60 seconds +# using InferenceX benchmark helpers, avoiding the full dataset download. +# The resulting AgentX artifacts validate CloudAI result extraction and DSE. set -euo pipefail export INFMAX_CONTAINER_WORKSPACE=/infmax-workspace source "$INFMAX_CONTAINER_WORKSPACE/benchmarks/benchmark_lib.sh" From 126af5c88f27468642610121da6ba9ba7cb4a87b Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Fri, 30 Oct 2026 13:25:50 -0600 Subject: [PATCH 18/20] Configure optional pinned models for SRT Slurm Signed-off-by: Jakob Johnson --- .../srt_slurm/agentx_dse/README.md | 17 ++++++++---- .../srt_slurm/agentx_dse/scenario.toml | 3 +++ .../srt_slurm/agentx_dse/srtslurm.yaml | 5 ++-- doc/workloads/srt_slurm.rst | 21 +++++++++++++++ src/cloudai/workloads/srt_slurm/srt_slurm.py | 23 +++++++++++++--- tests/workloads/srt_slurm/test_srt_slurm.py | 26 ++++++++++++++++++- 6 files changed, 84 insertions(+), 11 deletions(-) diff --git a/conf/experimental/srt_slurm/agentx_dse/README.md b/conf/experimental/srt_slurm/agentx_dse/README.md index b3bba31cd..c51f28ea7 100644 --- a/conf/experimental/srt_slurm/agentx_dse/README.md +++ b/conf/experimental/srt_slurm/agentx_dse/README.md @@ -1,6 +1,7 @@ # AgentX DSE smoke test -This example provisions pinned SRT Slurm and InferenceX repositories, then runs +This example provisions pinned SRT Slurm and InferenceX repositories, downloads +Qwen3-0.6B through CloudAI's model installer, then runs SGLang with Qwen3-0.6B on one serving GPU of a four-GPU GB200 node. Grid search tries context lengths 2048 and 4096 and maximizes `p90_intvty`. @@ -22,9 +23,15 @@ All `/shared/...` paths and the `gpu` partition are placeholders: `output_path = "/shared/cloudai/results"` in your system config. If you change the installation root, update `srtctl_root` and both InferenceX paths in `srtslurm.yaml` and `scenario.toml` as well. -- Provide a Qwen3-0.6B model directory at `/shared/models/Qwen3-0.6B` and a - compatible SGLang container at `/shared/images/sglang.sqsh`. The client uses - the model's mounted tokenizer without downloading it from Hugging Face. +- Set `hf_home_path = "/shared/huggingface"` in your system config. The scenario + declares Qwen3-0.6B and its commit in `hf_models`; `cloudai install` downloads that revision + into that cache. In `srtslurm.yaml`, `model_paths.qwen` must point to the + snapshot under the same cache root (`hf_home_path/hub/models--Qwen--Qwen3-0.6B/snapshots/`). + The snapshot revision must match the scenario's model revision. CloudAI does not + synchronize these paths. The cache must be + accessible on the submission and compute nodes. +- Provide a compatible SGLang container at `/shared/images/sglang.sqsh`. SRT + mounts the model snapshot at `/model`; the client uses its tokenizer offline. - Choose your partition in both system and SRT cluster configs. CloudAI uses its partition for the setup allocation; SRT Slurm owns benchmark scheduling. - Configure the Slurm CLI environment and any site-required container settings. @@ -41,7 +48,7 @@ After adapting the paths, run from the directory containing the input files. Pass the path to your adapted copy of the common system template: ```bash -mkdir -p /shared/cloudai/install /shared/cloudai/results +mkdir -p /shared/cloudai/install /shared/cloudai/results /shared/huggingface cloudai install --system-config /path/to/system.toml --test-scenario scenario.toml cloudai run --system-config /path/to/system.toml --test-scenario scenario.toml ``` diff --git a/conf/experimental/srt_slurm/agentx_dse/scenario.toml b/conf/experimental/srt_slurm/agentx_dse/scenario.toml index 0892063ed..9e8aeca40 100644 --- a/conf/experimental/srt_slurm/agentx_dse/scenario.toml +++ b/conf/experimental/srt_slurm/agentx_dse/scenario.toml @@ -21,6 +21,9 @@ id = "agentx" name = "qwen_agentx" description = "Two-step AgentX DSE correctness smoke test" test_template_name = "SRTSlurm" +hf_models = [ + { model_name = "Qwen/Qwen3-0.6B", revision = "c1899de289a04d12100db370d81485cdf75e47ca" }, +] agent = "grid_search" agent_steps = 2 agent_metrics = ["p90_intvty"] diff --git a/conf/experimental/srt_slurm/agentx_dse/srtslurm.yaml b/conf/experimental/srt_slurm/agentx_dse/srtslurm.yaml index 1e4eb5d8d..645f11e1c 100644 --- a/conf/experimental/srt_slurm/agentx_dse/srtslurm.yaml +++ b/conf/experimental/srt_slurm/agentx_dse/srtslurm.yaml @@ -5,9 +5,10 @@ default_partition: gpu gpus_per_node: 4 srtctl_root: /shared/cloudai/install/srt-slurm__098e15aca8504d4747c6f62c7a845e178b7f4b10 -model_paths: {qwen: /shared/models/Qwen3-0.6B} +# Match system.toml hf_home_path and the model revision pinned in scenario.toml. +model_paths: + qwen: /shared/huggingface/hub/models--Qwen--Qwen3-0.6B/snapshots/c1899de289a04d12100db370d81485cdf75e47ca containers: {sglang: /shared/images/sglang.sqsh} default_mounts: /shared/cloudai/install/InferenceX__f8dff19a6898e3ee0dd95232f867151d4d2d6392: /infmax-workspace /shared/cloudai/example: /agentx-test - /shared/models/Qwen3-0.6B: /model diff --git a/doc/workloads/srt_slurm.rst b/doc/workloads/srt_slurm.rst index 104951091..f784c7e8b 100644 --- a/doc/workloads/srt_slurm.rst +++ b/doc/workloads/srt_slurm.rst @@ -71,6 +71,27 @@ tools for the compute-node architecture. A checkout-local bootstrap config prevents interactive setup prompts; existing checkout configs are preserved. Trial submission always uses the explicit ``cluster_config``. +Model installation +~~~~~~~~~~~~~~~~~~ + +Optionally set ``hf_models = ["Qwen/Qwen3-0.6B"]`` on the CloudAI test definition +(outside ``cmd_args``) to download models through the existing ``HFModel`` +installer. The default empty list preserves support for pre-staged weights. +``system.toml``'s ``hf_home_path`` controls the cache root; it defaults to +``/huggingface``. Downloads are stored under +``/hub/models----/snapshots/``. + +Set the corresponding ``model_paths`` alias in ``srtslurm.yaml`` to that snapshot +directory. Its cache root must match ``hf_home_path`` and be accessible on the +submission and compute nodes. CloudAI does not infer models from recipes or +synchronize these paths. To pin a model commit, use +``hf_models = [{model_name = "Qwen/Qwen3-0.6B", revision = ""}]`` and put +the same commit in the SRT snapshot path. String entries remain unpinned. +Pinned models are checked and uninstalled by revision; uninstalling an unpinned +model retains the existing behavior of removing the entire cached model. If the HF +cache is inaccessible locally, the installer assumes models are already available +on execution nodes instead of downloading them. + Scheduling ownership ~~~~~~~~~~~~~~~~~~~~ diff --git a/src/cloudai/workloads/srt_slurm/srt_slurm.py b/src/cloudai/workloads/srt_slurm/srt_slurm.py index facb69c0f..b77eeae80 100644 --- a/src/cloudai/workloads/srt_slurm/srt_slurm.py +++ b/src/cloudai/workloads/srt_slurm/srt_slurm.py @@ -17,9 +17,9 @@ from pathlib import Path from typing import Any, Literal -from pydantic import Field, PrivateAttr, model_validator +from pydantic import BaseModel, Field, PrivateAttr, model_validator -from cloudai.core import Installable, JobStatusResult, TestRun +from cloudai.core import HFModel, Installable, JobStatusResult, TestRun from cloudai.models.workload import CmdArgs, TestDefinition from .installable import SRTSlurmExecutable @@ -46,10 +46,20 @@ def excluded_cmd_args(self) -> list[str]: ] +class SRTSlurmHFModel(BaseModel): + """An optional model prerequisite, independently configured from the SRT recipe.""" + + model_name: str + revision: str | None = None + + class SRTSlurmTestDefinition(TestDefinition): """Run one declarative srt-slurm recipe as one CloudAI trial.""" cmd_args: SRTSlurmCmdArgs + hf_models: list[str | SRTSlurmHFModel] = Field( + default_factory=list, description="Optional Hugging Face models to install." + ) _srtctl: SRTSlurmExecutable | None = PrivateAttr(default=None) @model_validator(mode="after") @@ -82,7 +92,14 @@ def srtctl_path(self) -> Path: @property def installables(self) -> list[Installable]: srtctl = self.srtctl - return [srtctl, *(repo for repo in self.git_repos if repo is not srtctl.git_repo)] + return [ + srtctl, + *(repo for repo in self.git_repos if repo is not srtctl.git_repo), + *( + HFModel(model_name=model) if isinstance(model, str) else HFModel(**model.model_dump()) + for model in self.hf_models + ), + ] @staticmethod def job_id_path(tr: TestRun) -> Path: diff --git a/tests/workloads/srt_slurm/test_srt_slurm.py b/tests/workloads/srt_slurm/test_srt_slurm.py index 4f00b828b..4c672405e 100644 --- a/tests/workloads/srt_slurm/test_srt_slurm.py +++ b/tests/workloads/srt_slurm/test_srt_slurm.py @@ -19,7 +19,7 @@ import pytest import yaml -from cloudai.core import GitRepo, TestRun +from cloudai.core import GitRepo, HFModel, TestRun from cloudai.workloads.srt_slurm import ( SRTSlurmCommandGenStrategy, SRTSlurmTestDefinition, @@ -135,3 +135,27 @@ def test_set_arguments_quote_literal_dotted_keys(): "--set", 'container_mounts."/models/model.v1"="/model"', ] + + +@pytest.mark.parametrize("models", [[], ["Qwen/Qwen3-0.6B"], ["Qwen/Qwen3-0.6B", "example/another-model"]]) +def test_optional_hf_model_installables(tmp_path: Path, make_tdef, models: list[str]): + config = make_tdef(tmp_path).model_dump() + config["hf_models"] = models + tdef = SRTSlurmTestDefinition.model_validate(config) + + assert tdef.installables == [tdef.srtctl, *(HFModel(model_name=name) for name in models)] + assert tdef.hf_models == models + tr = TestRun(name="srt", test=tdef, num_nodes=1, nodes=[]) + assert not any("hf_models" in name for name in tr.param_space) + + +def test_pinned_hf_model_installable(tmp_path: Path, make_tdef): + config = make_tdef(tmp_path).model_dump() + config["hf_models"] = [{"model_name": "Qwen/Qwen3-0.6B", "revision": "a" * 40}, "example/unpinned"] + tdef = SRTSlurmTestDefinition.model_validate(config) + assert tdef.installables == [ + tdef.srtctl, + HFModel("Qwen/Qwen3-0.6B", revision="a" * 40), + HFModel("example/unpinned"), + ] + assert SRTSlurmTestDefinition.model_validate(tdef.model_dump()).installables == tdef.installables From bf7192b7a2ea7463a1ba4d2a55e4de2d285448c0 Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Fri, 30 Oct 2026 13:39:20 -0600 Subject: [PATCH 19/20] make documentation clearer --- .../srt_slurm/agentx_dse/README.md | 15 ++++++++------ .../srt_slurm/agentx_dse/scenario.toml | 1 - .../srt_slurm/agentx_dse/srtslurm.yaml | 3 +-- doc/workloads/srt_slurm.rst | 20 ++++++++++++------- .../srt_slurm/slurm_command_gen_strategy.py | 1 + tests/ref_data/srt-slurm.sbatch | 1 + 6 files changed, 25 insertions(+), 16 deletions(-) diff --git a/conf/experimental/srt_slurm/agentx_dse/README.md b/conf/experimental/srt_slurm/agentx_dse/README.md index c51f28ea7..a54ece809 100644 --- a/conf/experimental/srt_slurm/agentx_dse/README.md +++ b/conf/experimental/srt_slurm/agentx_dse/README.md @@ -20,9 +20,10 @@ All `/shared/...` paths and the `gpu` partition are placeholders: `/shared/cloudai/example`, or update the scenario paths and the corresponding `/agentx-test` mount together. - Set `install_path = "/shared/cloudai/install"` and - `output_path = "/shared/cloudai/results"` in your system config. If you change - the installation root, update `srtctl_root` and both InferenceX paths in - `srtslurm.yaml` and `scenario.toml` as well. + `output_path = "/shared/cloudai/results"` in your system config. CloudAI exports + its absolute installation root as `CLOUDAI_INSTALL_FOLDER`; SRT expands that + variable in the InferenceX mount in `srtslurm.yaml`. SRT derives its own checkout + root, so changing `install_path` does not require editing either repo path. - Set `hf_home_path = "/shared/huggingface"` in your system config. The scenario declares Qwen3-0.6B and its commit in `hf_models`; `cloudai install` downloads that revision into that cache. In `srtslurm.yaml`, `model_paths.qwen` must point to the @@ -33,14 +34,16 @@ All `/shared/...` paths and the `gpu` partition are placeholders: - Provide a compatible SGLang container at `/shared/images/sglang.sqsh`. SRT mounts the model snapshot at `/model`; the client uses its tokenizer offline. - Choose your partition in both system and SRT cluster configs. CloudAI uses - its partition for the setup allocation; SRT Slurm owns benchmark scheduling. + its partition for the setup allocation; SRT generates and submits benchmark + jobs using its recipe and cluster config, and Slurm schedules them. - Configure the Slurm CLI environment and any site-required container settings. Dependencies require network access. uv caches and managed Python downloads use node-local `/tmp`; other application caches still need writable storage. The `srt-slurm` repository entry supplies srtctl, regardless of entry order. -The InferenceX entry installs its pinned submodules. `INFMAX_WORKSPACE` and the `/infmax-workspace` mount -connect that checkout to the custom benchmark client. +The InferenceX entry installs its pinned submodules. The `/infmax-workspace` mount +connects that checkout to the custom benchmark client. If the InferenceX commit +changes, update the checkout basename in `srtslurm.yaml` to match. ## Run diff --git a/conf/experimental/srt_slurm/agentx_dse/scenario.toml b/conf/experimental/srt_slurm/agentx_dse/scenario.toml index 9e8aeca40..5a9829288 100644 --- a/conf/experimental/srt_slurm/agentx_dse/scenario.toml +++ b/conf/experimental/srt_slurm/agentx_dse/scenario.toml @@ -47,7 +47,6 @@ agent_reward_function = "identity" context-length = [2048, 4096] [Tests.extra_env_vars] - INFMAX_WORKSPACE = "/shared/cloudai/install/InferenceX__f8dff19a6898e3ee0dd95232f867151d4d2d6392" # Keep host-side uv work off shared home storage. UV_CACHE_DIR = "/tmp/cloudai-agentx-uv-cache" diff --git a/conf/experimental/srt_slurm/agentx_dse/srtslurm.yaml b/conf/experimental/srt_slurm/agentx_dse/srtslurm.yaml index 645f11e1c..68dc25ad7 100644 --- a/conf/experimental/srt_slurm/agentx_dse/srtslurm.yaml +++ b/conf/experimental/srt_slurm/agentx_dse/srtslurm.yaml @@ -4,11 +4,10 @@ default_partition: gpu gpus_per_node: 4 -srtctl_root: /shared/cloudai/install/srt-slurm__098e15aca8504d4747c6f62c7a845e178b7f4b10 # Match system.toml hf_home_path and the model revision pinned in scenario.toml. model_paths: qwen: /shared/huggingface/hub/models--Qwen--Qwen3-0.6B/snapshots/c1899de289a04d12100db370d81485cdf75e47ca containers: {sglang: /shared/images/sglang.sqsh} default_mounts: - /shared/cloudai/install/InferenceX__f8dff19a6898e3ee0dd95232f867151d4d2d6392: /infmax-workspace + ${CLOUDAI_INSTALL_FOLDER}/InferenceX__f8dff19a6898e3ee0dd95232f867151d4d2d6392: /infmax-workspace /shared/cloudai/example: /agentx-test diff --git a/doc/workloads/srt_slurm.rst b/doc/workloads/srt_slurm.rst index f784c7e8b..2b7405b3f 100644 --- a/doc/workloads/srt_slurm.rst +++ b/doc/workloads/srt_slurm.rst @@ -1,11 +1,11 @@ SRT-Slurm ========= -The ``SRTSlurm`` workload submits an existing recipe with ``srtctl apply`` and -monitors the resulting Slurm job directly. CloudAI is a pass-through layer for -benchmark submission: srt-slurm owns scheduling and resource orchestration. -The pinned srtctl version validates -the recipe, including recipes without a ``schema`` field. Each CloudAI trial +The ``SRTSlurm`` workload delegates benchmark job generation and submission to +srt-slurm through ``srtctl apply``, using its recipe and cluster configuration. +Slurm schedules the job, while CloudAI monitors completion and processes results. +The pinned srtctl version validates the recipe, including recipes without a +``schema`` field. Each CloudAI trial passes selected ``recipe_overrides`` as typed ``--set KEY=VALUE`` arguments; the original YAML remains unchanged. @@ -13,8 +13,8 @@ Configuration ------------- Use ``conf/experimental/srt_slurm/scenario.toml`` as a starting point. Adapt -its cluster configuration to your system, including ``srtctl_root`` pointing -to the installed checkout. Paths are resolved from the CloudAI working directory +its cluster configuration to your system. Omit ``srtctl_root`` to let SRT derive +the installed checkout path. Paths are resolved from the CloudAI working directory and must be accessible on the shared filesystem. .. code-block:: toml @@ -63,6 +63,12 @@ Tools supplied through virtual-environment PATH entries are removed as part of this compatibility cleanup; supply required compute-node tools outside a Python virtual environment. +CloudAI also exports ``CLOUDAI_INSTALL_FOLDER`` as its absolute ``install_path``. +SRT cluster configs can use its existing environment expansion in mount sources, +for example ``${CLOUDAI_INSTALL_FOLDER}/InferenceX__: /infmax-workspace``. +This keeps the scenario independent of the installation root; the checkout name +must still match the configured repository and commit. + Installation runs ``make setup`` in a single-node, 15-minute ``srun`` using ``system.toml``'s ``default_partition``, ``account``, and ``extra_srun_args``. These settings apply only to the installation allocation; benchmark scheduling diff --git a/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py b/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py index 415c87933..88aee9765 100644 --- a/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py +++ b/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py @@ -111,6 +111,7 @@ def write_run_script(self) -> None: "# Configure submission environment", *[f"export {key}={shlex.quote(str(value))}" for key, value in self.final_env_vars.items()], f"export SRTSLURM_CONFIG={shlex.quote(str(cluster_config))}", + f"export CLOUDAI_INSTALL_FOLDER={shlex.quote(str(self.system.install_path.absolute()))}", *(["export AGENTIC_OUTPUT_DIR=/logs/agentic"] if self.tdef.cmd_args.result_format == "agentx" else []), "", "# Prevent login-node Python environments from leaking into compute jobs.", diff --git a/tests/ref_data/srt-slurm.sbatch b/tests/ref_data/srt-slurm.sbatch index 591bbd893..3a84de754 100644 --- a/tests/ref_data/srt-slurm.sbatch +++ b/tests/ref_data/srt-slurm.sbatch @@ -4,6 +4,7 @@ set -euo pipefail # Configure submission environment export TEST_VAR='value with spaces' export SRTSLURM_CONFIG=__CLOUDAI_DIR__/conf/experimental/srt_slurm/srtslurm.yaml +export CLOUDAI_INSTALL_FOLDER=__INSTALL_DIR__ # Prevent login-node Python environments from leaking into compute jobs. # A: uv searches PATH after VIRTUAL_ENV. Remove every virtualenv-backed From 35245c419342dfdbb6b59551247f188842c28828 Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Fri, 30 Oct 2026 17:15:09 -0600 Subject: [PATCH 20/20] Fix SRT example model references and cache mounts Signed-off-by: Jakob Johnson --- conf/experimental/srt_slurm/agentx_dse/README.md | 8 +++++--- conf/experimental/srt_slurm/agentx_dse/srtslurm.yaml | 2 ++ conf/experimental/srt_slurm/recipe.yaml | 2 +- doc/workloads/srt_slurm.rst | 6 ++++++ 4 files changed, 14 insertions(+), 4 deletions(-) diff --git a/conf/experimental/srt_slurm/agentx_dse/README.md b/conf/experimental/srt_slurm/agentx_dse/README.md index a54ece809..80335d161 100644 --- a/conf/experimental/srt_slurm/agentx_dse/README.md +++ b/conf/experimental/srt_slurm/agentx_dse/README.md @@ -28,9 +28,11 @@ All `/shared/...` paths and the `gpu` partition are placeholders: declares Qwen3-0.6B and its commit in `hf_models`; `cloudai install` downloads that revision into that cache. In `srtslurm.yaml`, `model_paths.qwen` must point to the snapshot under the same cache root (`hf_home_path/hub/models--Qwen--Qwen3-0.6B/snapshots/`). - The snapshot revision must match the scenario's model revision. CloudAI does not - synchronize these paths. The cache must be - accessible on the submission and compute nodes. + The snapshot revision must match the scenario's model revision. Keep the + accompanying `blobs` directory mounted at `/blobs`: SRT mounts the snapshot at + `/model`, where its `../../blobs` symlinks need that target. Both YAML paths must + use the same `hf_home_path` cache root. CloudAI does not synchronize these paths; + the cache must be accessible on the submission and compute nodes. - Provide a compatible SGLang container at `/shared/images/sglang.sqsh`. SRT mounts the model snapshot at `/model`; the client uses its tokenizer offline. - Choose your partition in both system and SRT cluster configs. CloudAI uses diff --git a/conf/experimental/srt_slurm/agentx_dse/srtslurm.yaml b/conf/experimental/srt_slurm/agentx_dse/srtslurm.yaml index 68dc25ad7..039c7c973 100644 --- a/conf/experimental/srt_slurm/agentx_dse/srtslurm.yaml +++ b/conf/experimental/srt_slurm/agentx_dse/srtslurm.yaml @@ -9,5 +9,7 @@ model_paths: qwen: /shared/huggingface/hub/models--Qwen--Qwen3-0.6B/snapshots/c1899de289a04d12100db370d81485cdf75e47ca containers: {sglang: /shared/images/sglang.sqsh} default_mounts: + # Snapshot links under /model resolve ../../blobs to /blobs inside the container. + /shared/huggingface/hub/models--Qwen--Qwen3-0.6B/blobs: /blobs ${CLOUDAI_INSTALL_FOLDER}/InferenceX__f8dff19a6898e3ee0dd95232f867151d4d2d6392: /infmax-workspace /shared/cloudai/example: /agentx-test diff --git a/conf/experimental/srt_slurm/recipe.yaml b/conf/experimental/srt_slurm/recipe.yaml index e79286e31..34f00ea49 100644 --- a/conf/experimental/srt_slurm/recipe.yaml +++ b/conf/experimental/srt_slurm/recipe.yaml @@ -1,7 +1,7 @@ schema: 2 name: sglang-example model: - path: Qwen/Qwen3-0.6B + path: hf:Qwen/Qwen3-0.6B container: docker.io/lmsysorg/sglang:v0.5.9 resources: gpu_type: h100 diff --git a/doc/workloads/srt_slurm.rst b/doc/workloads/srt_slurm.rst index 2b7405b3f..f431fd823 100644 --- a/doc/workloads/srt_slurm.rst +++ b/doc/workloads/srt_slurm.rst @@ -87,6 +87,12 @@ installer. The default empty list preserves support for pre-staged weights. ``/huggingface``. Downloads are stored under ``/hub/models----/snapshots/``. +HF cache snapshots contain relative links to ``../../blobs``. When SRT mounts a +snapshot at ``/model``, also mount that model cache's ``blobs`` directory at +``/blobs`` through ``srtslurm.yaml``'s ``default_mounts``. The snapshot and blob +mount sources must use the same ``hf_home_path`` cache root. See the AgentX DSE +example for a complete configuration. + Set the corresponding ``model_paths`` alias in ``srtslurm.yaml`` to that snapshot directory. Its cache root must match ``hf_home_path`` and be accessible on the submission and compute nodes. CloudAI does not infer models from recipes or