feat: run apply in current Slurm allocation

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
 

docs/cli.md

@@ -254,6 +254,7 @@ srtctl apply -f <config.yaml> [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 <job-id>` 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

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:

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