From 22adb1401e8a4fc9bdf03d9ecefc2c8845096785 Mon Sep 17 00:00:00 2001
From: eligotts <eligotts@users.noreply.github.com>
Date: Sun, 5 Jul 2026 20:18:51 +0000
Subject: [PATCH] fix(slurm): tear down inference-only nodes after clean
 trainer exit
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

`srun --kill-on-bad-exit=1` only reaps peer tasks on non-zero exits, so a
trainer that finishes cleanly left inference-only nodes running vLLM until
the allocation hit its time limit.

Add a run-completion sentinel for separate train/inference deployments:

- The trainer task wrapper touches `$OUTPUT_DIR/control/slurm_run_complete`
  on exit 0 (and propagates its exit code either way).
- Inference-only nodes run a watcher that exits cleanly once the sentinel
  appears; an inference component exiting *without* the sentinel is treated
  as a failure so kill-on-bad-exit still tears the job down.
- The orchestrator-hosting task re-waits when the orchestrator exits cleanly
  before the trainer finishes — otherwise its `wait -n` would return early
  and kill the trainer wrapper before it touches the sentinel.

Validated by rendering the template with and without inference nodes and
`bash -n` on both outputs; a live two-direction SLURM smoke (clean exit
teardown, failure propagation) is still pending.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 .../templates/multi_node_rl.sbatch.j2         | 64 ++++++++++++++-----
 1 file changed, 49 insertions(+), 15 deletions(-)

diff --git a/src/prime_rl/templates/multi_node_rl.sbatch.j2 b/src/prime_rl/templates/multi_node_rl.sbatch.j2
index 4da0c2120b..8a09bdd0ec 100755
--- a/src/prime_rl/templates/multi_node_rl.sbatch.j2
+++ b/src/prime_rl/templates/multi_node_rl.sbatch.j2
@@ -59,10 +59,14 @@ export PROJECT_DIR={{ project_dir }}
 export CONFIG_DIR={{ config_dir }}
 export OUTPUT_DIR={{ output_dir }}
 export ORCHESTRATOR_OUTPUT_DIR={{ orchestrator_output_dir }}
-mkdir -p $OUTPUT_DIR/logs/trainer $OUTPUT_DIR/logs/inference
+mkdir -p $OUTPUT_DIR/logs/trainer $OUTPUT_DIR/logs/inference $OUTPUT_DIR/control
 rm -f $OUTPUT_DIR/logs/inference/*.log
 ln -sfn trainer/node_0.log $OUTPUT_DIR/logs/trainer.log
 ln -sfn inference/node_0.log $OUTPUT_DIR/logs/inference.log
+{% if num_infer_nodes > 0 -%}
+export RUN_DONE_FILE="$OUTPUT_DIR/control/slurm_run_complete"
+rm -f "$RUN_DONE_FILE"
+{% endif %}
 
 export WANDB_SHARED_MODE=1
 export WANDB_SHARED_RUN_ID=${WANDB_SHARED_RUN_ID:-$(python3 -c "import uuid; print(uuid.uuid4().hex)")}
@@ -391,20 +395,29 @@ else
 
     echo $HOSTNAMES_STR | tee $OUTPUT_DIR/logs/trainer/node_${TRAIN_NODE_RANK}.log
 
-    WANDB_SHARED_LABEL=trainer uv run torchrun \
-        --role=trainer \
-        --nnodes=$NUM_TRAIN_NODES \
-        --nproc-per-node=$GPUS_PER_NODE \
-        --node-rank=$TRAIN_NODE_RANK \
-        --rdzv-endpoint=$MASTER_ADDR:$MASTER_PORT \
-        --rdzv-id=job_$SLURM_JOB_ID \
-        --log-dir=$OUTPUT_DIR/logs/trainer/torchrun \
-        --tee=3 \
-        --redirects=3 \
-        --local-ranks-filter={{ ranks_filter }} \
-        -m prime_rl.trainer.rl.train \
-        @ $CONFIG_DIR/trainer.toml \
-        2>&1 | sed -u 's/^\[[a-zA-Z]*[0-9]*\]://' | tee -a $OUTPUT_DIR/logs/trainer/node_${TRAIN_NODE_RANK}.log &
+    (
+        WANDB_SHARED_LABEL=trainer uv run torchrun \
+            --role=trainer \
+            --nnodes=$NUM_TRAIN_NODES \
+            --nproc-per-node=$GPUS_PER_NODE \
+            --node-rank=$TRAIN_NODE_RANK \
+            --rdzv-endpoint=$MASTER_ADDR:$MASTER_PORT \
+            --rdzv-id=job_$SLURM_JOB_ID \
+            --log-dir=$OUTPUT_DIR/logs/trainer/torchrun \
+            --tee=3 \
+            --redirects=3 \
+            --local-ranks-filter={{ ranks_filter }} \
+            -m prime_rl.trainer.rl.train \
+            @ $CONFIG_DIR/trainer.toml \
+            2>&1 | sed -u 's/^\[[a-zA-Z]*[0-9]*\]://' | tee -a $OUTPUT_DIR/logs/trainer/node_${TRAIN_NODE_RANK}.log
+        TRAIN_EXIT_CODE=$?
+        {% if num_infer_nodes > 0 -%}
+        if [ "$TRAIN_EXIT_CODE" -eq 0 ]; then
+            touch "$RUN_DONE_FILE"
+        fi
+        {% endif -%}
+        exit "$TRAIN_EXIT_CODE"
+    ) &
 {% if num_infer_nodes > 0 %}    fi{% endif %}
 {% if num_infer_nodes > 0 %}
 # Launch the orchestrator on one node. Default is trainer rank 0 (SLURM_PROCID ==
@@ -427,6 +440,15 @@ if [ "$SLURM_PROCID" -eq "$ORCH_PROCID" ]; then
         2>&1 | tee $OUTPUT_DIR/logs/orchestrator.log &
 fi
 {% endif %}
+{% if num_infer_nodes > 0 %}
+if [ "$SLURM_PROCID" -lt "$NUM_INFER_NODES" ]; then
+    (
+        while [ ! -f "$RUN_DONE_FILE" ]; do
+            sleep 5
+        done
+    ) &
+fi
+{% endif %}
 
 # Wait for the first background job to exit and propagate its exit code. This
 # turns background process crashes (router, orchestrator) into task failures
@@ -435,6 +457,18 @@ fi
 wait -n
 EXIT_CODE=$?
 REMAINING_PIDS=$(jobs -p)
+{%- if num_infer_nodes > 0 %}
+if [ "$SLURM_PROCID" -eq "$ORCH_PROCID" ] && [ "$EXIT_CODE" -eq 0 ] && [ ! -f "$RUN_DONE_FILE" ]; then
+    echo "[$(hostname)] orchestrator exited before trainer completion; waiting for trainer sentinel"
+    wait -n
+    EXIT_CODE=$?
+    REMAINING_PIDS=$(jobs -p)
+fi
+if [ "$SLURM_PROCID" -lt "$NUM_INFER_NODES" ] && [ "$EXIT_CODE" -eq 0 ] && [ ! -f "$RUN_DONE_FILE" ]; then
+    echo "[$(hostname)] inference component exited before trainer completion; treating as failure"
+    EXIT_CODE=1
+fi
+{% endif -%}
 {%- if cleanup_grace_period %}
 # Cross-node teardown is driven by `srun --kill-on-bad-exit=1`, which reaps peer
 # tasks only once a task *exits* non-zero. So on a non-zero exit (crash, SIGTERM,