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..80335d161 --- /dev/null +++ b/conf/experimental/srt_slurm/agentx_dse/README.md @@ -0,0 +1,63 @@ +# AgentX DSE smoke test + +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`. + +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: + +- 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. 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 + 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. 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 + 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. 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 + +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 /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 +``` + +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..8c0cd439d --- /dev/null +++ b/conf/experimental/srt_slurm/agentx_dse/agentx_client.sh @@ -0,0 +1,52 @@ +#!/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. + +# 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" +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..5a9829288 --- /dev/null +++ b/conf/experimental/srt_slurm/agentx_dse/scenario.toml @@ -0,0 +1,53 @@ +# 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" +hf_models = [ + { model_name = "Qwen/Qwen3-0.6B", revision = "c1899de289a04d12100db370d81485cdf75e47ca" }, +] +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] + + # 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..039c7c973 --- /dev/null +++ b/conf/experimental/srt_slurm/agentx_dse/srtslurm.yaml @@ -0,0 +1,15 @@ +# 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 +# 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: + # 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 new file mode 100644 index 000000000..34f00ea49 --- /dev/null +++ b/conf/experimental/srt_slurm/recipe.yaml @@ -0,0 +1,23 @@ +schema: 2 +name: sglang-example +model: + path: hf:Qwen/Qwen3-0.6B + container: docker.io/lmsysorg/sglang:v0.5.9 +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..b0e6afe7c --- /dev/null +++ b/conf/experimental/srt_slurm/scenario.toml @@ -0,0 +1,41 @@ +# 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 = "srt_slurm" + +[[Tests]] +id = "srt.sglang" +# Benchmark node allocation is defined by recipe.yaml, not this value. +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"] + + [[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" + # Owns benchmark scheduling defaults; CloudAI allocation settings are not merged. + 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..16927cf15 --- /dev/null +++ b/conf/experimental/srt_slurm/srtslurm.yaml @@ -0,0 +1,6 @@ +# Benchmark scheduling defaults belong here and in the recipe. +# CloudAI system/scenario allocation settings are not forwarded to srt-slurm. +# 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 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..f431fd823 --- /dev/null +++ b/doc/workloads/srt_slurm.rst @@ -0,0 +1,200 @@ +SRT-Slurm +========= + +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. + +Configuration +------------- + +Use ``conf/experimental/srt_slurm/scenario.toml`` as a starting point. Adapt +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 + + name = "srt_sglang" + description = "SGLang with srt-slurm" + test_template_name = "SRTSlurm" + agent_metrics = ["tput_per_gpu"] + + [[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 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 +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 +``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 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. + +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 +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``. + +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/``. + +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 +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 +~~~~~~~~~~~~~~~~~~~~ + +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. + +Results +------- + +``result_format = "sa-bench"`` reads job-local ``sa-bench_*/results_*.json`` +files, falling back to ``benchmark-rollup.json`` when raw results are absent. +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. 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 +``conc_/aiperf_artifacts/profile_export.jsonl``. CloudAI reads the existing +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 +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..1119c007a --- /dev/null +++ b/src/cloudai/workloads/srt_slurm/__init__.py @@ -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. + +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..ca8d5a6ce --- /dev/null +++ b/src/cloudai/workloads/srt_slurm/installable.py @@ -0,0 +1,115 @@ +# 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 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.""" + + _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"] + 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: + # 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) + + 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, + 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}") + 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..d9dc23f69 --- /dev/null +++ b/src/cloudai/workloads/srt_slurm/report_generation_strategy.py @@ -0,0 +1,80 @@ +# 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 csv +import logging +from typing import ClassVar, cast + +from cloudai.core import METRIC_ERROR, MetricValue, ReportGenerationStrategy + +from .results import load_srt_slurm_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", + "total_token_throughput", + "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]]: + try: + 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 [] + + 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..cbaf30214 --- /dev/null +++ b/src/cloudai/workloads/srt_slurm/results.py @@ -0,0 +1,222 @@ +# 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 +import math +from dataclasses import dataclass +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} + + +@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(): + if old in metrics: + metrics[new] = metrics[old] + if "throughput_toks" not in metrics: + return {} + metrics["output_throughput"] = metrics["throughput_toks"] + # 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 + + +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], gpus: JobGPUCounts) -> dict[str, float]: + request = data.get("request_metrics") + if not isinstance(request, dict): + return {} + throughput = request.get("throughput") + if not isinstance(throughput, dict): + return {} + metrics = numeric_fields(data) + if "conc" in metrics: + metrics["concurrency"] = metrics["conc"] + # 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"), + ): + 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) + 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_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: + 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), gpus) + 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), 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, 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 new file mode 100644 index 000000000..88aee9765 --- /dev/null +++ b/src/cloudai/workloads/srt_slurm/slurm_command_gen_strategy.py @@ -0,0 +1,207 @@ +# 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 +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() + return { + "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": output / "run_srtctl.sh", + "job_id": self.tdef.job_id_path(self.test_run).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.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) + apply_command = shlex.join( + [ + str(self.tdef.srtctl_path), + "apply", + "-f", + str(recipe), + "-o", + str(paths["nested_outputs"]), + "-y", + *overrides, + *self._extra_cmd_arg_parts(), + ] + ) + script = "\n".join( + [ + "#!/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))}", + 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.", + "# 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", + 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: 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", + " exit 1", + "fi", + f"cat {shlex.quote(str(paths['run_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' " + 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"', + "", + ] + ) + 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(): + args.append(f"--{name.replace('_', '-')}") + if value: + args.append(value) + return args + + 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"]) + 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..b77eeae80 --- /dev/null +++ b/src/cloudai/workloads/srt_slurm/srt_slurm.py @@ -0,0 +1,139 @@ +# 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. + +from pathlib import Path +from typing import Any, Literal + +from pydantic import BaseModel, Field, PrivateAttr, model_validator + +from cloudai.core import HFModel, Installable, JobStatusResult, TestRun +from cloudai.models.workload import CmdArgs, TestDefinition + +from .installable import SRTSlurmExecutable +from .results import load_srt_slurm_results + + +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" + # 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", + "cmd_args.recipe_override_mode", + ] + + +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") + 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: + 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 + 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]: + srtctl = self.srtctl + 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: + """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: + return None + if not job_id.isdigit(): + return None + 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: + try: + results = load_srt_slurm_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: + output = self.job_output_path(tr) + return output / "logs" if output else None diff --git a/tests/ref_data/srt-slurm.sbatch b/tests/ref_data/srt-slurm.sbatch new file mode 100644 index 000000000..3a84de754 --- /dev/null +++ b/tests/ref_data/srt-slurm.sbatch @@ -0,0 +1,45 @@ +#!/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 +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 +# 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 +# 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: 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 + 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 3f6b38b50..a6852cf36 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,29 @@ 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/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"), + 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", 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..9b260c3eb 100644 --- a/tests/test_test_scenario.py +++ b/tests/test_test_scenario.py @@ -76,6 +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.triton_inference import TritonInferenceReportGenerationStrategy, TritonInferenceTestDefinition from cloudai.workloads.ucc_test import UCCTestDefinition, UCCTestReportGenerationStrategy @@ -757,7 +761,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 +802,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/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_installable.py b/tests/workloads/srt_slurm/test_installable.py new file mode 100644 index 000000000..2e3beb1ff --- /dev/null +++ b/tests/workloads/srt_slurm/test_installable.py @@ -0,0 +1,114 @@ +# 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 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"), + ) 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(): + 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..4bdec5018 --- /dev/null +++ b/tests/workloads/srt_slurm/test_report_generation_strategy.py @@ -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. + +import json +from pathlib import Path + +from cloudai.core import TestRun +from cloudai.systems.slurm import SlurmSystem +from cloudai.workloads.srt_slurm import SRTSlurmReportGenerationStrategy + + +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) + 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( + 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..5b9ab5d0e --- /dev/null +++ b/tests/workloads/srt_slurm/test_results.py @@ -0,0 +1,307 @@ +# 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 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, SRTSlurmTestDefinition +from cloudai.workloads.srt_slurm.results import ( + JobGPUCounts, + agentx_metrics, + load_gpu_counts, + load_srt_slurm_results, + sa_bench_metrics, +) + + +def test_fixed_sequence_uses_serving_gpu_counts(tmp_path: Path): + 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}, + gpus, + ) + 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 + + +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}}]], +) +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") + write_job_metadata(tdef, tr) + 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, 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) + 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" + 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": 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, 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" + profile.parent.mkdir(parents=True) + profile.write_text('{"metrics": {}}\n') + assert tdef.was_run_successful(tr).is_successful + 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 + 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} + + +@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, 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") + 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, 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") + 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_srt_slurm_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, 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") + 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_srt_slurm_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, 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") + 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, 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( + {"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_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}}} + + +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 new file mode 100644 index 000000000..1048a8b99 --- /dev/null +++ b/tests/workloads/srt_slurm/test_slurm_command_gen_strategy.py @@ -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. + +from pathlib import Path + +import pytest +import yaml + +from cloudai.core import TestRun +from cloudai.systems.slurm import SlurmSystem +from cloudai.workloads.srt_slurm import SRTSlurmCommandGenStrategy + + +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"): + SRTSlurmCommandGenStrategy(slurm_system, tr).gen_srun_command() + + +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" + 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 / "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() + 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 new file mode 100644 index 000000000..4c672405e --- /dev/null +++ b/tests/workloads/srt_slurm/test_srt_slurm.py @@ -0,0 +1,161 @@ +# 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. + +from pathlib import Path + +import pytest +import yaml + +from cloudai.core import GitRepo, HFModel, TestRun +from cloudai.workloads.srt_slurm import ( + SRTSlurmCommandGenStrategy, + SRTSlurmTestDefinition, +) + + +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", + "cmd_args.recipe_override_mode", + ] + assert "cmd_args.recipe_overrides" not in tdef.dse_excluded_args + + +@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 srt_repo + assert tdef.srtctl.git_repo.commit == "updated-revision" + assert "updated-revision" in tdef.srtctl.venv_name + assert tdef.installables == [tdef.srtctl, other_repo] + assert isinstance(tdef.installables[1], GitRepo) + + +@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: + config.pop("git_repos") + else: + config["git_repos"] = git_repos + + 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)"]) +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"', + ] + + +@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