Update SLURM skills to handle launching from the login nodes

.agents/skills/add-new-model/reference/testing.md

@@ -12,6 +12,8 @@ nvidia-smi --query-gpu=index,memory.used,memory.free --format=csv
 
 Pick indices whose `memory.used` is under ~1000 MiB. Do this check **again** before each command, not once at the start of the session - another user can grab GPUs between your runs. When running multiple torchruns in sequence, reset the chosen GPU set each time.
 
+If no GPUs are visible (e.g., on a login node), dispatch each command as an srun step via the **launch-with-slurm** skill and keep the `CUDA_VISIBLE_DEVICES` selection inside the step.
+
 ## Timeouts (keep them short)
 
 - Single-GPU sanity: 120s

.agents/skills/capture-nsys-profile/SKILL.md

@@ -38,7 +38,7 @@ Target DP=1 (smallest world that satisfies the parallelism plan). Assuming 8 GPU
 
 ### If running under SLURM
 
-If `$SLURM_JOB_ID` is set, use the **launch-with-slurm** skill to read the allocation's node count and compare it to the minimum above. If the allocation is short, surface that to the user instead of launching.
+Use the **launch-with-slurm** skill to find the allocation and read its node count; compare it to the minimum above. If the allocation is short, surface that to the user instead of launching.
 
 ## Step 3: Capture the profile
 

.agents/skills/capture-nsys-profile/scripts/launch_capture.sh

@@ -21,6 +21,7 @@ NSYS_ARGS+=(--delay=0)
 
 TORCHRUN_ARGS=()
 TORCHRUN_ARGS+=(--nnodes=${SLURM_NNODES:-1} --node-rank=${SLURM_NODEID:-0} --nproc-per-node=gpu)
-TORCHRUN_ARGS+=(--rdzv-backend=c10d --rdzv-endpoint=${SLURM_LAUNCH_NODE_IPADDR:-localhost}:15213)
+RDZV_HOST=$(scontrol show hostnames "${SLURM_STEP_NODELIST:-localhost}" | head -1 || echo localhost)
+TORCHRUN_ARGS+=(--rdzv-backend=c10d --rdzv-endpoint=$RDZV_HOST:15213)
 
 nsys ${NSYS_ARGS[@]} torchrun ${TORCHRUN_ARGS[@]} $SCRIPT $@

.agents/skills/launch-with-slurm/SKILL.md

@@ -1,20 +1,33 @@
 ---
 name: launch-with-slurm
-description: Reference for launching jobs inside a SLURM allocation via srun (single-node or multi-node). Use whenever $SLURM_JOB_ID is set and work needs to run on the allocated compute — from direct user requests ("run on the cluster", "launch on slurm", "train across N nodes", "dispatch the job") OR from within another skill's workflow (e.g., validate-correctness running validation on the allocation, add-new-model reaching pp=2/ep=2). Covers how to read the allocation context from $SLURM_JOB_ID via scontrol, the srun flags that matter (-W 0, -N, -o, --open-mode, --nodelist), and gotchas like the executable bit requirement and distributed-aware output redirection.
+description: Reference for launching jobs inside a SLURM allocation via srun (single-node or multi-node). Use whenever work needs to run on allocated compute — from direct user requests ("run on the cluster", "use my running job", "on my allocation", "launch on slurm", "train across N nodes", "dispatch the job") OR from within another skill's workflow (e.g., validate-correctness running validation on the allocation, add-new-model reaching pp=2/ep=2). Covers finding the target job ($SLURM_JOB_ID or squeue, including array-job ID resolution), reading the allocation with scontrol, the srun flags that matter (--jobid, -W 0, -N, -o, --open-mode, --nodelist, --overlap), and gotchas like GPU sharing with live steps and distributed-aware output redirection.
 ---
 
 # Launch with SLURM
 
-Reference for dispatching work inside a SLURM allocation via `srun`. Most of the time the agent is invoked from inside an existing allocation — the compute is already reserved, and the job is to use it well. Applies to both single-node and multi-node allocations.
+`srun --jobid=<jobid>` attaches a step to a running allocation regardless of where it is invoked — login node, `ssh`'d compute node, or inside the job itself. The scheduler draws nodes from the allocation pool and may pick a node other than the invoking one. Default to it over raw `torchrun`/`bash`: `srun` propagates env vars, handles distributed-aware I/O, and manages signals across ranks correctly — even on a single node with multiple GPUs. The `examples/*/launch.sh` scripts read `SLURM_NNODES`/`SLURM_NODEID`, which `srun` sets inside a step, so they work unchanged.
 
-Default to `srun` over raw `torchrun`/`bash` whenever `$SLURM_JOB_ID` is set. `srun` propagates env vars, handles distributed-aware I/O, and manages signals across ranks correctly — even on a single node with multiple GPUs.
+## Step 1: Find the Job ID
 
-## Step 1: Read the Allocation
+Two sources, one answer — the numeric JobId of the job (e.g. `2193449`):
 
-If `$SLURM_JOB_ID` is set, the compute environment is already encoded — don't guess, ask SLURM:
+- **`$SLURM_JOB_ID` is set** — you are inside the job, or in an `ssh` session on an allocated node that inherits the job environment. Use it as-is; it is already the numeric JobId.
+- **Otherwise** — list the user's jobs and pick a RUNNING one: by name if the user named it, ask if several are ambiguous or none is RUNNING, skip PENDING (no compute yet):
 
 ```bash
-scontrol show job $SLURM_JOB_ID
+squeue -u $USER -o "%.10i %.9T %.6D %.24N %.24j %.10L"
+```
+
+**Array jobs: resolve to the numeric JobId.** `squeue` reports elements as `<master>_<task>` (e.g. `2193448_1`), but `--jobid` wants the element's own numeric JobId — given `2193448_1`, it strips the `_1`, resolves to the array master, and fails with "Job is pending execution" when the master has pending elements:
+
+```bash
+scontrol show job 2193448_1 | grep -oP '^JobId=\K[0-9]+'    # → 2193449
+```
+
+Then read the allocation — don't guess, ask SLURM:
+
+```bash
+scontrol show job $JOBID
 ```
 
 Key fields to extract:
@@ -27,22 +40,23 @@ Key fields to extract:
 For a quick remaining-time check, use `squeue` directly — it returns `D-HH:MM:SS` without needing to parse timestamps:
 
 ```bash
-squeue -h -j $SLURM_JOB_ID -o %L
+squeue -h -j $JOBID -o %L
 ```
 
 Before launching anything long-running, compare this against the estimated runtime. If the budget is too tight, surface this to the user instead of launching and getting killed mid-run.
 
-If `$SLURM_JOB_ID` is not set, you're not inside an allocation — fall back to direct `torchrun` and do not attempt to invoke `srun`.
-
 ## Step 2: Build the srun Command
 
-### Flags that matter
-
-- **`-N <n>`** — number of nodes to dispatch to. In most training runs this matches the pipeline-parallel degree (PP), but not always: the full parallelism plan and GPUs-per-node determine total nodes (e.g., PP=1 with EP=16 on 8-GPU nodes still needs 2 nodes).
+- **`--jobid=<jobid>`** — anchor the step to the allocation. Required when `$SLURM_JOB_ID` is unset or holds a different job; redundant but harmless when it already holds the target.
+- **`-N <n>`** — number of nodes to dispatch to. In most training runs this matches PP, but the full parallelism plan and GPUs-per-node determine total nodes (e.g., PP=1 with EP=16 on 8-GPU nodes still needs 2). `-N1` borrows one node of a multi-node allocation for probes or single-node tests.
 - **`-W 0`** — wait indefinitely for stragglers after the first task exits. The default behavior terminates remaining tasks shortly after the first one ends, which kills workers that are still cleanly shutting down. Always use `-W 0` for training and evaluation runs.
 - **`-o <file>`** — stdout redirection. Use this instead of piping through `tee`. On multi-node, `tee`ing srun output collapses concurrent writes from all ranks. `-o` is distributed-aware — srun collects output from every rank into the single specified file, preserving the one-command-one-log abstraction. By convention, PithTrain runs log under `logging/<descriptive-name>.log`.
 - **`--open-mode=append`** vs **`--open-mode=truncate`** — for resumed training, `append` preserves history across restarts. Use `truncate` for fresh runs where overwriting is intended.
 - **`--nodelist=<hosts>`** — restrict dispatch to specific nodes. Useful for debugging at a smaller scale (e.g., 4 nodes allocated, but debug with 2 specific ones).
+- **`--overlap`** — share the allocation's CPUs, memory, and GPUs with the job's other steps. Without it, a step reserves the whole allocation, so any later step silently pends until it finishes — pass `--overlap` whenever launching alongside a running step.
+- **`--gres=gpu:<k>`** — request a subset of the job's GPUs for the step. By default a step sees every GPU the job holds on the node, so `torchrun --nproc-per-node=gpu` just works. Sharing means contention: probe `nvidia-smi` utilization before launching heavy work onto a job whose GPUs are already busy; if it's saturated, surface this to the user instead of piling on.
+
+`srun` execs the command directly, not through a shell — invoke scripts as `bash <script>` or ensure the `+x` bit.
 
 ## References
 

.agents/skills/setup-benchmark-inputs/SKILL.md

@@ -19,6 +19,6 @@ mkdir -p logging
 # Single-node (DeepSeek-V2-Lite)
 bash .agents/skills/setup-benchmark-inputs/scripts/launch_setup.sh --model deepseek-v2-lite 2>&1 | tee logging/setup-deepseek-v2-lite.log
 
-# Multi-node via SLURM (Qwen3-30B-A3B)
+# Multi-node via SLURM (Qwen3-30B-A3B) — anchor the step with the launch-with-slurm skill
 srun -W 0 -o logging/setup-qwen3-30b-a3b.log .agents/skills/setup-benchmark-inputs/scripts/launch_setup.sh --model qwen3-30b-a3b
 ```

.agents/skills/validate-correctness/scripts/launch_validate.sh

@@ -9,6 +9,7 @@ SCRIPT=.agents/skills/validate-correctness/scripts/validate.py
 
 TORCHRUN_ARGS=()
 TORCHRUN_ARGS+=(--nnodes=${SLURM_NNODES:-1} --node-rank=${SLURM_NODEID:-0} --nproc-per-node=gpu)
-TORCHRUN_ARGS+=(--rdzv-backend=c10d --rdzv-endpoint=${SLURM_LAUNCH_NODE_IPADDR:-localhost}:15213)
+RDZV_HOST=$(scontrol show hostnames "${SLURM_STEP_NODELIST:-localhost}" | head -1 || echo localhost)
+TORCHRUN_ARGS+=(--rdzv-backend=c10d --rdzv-endpoint=$RDZV_HOST:15213)
 
 torchrun ${TORCHRUN_ARGS[@]} $SCRIPT $@

benchmarks/pretraining/deepseek-v2-lite/h100-1n8g/pp1-dp1-cp1-ep8-seq2048-bf16.sh

@@ -1,8 +1,11 @@
 #!/bin/bash
 # Launch the training.
 
+set -euo pipefail
+
 TRUN_ARGS=()
 TRUN_ARGS+=(--nnodes=${SLURM_NNODES:-1} --node-rank=${SLURM_NODEID:-0} --nproc-per-node=gpu)
-TRUN_ARGS+=(--rdzv-backend=c10d --rdzv-endpoint=${SLURM_LAUNCH_NODE_IPADDR:-localhost}:15213)
+RDZV_HOST=$(scontrol show hostnames "${SLURM_STEP_NODELIST:-localhost}" | head -1 || echo localhost)
+TRUN_ARGS+=(--rdzv-backend=c10d --rdzv-endpoint=$RDZV_HOST:15213)
 
 torchrun ${TRUN_ARGS[@]} benchmarks/pretraining/deepseek-v2-lite/h100-1n8g/pp1-dp1-cp1-ep8-seq2048-bf16.py

benchmarks/pretraining/gpt-oss-120b/h100-8n8g/pp8-dp1-cp1-ep8-seq2048-bf16.sh

@@ -1,8 +1,11 @@
 #!/bin/bash
 # Launch the training.
 
+set -euo pipefail
+
 TRUN_ARGS=()
 TRUN_ARGS+=(--nnodes=${SLURM_NNODES:-1} --node-rank=${SLURM_NODEID:-0} --nproc-per-node=gpu)
-TRUN_ARGS+=(--rdzv-backend=c10d --rdzv-endpoint=${SLURM_LAUNCH_NODE_IPADDR:-localhost}:15213)
+RDZV_HOST=$(scontrol show hostnames "${SLURM_STEP_NODELIST:-localhost}" | head -1 || echo localhost)
+TRUN_ARGS+=(--rdzv-backend=c10d --rdzv-endpoint=$RDZV_HOST:15213)
 
 torchrun ${TRUN_ARGS[@]} benchmarks/pretraining/gpt-oss-120b/h100-8n8g/pp8-dp1-cp1-ep8-seq2048-bf16.py

benchmarks/pretraining/qwen3-30b-a3b/h100-2n8g/pp2-dp1-cp1-ep8-seq2048-bf16.sh

@@ -1,8 +1,11 @@
 #!/bin/bash
 # Launch the training.
 
+set -euo pipefail
+
 TRUN_ARGS=()
 TRUN_ARGS+=(--nnodes=${SLURM_NNODES:-1} --node-rank=${SLURM_NODEID:-0} --nproc-per-node=gpu)
-TRUN_ARGS+=(--rdzv-backend=c10d --rdzv-endpoint=${SLURM_LAUNCH_NODE_IPADDR:-localhost}:15213)
+RDZV_HOST=$(scontrol show hostnames "${SLURM_STEP_NODELIST:-localhost}" | head -1 || echo localhost)
+TRUN_ARGS+=(--rdzv-backend=c10d --rdzv-endpoint=$RDZV_HOST:15213)
 
 torchrun ${TRUN_ARGS[@]} benchmarks/pretraining/qwen3-30b-a3b/h100-2n8g/pp2-dp1-cp1-ep8-seq2048-bf16.py

benchmarks/pretraining/qwen3-30b-a3b/h100-4n8g/pp4-dp1-cp1-ep8-seq4096-bf16.sh

@@ -1,8 +1,11 @@
 #!/bin/bash
 # Launch the training.
 
+set -euo pipefail
+
 TRUN_ARGS=()
 TRUN_ARGS+=(--nnodes=${SLURM_NNODES:-1} --node-rank=${SLURM_NODEID:-0} --nproc-per-node=gpu)
-TRUN_ARGS+=(--rdzv-backend=c10d --rdzv-endpoint=${SLURM_LAUNCH_NODE_IPADDR:-localhost}:15213)
+RDZV_HOST=$(scontrol show hostnames "${SLURM_STEP_NODELIST:-localhost}" | head -1 || echo localhost)
+TRUN_ARGS+=(--rdzv-backend=c10d --rdzv-endpoint=$RDZV_HOST:15213)
 
 torchrun ${TRUN_ARGS[@]} benchmarks/pretraining/qwen3-30b-a3b/h100-4n8g/pp4-dp1-cp1-ep8-seq4096-bf16.py

benchmarks/pretraining/qwen3-30b-a3b/h100-8n8g/pp4-dp2-cp1-ep8-seq4096-bf16.sh

@@ -1,8 +1,11 @@
 #!/bin/bash
 # Launch the training.
 
+set -euo pipefail
+
 TRUN_ARGS=()
 TRUN_ARGS+=(--nnodes=${SLURM_NNODES:-1} --node-rank=${SLURM_NODEID:-0} --nproc-per-node=gpu)
-TRUN_ARGS+=(--rdzv-backend=c10d --rdzv-endpoint=${SLURM_LAUNCH_NODE_IPADDR:-localhost}:15213)
+RDZV_HOST=$(scontrol show hostnames "${SLURM_STEP_NODELIST:-localhost}" | head -1 || echo localhost)
+TRUN_ARGS+=(--rdzv-backend=c10d --rdzv-endpoint=$RDZV_HOST:15213)
 
 torchrun ${TRUN_ARGS[@]} benchmarks/pretraining/qwen3-30b-a3b/h100-8n8g/pp4-dp2-cp1-ep8-seq4096-bf16.py

benchmarks/pretraining/qwen3.5-35b-a3b/h100-2n8g/pp2-dp1-cp1-ep8-seq1536-bf16.sh

@@ -1,8 +1,11 @@
 #!/bin/bash
 # Launch the training.
 
+set -euo pipefail
+
 TRUN_ARGS=()
 TRUN_ARGS+=(--nnodes=${SLURM_NNODES:-1} --node-rank=${SLURM_NODEID:-0} --nproc-per-node=gpu)
-TRUN_ARGS+=(--rdzv-backend=c10d --rdzv-endpoint=${SLURM_LAUNCH_NODE_IPADDR:-localhost}:15213)
+RDZV_HOST=$(scontrol show hostnames "${SLURM_STEP_NODELIST:-localhost}" | head -1 || echo localhost)
+TRUN_ARGS+=(--rdzv-backend=c10d --rdzv-endpoint=$RDZV_HOST:15213)
 
 torchrun ${TRUN_ARGS[@]} benchmarks/pretraining/qwen3.5-35b-a3b/h100-2n8g/pp2-dp1-cp1-ep8-seq1536-bf16.py

benchmarks/pretraining/qwen3.5-35b-a3b/h100-4n8g/pp4-dp1-cp1-ep8-seq3072-bf16.sh

@@ -1,8 +1,11 @@
 #!/bin/bash
 # Launch the training.
 
+set -euo pipefail
+
 TRUN_ARGS=()
 TRUN_ARGS+=(--nnodes=${SLURM_NNODES:-1} --node-rank=${SLURM_NODEID:-0} --nproc-per-node=gpu)
-TRUN_ARGS+=(--rdzv-backend=c10d --rdzv-endpoint=${SLURM_LAUNCH_NODE_IPADDR:-localhost}:15213)
+RDZV_HOST=$(scontrol show hostnames "${SLURM_STEP_NODELIST:-localhost}" | head -1 || echo localhost)
+TRUN_ARGS+=(--rdzv-backend=c10d --rdzv-endpoint=$RDZV_HOST:15213)
 
 torchrun ${TRUN_ARGS[@]} benchmarks/pretraining/qwen3.5-35b-a3b/h100-4n8g/pp4-dp1-cp1-ep8-seq3072-bf16.py

examples/pretrain_lm/launch.sh

@@ -20,7 +20,8 @@ fi
 # Setup distributed.
 LAUNCH_ARGS=()
 LAUNCH_ARGS+=(--nnodes=${SLURM_NNODES:-1} --node-rank=${SLURM_NODEID:-0} --nproc-per-node=gpu)
-LAUNCH_ARGS+=(--rdzv-backend=c10d --rdzv-endpoint=${SLURM_LAUNCH_NODE_IPADDR:-localhost}:15213)
+RDZV_HOST=$(scontrol show hostnames "${SLURM_STEP_NODELIST:-localhost}" | head -1 || echo localhost)
+LAUNCH_ARGS+=(--rdzv-backend=c10d --rdzv-endpoint=$RDZV_HOST:15213)
 
 # Launch the training.
 SCRIPT=examples/pretrain_lm/$1/script.py