From 55f0872a1952350975e233f2cb8bbd8736f7eb2a Mon Sep 17 00:00:00 2001 From: Anish Shanbhag Date: Thu, 27 Aug 2026 17:38:29 -0700 Subject: [PATCH 1/3] feat: run apply in current Slurm allocation Signed-off-by: Anish Shanbhag --- README.md | 3 ++ docs/cli.md | 10 ++++ src/srtctl/cli/submit.py | 103 +++++++++++++++++++++++++++++++++++++++ tests/test_submit_cli.py | 66 +++++++++++++++++++++++++ 4 files changed, 182 insertions(+) diff --git a/README.md b/README.md index e3d91ecf0..a11dc8068 100644 --- a/README.md +++ b/README.md @@ -35,6 +35,9 @@ srtctl apply -f config.yaml # Deploy an inference endpoint without running a benchmark srtctl apply -f config.yaml --serve-only +# Run a single recipe in the current SLURM allocation +srtctl apply -f config.yaml --current-allocation + # Submit with custom setup script srtctl apply -f config.yaml --setup-script custom-setup.sh diff --git a/docs/cli.md b/docs/cli.md index 06142c945..6c0d45e3f 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -254,6 +254,7 @@ srtctl apply -f [options] | `--tags` | Comma-separated tags for the run | | `--serve-only` | Deploy the endpoint without running a benchmark; serve until cancellation | | `--bash` | Print a direct single-node lifecycle script to stdout without submitting | +| `--current-allocation` | Run a single-job recipe in the current Slurm allocation without submitting another job | | `-y, --yes` | Skip confirmation prompts | **Examples:** @@ -265,6 +266,9 @@ srtctl apply -f recipes/gb200-fp8/sglang-1p4d.yaml # Serve the same recipe without running its configured benchmark srtctl apply -f recipes/gb200-fp8/sglang-1p4d.yaml --serve-only +# Run in an allocation obtained with salloc or sbatch +srtctl apply -f recipes/gb200-fp8/sglang-1p4d.yaml --current-allocation + # Submit sweep (auto-detected from sweep: section) srtctl apply -f configs/my-sweep.yaml @@ -291,6 +295,12 @@ the frontend URL in the sweep log, and keeps the service running until the job i time limit. It ignores the recipe's configured benchmark for that submission. Use `scancel ` to stop the service; srtctl then cleans up the processes it started. +`--current-allocation` requires `SLURM_JOB_ID` in the environment. It skips `sbatch` and runs the normal +orchestrator in the foreground, so its `srun` commands become job steps in that allocation. The initial implementation +accepts one regular recipe at a time; directories, sweeps, override configs, `--bash`, `--mock`, `--json`, and `--tags` +are not supported. Allocation settings such as account, partition, and time limit have no effect because the allocation +already exists. + `--bash` renders a small direct-host launcher; it is not an sbatch script. The launcher owns a Docker serving container and runs the serving lifecycle inside the selected SGLang image. It currently supports a one-node SGLang backend with the Dynamo frontend, `frontend.enable_multiple_frontends: false`, and one diff --git a/src/srtctl/cli/submit.py b/src/srtctl/cli/submit.py index a1008cff5..030f7de82 100755 --- a/src/srtctl/cli/submit.py +++ b/src/srtctl/cli/submit.py @@ -9,6 +9,7 @@ Usage: srtctl apply -f config.yaml # Submit job + srtctl apply -f config.yaml --current-allocation # Run in the current allocation srtctl apply -f config.yaml -o /path/to/logs # Submit with custom output dir srtctl dry-run -f sweep.yaml --sweep # Dry run sweep """ @@ -56,6 +57,7 @@ ) from srtctl.core.lockfile import load_lockfile_fingerprints from srtctl.core.schema import SrtConfig, installs_dynamo +from srtctl.core.slurm import get_slurm_job_id from srtctl.core.status import create_job_record from srtctl.core.validation import preflight_config_variants from srtctl.ports import MOONCAKE_MASTER_PORT @@ -926,6 +928,72 @@ def submit_single( ) +def run_in_current_allocation( + config_path: Path, + *, + setup_script: str | None = None, + output_dir: Path | None = None, + enforce_preflight: bool = True, + serve_only: bool = False, +) -> int: + """Run one config as job steps in the current SLURM allocation.""" + job_id = get_slurm_job_id() + if not job_id: + raise ValueError("--current-allocation requires SLURM_JOB_ID to be set") + + config = load_config(config_path) + if enforce_preflight: + with open(config_path) as f: + raw_config = yaml.safe_load(f) + _assert_preflight_passed(raw_config, label=str(config_path)) + + configured_source = os.environ.get("SRTCTL_SOURCE_DIR") or get_srtslurm_setting("srtctl_root") + srtctl_source = ( + Path(configured_source).resolve() if configured_source else Path(__file__).parent.parent.parent.parent.resolve() + ) + validate_setup(srtctl_source) + + if output_dir: + output_base = output_dir.resolve() + else: + configured_output = get_srtslurm_setting("output_dir") + output_base = ( + Path(os.path.expandvars(configured_output)).resolve() if configured_output else srtctl_source / "outputs" + ) + + job_output_dir = output_base / job_id + job_output_dir.mkdir(parents=True, exist_ok=True) + runtime_config_path = job_output_dir / "config.yaml" + if config_path.resolve() != runtime_config_path.resolve(): + shutil.copy2(config_path, runtime_config_path) + + env = os.environ.copy() + env["SRTCTL_OUTPUT_DIR"] = str(job_output_dir) + env["SRTCTL_SOURCE_DIR"] = str(srtctl_source) + dynamo_environment = config.dynamo.get_wheel_environment() + env.update(dynamo_environment) + env.update({key: str(value) for key, value in config.environment.items()}) + if setup_script: + env["SRTCTL_SETUP_SCRIPT"] = setup_script + + if dynamo_environment or env.get("SRTCTL_PREFETCH_AI_DYNAMO") == "1": + prefetch_script = srtctl_source / "src" / "srtctl" / "runtime_scripts" / "dynamo_wheels.py" + if not prefetch_script.exists(): + raise FileNotFoundError(f"Dynamo wheel prefetch script not found: {prefetch_script}") + subprocess.run([sys.executable, str(prefetch_script), "prefetch"], env=env, check=True) + + command = [sys.executable, "-m", "srtctl.cli.do_sweep", str(runtime_config_path)] + if serve_only: + command.append("--serve-only") + + console.print(f"[bold cyan]Running in SLURM allocation {job_id}:[/] {config.name}") + console.print(f"[dim]Outputs:[/] {job_output_dir}") + _print_running_summary(config, console, serve_only=serve_only) + + result = subprocess.run(command, env=env, check=False) + return result.returncode + + def is_sweep_config(config_path: Path) -> bool: """Check if config file is a sweep config by looking for 'sweep' section.""" try: @@ -1456,6 +1524,7 @@ def main(): epilog="""Examples: srtctl # Interactive mode srtctl apply -f config.yaml # Submit job + srtctl apply -f config.yaml --current-allocation # Run in the current SLURM allocation srtctl apply -f config.yaml --serve-only # Serve until cancelled; do not benchmark srtctl apply -f config.yaml --bash # Print a direct single-node Bash lifecycle script srtctl apply -f ./configs/ # Submit all YAMLs in directory @@ -1501,6 +1570,11 @@ def add_common_args(p): dest="bash_output", help="Print a direct single-node Bash lifecycle script to stdout and exit without submitting.", ) + apply_parser.add_argument( + "--current-allocation", + action="store_true", + help="Run as job steps in the current SLURM allocation instead of submitting with sbatch.", + ) apply_parser.add_argument( "--json", action="store_true", @@ -1597,6 +1671,7 @@ def add_common_args(p): json_mode = bool(getattr(args, "json_output", False)) mock_mode = bool(getattr(args, "mock_mode", False)) bash_mode = bool(getattr(args, "bash_output", False)) + current_allocation = bool(getattr(args, "current_allocation", False)) serve_only = bool(getattr(args, "serve_only", False)) if bash_mode and json_mode: parser.error("--bash cannot be combined with --json") @@ -1610,6 +1685,16 @@ def add_common_args(p): parser.error("--serve-only cannot be combined with --mock") if serve_only and getattr(args, "sweep", False): parser.error("--serve-only does not support sweeps") + if current_allocation and bash_mode: + parser.error("--current-allocation cannot be combined with --bash") + if current_allocation and mock_mode: + parser.error("--current-allocation cannot be combined with --mock") + if current_allocation and json_mode: + parser.error("--current-allocation cannot be combined with --json") + if current_allocation and getattr(args, "sweep", False): + parser.error("--current-allocation supports single-job configs only") + if current_allocation and getattr(args, "tags", None): + parser.error("--current-allocation cannot be combined with --tags") # Always rebind the module console on each invocation so json-mode prose # goes to stderr and non-json prose returns to stdout. Save the original @@ -1778,6 +1863,24 @@ def restore_console() -> None: no_preflight = getattr(args, "no_preflight", False) enforce_preflight = not (mock_mode or is_dry_run or no_preflight) + if current_allocation: + if effective_config_path.is_dir(): + raise ValueError("--current-allocation expects a single config file, not a directory") + if selector or is_override_config(effective_config_path): + raise ValueError("--current-allocation does not support override configs") + if is_sweep_config(effective_config_path): + raise ValueError("--current-allocation supports single-job configs only") + + exit_code = run_in_current_allocation( + effective_config_path, + setup_script=setup_script, + output_dir=output_dir, + enforce_preflight=enforce_preflight, + serve_only=serve_only, + ) + restore_console() + sys.exit(exit_code) + # Handle directory input if effective_config_path.is_dir(): if serve_only: diff --git a/tests/test_submit_cli.py b/tests/test_submit_cli.py index 6463cc768..05f0034f7 100644 --- a/tests/test_submit_cli.py +++ b/tests/test_submit_cli.py @@ -6,6 +6,7 @@ import sys from io import StringIO from pathlib import Path +from subprocess import CompletedProcess import pytest import yaml @@ -117,6 +118,71 @@ def fail_subprocess_run(*_args, **_kwargs): assert "srtctl.cli.run_benchmark" not in output +def test_apply_current_allocation_runs_existing_orchestrator(monkeypatch, tmp_path: Path, capsys) -> None: + config_path = tmp_path / "config.yaml" + config_path.write_text(yaml.safe_dump(MINIMAL_DRY_RUN_CONFIG)) + output_dir = tmp_path / "outputs" + calls = [] + + def fake_run(command, **kwargs): + calls.append((command, kwargs)) + return CompletedProcess(command, 0) + + monkeypatch.setenv("SLURM_JOB_ID", "12345") + monkeypatch.setenv("SLURM_NODELIST", "gpu-01") + monkeypatch.setattr(submit_cli, "validate_setup", lambda _source: None) + monkeypatch.setattr(submit_cli, "_assert_preflight_passed", lambda *_args, **_kwargs: None) + monkeypatch.setattr(submit_cli.subprocess, "run", fake_run) + monkeypatch.setattr( + sys, + "argv", + [ + "srtctl", + "apply", + "-f", + str(config_path), + "-o", + str(output_dir), + "--current-allocation", + "--setup-script", + "custom.sh", + ], + ) + + with pytest.raises(SystemExit) as exc_info: + submit_cli.main() + + assert exc_info.value.code == 0 + assert len(calls) == 1 + command, kwargs = calls[0] + runtime_config = output_dir / "12345" / "config.yaml" + assert command == [sys.executable, "-m", "srtctl.cli.do_sweep", str(runtime_config)] + assert kwargs["check"] is False + assert kwargs["env"]["SLURM_JOB_ID"] == "12345" + assert kwargs["env"]["SRTCTL_OUTPUT_DIR"] == str(output_dir / "12345") + assert kwargs["env"]["SRTCTL_SETUP_SCRIPT"] == "custom.sh" + assert yaml.safe_load(runtime_config.read_text()) == MINIMAL_DRY_RUN_CONFIG + assert "Running in SLURM allocation 12345" in capsys.readouterr().out + + +def test_apply_current_allocation_requires_slurm_job(monkeypatch, tmp_path: Path, capsys) -> None: + config_path = tmp_path / "config.yaml" + config_path.write_text(yaml.safe_dump(MINIMAL_DRY_RUN_CONFIG)) + monkeypatch.delenv("SLURM_JOB_ID", raising=False) + monkeypatch.delenv("SLURM_JOBID", raising=False) + monkeypatch.setattr( + sys, + "argv", + ["srtctl", "apply", "-f", str(config_path), "--current-allocation"], + ) + + with pytest.raises(SystemExit) as exc_info: + submit_cli.main() + + assert exc_info.value.code == 1 + assert "--current-allocation requires SLURM_JOB_ID to be set" in capsys.readouterr().out + + def test_load_config_rejects_empty_yaml(tmp_path: Path) -> None: path = tmp_path / "empty.yaml" path.write_text("") From 0120dd7a4d0206fe135d5e893e0e041c6f0190de Mon Sep 17 00:00:00 2001 From: Anish Shanbhag Date: Thu, 27 Aug 2026 17:46:16 -0700 Subject: [PATCH 2/3] test: patch refactored health helper Signed-off-by: Anish Shanbhag --- tests/test_benchmarks.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_benchmarks.py b/tests/test_benchmarks.py index 3681d01f1..af29a6e67 100644 --- a/tests/test_benchmarks.py +++ b/tests/test_benchmarks.py @@ -1539,7 +1539,7 @@ def test_eval_only_health_check_fails(self): orch = self._make_orchestrator() stop = threading.Event() with patch.dict(os.environ, {"EVAL_ONLY": "true"}, clear=False): - with patch("srtctl.core.health.wait_for_model", return_value=False): + with patch("srtctl.cli.mixins.benchmark_stage.wait_for_model", return_value=False): result = orch._run_post_eval(stop) assert result == 1 @@ -1590,7 +1590,7 @@ def test_eval_only_successful(self): mock_proc.returncode = 0 with patch.dict(os.environ, {"EVAL_ONLY": "true"}, clear=False): - with patch("srtctl.core.health.wait_for_model", return_value=True): + with patch("srtctl.cli.mixins.benchmark_stage.wait_for_model", return_value=True): with patch("srtctl.cli.do_sweep.start_srun_process", return_value=mock_proc): result = orch._run_post_eval(stop) assert result == 0 From 64e4310bf3763893b8257e9eff9d45ba6f34cd46 Mon Sep 17 00:00:00 2001 From: Anish Shanbhag Date: Thu, 27 Aug 2026 18:09:21 -0700 Subject: [PATCH 3/3] test: cover current allocation modes Signed-off-by: Anish Shanbhag --- tests/test_submit_cli.py | 120 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 120 insertions(+) diff --git a/tests/test_submit_cli.py b/tests/test_submit_cli.py index 05f0034f7..dbb9a0bc6 100644 --- a/tests/test_submit_cli.py +++ b/tests/test_submit_cli.py @@ -183,6 +183,126 @@ def test_apply_current_allocation_requires_slurm_job(monkeypatch, tmp_path: Path assert "--current-allocation requires SLURM_JOB_ID to be set" in capsys.readouterr().out +def test_current_allocation_uses_configured_output_prefetch_and_serve_only(monkeypatch, tmp_path: Path) -> None: + config_path = tmp_path / "config.yaml" + config_path.write_text(yaml.safe_dump(MINIMAL_DRY_RUN_CONFIG)) + source_dir = tmp_path / "source" + prefetch_script = source_dir / "src" / "srtctl" / "runtime_scripts" / "dynamo_wheels.py" + prefetch_script.parent.mkdir(parents=True) + prefetch_script.write_text("") + output_dir = tmp_path / "configured-outputs" + calls = [] + + def fake_run(command, **kwargs): + calls.append((command, kwargs)) + return CompletedProcess(command, 0) + + monkeypatch.setenv("SLURM_JOB_ID", "23456") + monkeypatch.setenv("SRTCTL_SOURCE_DIR", str(source_dir)) + monkeypatch.setenv("SRTCTL_PREFETCH_AI_DYNAMO", "1") + monkeypatch.setattr(submit_cli, "validate_setup", lambda _source: None) + monkeypatch.setattr( + submit_cli, + "get_srtslurm_setting", + lambda key: str(output_dir) if key == "output_dir" else None, + ) + monkeypatch.setattr(submit_cli.subprocess, "run", fake_run) + + result = submit_cli.run_in_current_allocation(config_path, enforce_preflight=False, serve_only=True) + + runtime_config = output_dir / "23456" / "config.yaml" + assert result == 0 + assert calls[0][0] == [sys.executable, str(prefetch_script), "prefetch"] + assert calls[0][1]["check"] is True + assert calls[1][0] == [ + sys.executable, + "-m", + "srtctl.cli.do_sweep", + str(runtime_config), + "--serve-only", + ] + assert calls[1][1]["env"]["SRTCTL_SOURCE_DIR"] == str(source_dir) + + +def test_current_allocation_prefetch_requires_helper(monkeypatch, tmp_path: Path) -> None: + config_path = tmp_path / "config.yaml" + config_path.write_text(yaml.safe_dump(MINIMAL_DRY_RUN_CONFIG)) + source_dir = tmp_path / "source" + + monkeypatch.setenv("SLURM_JOB_ID", "34567") + monkeypatch.setenv("SRTCTL_SOURCE_DIR", str(source_dir)) + monkeypatch.setenv("SRTCTL_PREFETCH_AI_DYNAMO", "1") + monkeypatch.setattr(submit_cli, "validate_setup", lambda _source: None) + + with pytest.raises(FileNotFoundError, match="Dynamo wheel prefetch script not found"): + submit_cli.run_in_current_allocation( + config_path, + output_dir=tmp_path / "outputs", + enforce_preflight=False, + ) + + +@pytest.mark.parametrize( + ("extra_args", "message"), + [ + (["--bash"], "cannot be combined with --bash"), + (["--mock"], "cannot be combined with --mock"), + (["--json"], "cannot be combined with --json"), + (["--sweep"], "supports single-job configs only"), + (["--tags", "smoke"], "cannot be combined with --tags"), + ], +) +def test_current_allocation_rejects_alternate_execution_modes( + monkeypatch, + tmp_path: Path, + capsys, + extra_args: list[str], + message: str, +) -> None: + config_path = tmp_path / "config.yaml" + config_path.write_text(yaml.safe_dump(MINIMAL_DRY_RUN_CONFIG)) + monkeypatch.setattr( + sys, + "argv", + ["srtctl", "apply", "-f", str(config_path), "--current-allocation", *extra_args], + ) + + with pytest.raises(SystemExit) as exc_info: + submit_cli.main() + + assert exc_info.value.code == 2 + assert message in capsys.readouterr().err + + +@pytest.mark.parametrize("input_kind", ["directory", "selector", "sweep"]) +def test_current_allocation_rejects_multi_config_inputs( + monkeypatch, + tmp_path: Path, + capsys, + input_kind: str, +) -> None: + config_path = tmp_path / "config.yaml" + config = dict(MINIMAL_DRY_RUN_CONFIG) + if input_kind == "sweep": + config["sweep"] = {"parameters": {}} + config_path.write_text(yaml.safe_dump(config)) + config_arg = str(tmp_path) if input_kind == "directory" else str(config_path) + if input_kind == "selector": + config_arg += ":base" + monkeypatch.setenv("SLURM_JOB_ID", "45678") + monkeypatch.setattr( + sys, + "argv", + ["srtctl", "apply", "-f", config_arg, "--current-allocation"], + ) + + with pytest.raises(SystemExit) as exc_info: + submit_cli.main() + + assert exc_info.value.code == 1 + assert "--current-allocation" in capsys.readouterr().out + + def test_load_config_rejects_empty_yaml(tmp_path: Path) -> None: path = tmp_path / "empty.yaml" path.write_text("")