NVIDIA has released srt-slurm, a framework designed to convert declarative YAML configurations into reproducible SLURM benchmark workflows for distributed large language model serving. The tool, accessed via the srtctl command-line interface, allows developers to define cluster setups, run dry-runs of built-in and custom recipes, and generate parameter sweeps without needing immediate access to a physical GPU cluster. The tutorial demonstrates these capabilities within Google Colab, using the environment as a development workspace to validate production-grade recipes before deployment to actual hardware.
Setup and environment
The process begins by cloning the NVIDIA repository and installing the package in editable mode. A helper script handles shell command execution and output formatting. Once the source directory is exposed to the Python runtime, the code switches to the repository root and verifies the srtctl interface is functioning correctly.
import os, sys, subprocess, textwrap, json, shutil, importlib
from pathlib import Path
def run(cmd, check=True, quiet=False):
"""Run a shell command, stream output."""
print(f"\n$ {cmd}")
r = subprocess.run(cmd, shell=True, text=True, capture_output=True)
out = (r.stdout or "") + (r.stderr or "")
if not quiet:
print(out[-6000:])
if check and r.returncode != 0:
raise RuntimeError(f"Command failed ({r.returncode}): {cmd}")
return out
def section(title):
print("\n" + "═"*78 + f"\n {title}\n" + "═"*78)
section("1. Install srt-slurm")
REPO = Path("/content/srt-slurm") if Path("/content").exists() else Path.cwd()/"srt-slurm"
if not REPO.exists():
run(f"git clone --depth 1 https://github.com/NVIDIA/srt-slurm.git {REPO}", quiet=True)
run(f"{sys.executable} -m pip install -q -e {REPO}", quiet=True)
sys.path.insert(0, str(REPO / "src"))
importlib.invalidate_caches()
os.chdir(REPO)
run("srtctl --help")
Repository structure
Inspecting the codebase reveals a modular design. The src/srtctl directory contains the CLI tools for submitting, monitoring, and sweeping jobs. Core logic lives in core, handling schemas, validation, and SLURM script generation. Backend adapters for engines like SGLang and vLLM sit in backends, while Jinja2 templates in templates generate the actual submission scripts. Pre-built benchmark recipes for platforms such as GB200 and H100 are stored in recipes.
section("2. Repository architecture")
print(textwrap.dedent("""
src/srtctl/
cli/ submit.py (apply/dry-run/preflight/monitor), do_sweep, interactive
core/ schema.py (typed config), sweep.py, slurm.py (sbatch gen),
validation.py, health.py, topology.py, fingerprint.py
backends/ sglang.py, trtllm.py, vllm.py, mocker.py ← engine adapters
frontends/ Dynamo / router frontends
templates/ Jinja2 → sbatch + orchestrator scripts
recipes/ ready-made benchmarks per platform (gb200-fp4, h100, b200-fp8,
qwen3-32b, dsv4-pro, mocker smoke tests, ...)
analysis/ srtlog (log parsers) + Streamlit dashboard (Pareto, latency...)
docs/ sweeps.md, profiling.md, analyzing.md, config-reference.md
"""))
for d in ["recipes", "docs"]:
print(f"{d}/ →", ", ".join(sorted(p.name for p in (REPO/d).iterdir()))[:300])
Cluster configuration
Developers create a local srtslurm.yaml file to define simulated cluster defaults. This single file sets the account, partition, time limits, and GPU counts per node. It also lists container aliases for different engine versions and specifies the path to the model directory. This configuration allows the tool to resolve recipe references in Colab without a live SLURM connection.
section("3. Cluster configuration (srtslurm.yaml)")
(REPO/"srtslurm.yaml").write_text(textwrap.dedent("""\
cluster: "colab-demo"
default_account: "demo-account"
default_partition: "gpu"
default_time_limit: "01:00:00"
gpus_per_node: 4
use_gpus_per_node_directive: true
use_segment_sbatch_directive: true
containers:
dynamo-sglang: "/containers/dynamo-sglang.sqsh"
lmsysorg+sglang+v0.5.5.post2.sqsh: "/containers/sglang-v0.5.5.sqsh"
model_paths:
deepseek-r1: "/models/DeepSeek-R1"
"""))
print((REPO/"srtslurm.yaml").read_text())
Dry runs and disaggregated deployment
Running a dry-run on a built-in recipe, such as a mocker smoke test, shows how srtctl validates inputs and generates SLURM submission artifacts without executing real workloads. A custom recipe is then defined to separate prefill and decode workloads. This setup assigns prefill tasks to one node and decode tasks to two others, using the SGLang engine in fp8 precision. The configuration maps specific environment variables and model parameters to each stage.
section("4. Dry-run: mocker smoke test → generated sbatch script")
run("srtctl dry-run -f recipes/mocker/agg.yaml", check=False)
section("5. Custom disaggregated recipe (prefill/decode split)")
(REPO/"my-disagg.yaml").write_text(textwrap.dedent("""\
name: "colab-disagg-demo"
model:
path: "deepseek-r1"
container: "lmsysorg+sglang+v0.5.5.post2.sqsh"
precision: "fp8"
resources:
gpu_type: "gb200"
gpus_per_node: 4
prefill_nodes: 1
decode_nodes: 2
prefill_workers: 1
decode_workers: 2
backend:
prefill_environment: { PYTHONUNBUFFERED: "1" }
decode_environment: { PYTHONUNBUFFERED: "1" }
sglang_config:
prefill:
served-model-name: "deepseek-ai/DeepSeek-R1"
model-path: "/model/"
trust-remote-code: true
kv-cache-dtype: "fp8_e4m3"
tensor-parallel-size: 4
disaggregation-mode: "prefill"
decode:
served-model-name: "deepseek-ai/DeepSeek-R1"
model-path: "/model/"
trust-remote-code: true
kv-cache-dtype: "fp8_e4m3"
tensor-parallel-size: 4
disaggregation-mode: "decode"
benchmark:
type: "sa-bench"
isl: 1024
osl: 1024
concurrencies: [64, 128, 256]
req_rate: "inf"
"""))
run("srtctl dry-run -f my-disagg.yaml", check=False)
Parameter sweeps and API usage
The framework handles grid searches by expanding a single configuration file into multiple individual job definitions. Running the sweep command generates a directory of config files on disk. Developers can also use the typed Python API to load configurations directly, inspecting model paths, precision settings, and benchmark parameters. The API exposes enums for benchmarks, precisions, and GPU types, allowing code to validate inputs before generation.
section("6. Parameter sweep (grid search) — dry-run + expansion on disk")
run("srtctl dry-run -f examples/example-sweep.yaml", check=False)
sweep_dirs = sorted((REPO/"dry-runs").glob("example-sweep_sweep_*"))
if sweep_dirs:
latest = sweep_dirs[-1]
print("Per-job configs generated by the sweep expander:")
for p in sorted(latest.rglob("config.yaml")):
print(" ", p.relative_to(REPO))section("7. Programmatic use of srtctl's Python API")
import yaml
from srtctl.core.config import load_config
from srtctl.core.sweep import generate_sweep_configs, expand_template
from srtctl.core.schema import BenchmarkType, Precision, GpuType
cfg = load_config("my-disagg.yaml")
print(f"Loaded : {cfg.name}")
print(f"Model : {cfg.model.path} ({cfg.model.precision}) on {cfg.resources.gpu_type}")
print(f"Layout : {cfg.resources.prefill_nodes}P + {cfg.resources.decode_nodes}D nodes, "
f"{cfg.resources.gpus_per_node} GPUs/node")
print(f"Bench : {cfg.benchmark.type} isl={cfg.benchmark.isl} osl={cfg.benchmark.osl} "
f"concurrencies={cfg.benchmark.concurrencies}")
print(f"Enums : benchmarks={[b.value for b in BenchmarkType]}")
print(f" precisions={[p.value for p in Precision]}, gpus={[g.value for g in GpuType]}")
raw_sweep = yaml.safe_load(Path("examples/example-sweep.yaml").read_text())
jobs = generate_sweep_configs(raw_sweep)
print(f"\nSweep expands to {len(jobs)} jobs:")
for job_cfg, params in jobs:
pf = job_cfg["backend"]["sglang_config"]["prefill"]
print(f" {params} → chunked-prefill-size={pf['chunked-prefill-size']}, "
f"max-total-tokens={pf
Source Read original →Related articles
