refactor: standardize srt-slurm cluster setup hooks

AGENTS.md

@@ -66,6 +66,15 @@ check_env_vars IS_MULTINODE MODEL_NAME PRECISION
 - Do not add launcher-name aliases to `runners/runtime_settings.sh` or elsewhere for scripts that no runner resolves to. A launcher without a pool is dead code; a pool without a launcher fails at job start.
 - When a pool is retired, delete its launcher in the same PR rather than keeping it as a fallback for another pool.
 
+## SRT Slurm cluster hooks
+
+- Put reusable host-check functions in `runners/srt-slurm/hooks/common.sh`. Sourcing it must only define functions, without running checks, changing environment variables, or initializing benchmarks. Cluster-only helpers stay beside their setup script.
+- Keep cluster-specific host prerequisites in `runners/srt-slurm/hooks/<cluster>/setup.sh`, invoked explicitly by the matching cluster profile's `default_host_setup`. These run after allocation, before services and workers start.
+- Hooks are only for checks and setup required by that cluster's hosts or fabric. Keep them small, workload-independent, and safe to run repeatedly. Prefer native srt-slurm configuration whenever it can express the requirement.
+- Do not put benchmark execution, model selection, engine flags, concurrency tuning, evaluation, result collection, or job orchestration in hooks. Those belong in recipes, benchmark scripts, or the existing orchestration layer.
+- Do not use hooks to patch engines or containers, bypass failed checks, or hide runtime bugs behind retries and ad hoc workarounds. Fix problems in the component that owns them.
+- Pass settings explicitly from the cluster profile. Scope mutations to the allocated nodes, preserve other jobs' resources, and register teardown for temporary state that needs restoring. See [cluster profiles](docs/configuration-procedures.md#cluster-profiles).
+
 ## SRT Slurm synthetic acceptance
 
 - **Do not hard-code synthetic acceptance lengths in SRT recipes, master configs, or launchers.** InferenceX automatically selects the measured value from [`golden_al_distribution/`](golden_al_distribution/) for speculative AgentX throughput runs. Do not add manual `SYNTHETIC_ACCEPTANCE_LENGTH`, vLLM `synthetic_acceptance_length`, SGLang `SGLANG_SIMULATE_ACC_LEN`, or TRT-LLM `TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS` settings.

benchmarks/benchmark_lib.sh

@@ -558,33 +558,8 @@ _write_amd_smi_sidecar() {
     fi
 }
 
-# Poll rocm-smi VRAM% every 10s for up to 15 min until the busiest GPU is at or
-# below the threshold percent (default 10); return 1 otherwise so the caller
-# aborts instead of starting on GPUs still draining the previous job.
-# Pass a stricter threshold when the run sizes its KV cache from device-wide free
-# memory (torch.cuda.mem_get_info): on 288 GB parts the 10% gate admits ~28.8 GB
-# of residual, which the engine folds into non_torch and subtracts from the KV
-# pool, so the pool drifts run to run.
-wait_for_amd_gpu_clean() {
-    local threshold="${1:-10}"
-    local gpu_clean=false vram_max i
-    for i in $(seq 1 90); do
-        vram_max=$(rocm-smi --showmemuse 2>/dev/null \
-            | grep -oE "GPU Memory Allocated \(VRAM%\): [0-9]+" \
-            | awk '{if ($NF > m) m = $NF} END {print m+0}')
-        if [ "${vram_max:-0}" -le "$threshold" ]; then
-            echo "GPUs clean (vram%max=$vram_max <= $threshold after $((i * 10))s)"
-            gpu_clean=true
-            break
-        fi
-        echo "waiting for prior-job GPU memory reclaim: vram%max=$vram_max (target <= $threshold)"
-        sleep 10
-    done
-    if [ "$gpu_clean" != "true" ]; then
-        echo "Error: GPUs still draining prior job's memory after 15min" >&2
-        return 1
-    fi
-}
+# shellcheck source=runners/srt-slurm/hooks/common.sh
+source "$(dirname "${BASH_SOURCE[0]}")/../runners/srt-slurm/hooks/common.sh" || return 1
 
 # Return success only while a PID exists and is not a zombie waiting to be
 # reaped. `kill -0` alone treats zombies as live processes.

docs/configuration-procedures.md

@@ -67,6 +67,32 @@ parsed YAML scalars so quotes and punctuation remain data, not YAML or shell syn
 Keep model selection, cache preparation, and workload-dependent time limits in the
 launcher. Do not add profiles for non-srt-slurm launchers or change their routing here.
 
+Put per-allocation host checks and setup in
+`runners/srt-slurm/hooks/<cluster>/setup.sh`, with cluster-specific helpers beside it.
+The directory name matches the cluster profile's filename stem. Invoke the script
+explicitly through `default_host_setup.commands` in that profile; scripts are not
+auto-discovered. srt-slurm runs them on the selected allocated nodes, outside containers,
+before starting services and workers. A failed check stops startup by default.
+Pass configuration explicitly from the profile. If setup needs an undo step, keep it
+in `teardown.sh` beside `setup.sh` and register it in `default_host_setup.teardown`.
+These are job-owned hooks, not administrator-installed Slurm Prolog/Epilog scripts.
+Only add hooks for clusters that need them; do not add empty scripts for every profile.
+
+Put reusable host-check functions in `runners/srt-slurm/hooks/common.sh`; keep
+cluster-only helpers beside `setup.sh`. The common file only defines functions:
+sourcing it must not run checks, change environment variables, or initialize
+benchmarks. Both setup hooks and benchmark scripts can reuse these functions
+without importing benchmark initialization into host setup.
+
+Hooks inject **cluster-specific host prerequisites only**, such as fabric checks or
+required host-state preparation. Keep them small, workload-independent, and safe to
+run repeatedly. Prefer native srt-slurm settings over shell code where possible.
+Benchmark execution, model selection, engine flags, concurrency tuning, evaluation,
+result collection, and job orchestration do not belong here. Do not patch engines or
+containers, bypass failed checks, or hide runtime bugs with retries and ad hoc
+workarounds; fix the owning component instead. Limit host changes to allocated nodes
+and preserve resources used by other jobs.
+
 ## Procedure index
 
 1. [Prepare a worktree](#prepare-a-worktree)

runners/launch_mi300x-amd.sh

@@ -20,7 +20,7 @@ if [[ "$EXECUTION_PATH" == native-single-node ]]; then
     export SALLOC_TIME_LIMIT=180
     export SRT_SRUN_OPTIONS='{"container-remap-root":"", "container-writable":""}'
     SRT_SQUASH_FILE="/raid/inferencex/squash/$(printf '%s' "$IMAGE" | sed 's/[\/:@#]/_/g').sqsh"
-    launch_srt_single_node mi300x-amd
+    launch_srt_single_node mi300x-amd --var GITHUB_WORKSPACE "$GITHUB_WORKSPACE"
     exit $?
 fi
 

runners/launch_mi355x-amds.sh

@@ -19,7 +19,7 @@ if [[ "$EXECUTION_PATH" == native-single-node ]]; then
     export SALLOC_TIME_LIMIT=500
     export SRT_SRUN_OPTIONS='{"container-remap-root":"", "container-writable":""}'
     SRT_SQUASH_FILE="/var/lib/squash/$(printf '%s' "$IMAGE" | sed 's/[\/:@#]/_/g').sqsh"
-    launch_srt_single_node mi355x-amds
+    launch_srt_single_node mi355x-amds --var GITHUB_WORKSPACE "$GITHUB_WORKSPACE"
     exit $?
 fi
 

runners/srt-slurm/hooks/common.sh

@@ -0,0 +1,26 @@
+#!/usr/bin/env bash
+
+# Shared host checks. Sourcing this file only defines functions.
+
+# Poll VRAM usage every 10s for up to 15 minutes. A stricter threshold is useful
+# when the engine sizes its KV cache from device-wide free memory.
+wait_for_amd_gpu_clean() {
+    local threshold="${1:-10}"
+    local gpu_clean=false vram_max i
+    for i in $(seq 1 90); do
+        vram_max=$(rocm-smi --showmemuse 2>/dev/null \
+            | grep -oE "GPU Memory Allocated \(VRAM%\): [0-9]+" \
+            | awk '{if ($NF > m) m = $NF} END {print m+0}')
+        if [ "${vram_max:-0}" -le "$threshold" ]; then
+            echo "GPUs clean (vram%max=$vram_max <= $threshold after $((i * 10))s)"
+            gpu_clean=true
+            break
+        fi
+        echo "waiting for prior-job GPU memory reclaim: vram%max=$vram_max (target <= $threshold)"
+        sleep 10
+    done
+    if [ "$gpu_clean" != "true" ]; then
+        echo "Error: GPUs still draining prior job's memory after 15min" >&2
+        return 1
+    fi
+}

runners/srt-slurm/hooks/mi300x-amd/setup.sh

@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+set -eo pipefail
+
+# RCCL cannot reclaim scratch memory on MEC firmware older than 177 and crashes.
+# See https://rocm.docs.amd.com/en/docs-6.4.3/about/release-notes.html#amdgpu-driver-updates
+minimum=177
+firmware=$(rocm-smi --showfw | awk '/MEC firmware version/ {print $NF}' | sort -n | head -n 1)
+if [[ -z "$firmware" || "$firmware" -lt "$minimum" ]]; then
+    echo "[$(hostname -s)] MEC firmware ${firmware:-unknown} is older than $minimum" >&2
+    exit 1
+fi
+echo "[$(hostname -s)] MEC firmware $firmware"

runners/srt-slurm/hooks/mi355x-amds/check-rdma.sh

@@ -0,0 +1,54 @@
+#!/usr/bin/env bash
+set -eo pipefail
+
+# Fast, per-node fabric preflight for MI355X srt-slurm allocations. This keeps
+# the meaningful QoS/DCQCN gate from the retired amd_utils launcher without its
+# Docker or job-control plumbing.
+
+log() { printf '[%s] %s\n' "$(hostname -s)" "$*"; }
+fail() { log "RDMA preflight failed: $*" >&2; exit 1; }
+
+source "$(dirname "${BASH_SOURCE[0]}")/../../../../benchmarks/benchmark_lib.sh" --validation-only
+check_env_vars IBDEVICES
+expected_devices="$IBDEVICES"
+IFS=',' read -r -a devices <<< "$expected_devices"
+for device in "${devices[@]}"; do
+    [[ -d "/sys/class/infiniband/${device}" ]] || fail "missing device ${device}"
+done
+log "found all ${#devices[@]} expected RDMA devices: ${expected_devices}"
+
+if ! command -v nicctl >/dev/null 2>&1; then
+    log "nicctl is unavailable; device presence passed, QoS/DCQCN checks skipped"
+    exit 0
+fi
+
+probe=$(sudo -n nicctl show version firmware 2>&1 || true)
+if grep -qiE 'No AMD NICs|Invalid card handle|Failed to get NIC' <<< "$probe"; then
+    fail "nicctl cannot access the AMD NICs"
+fi
+
+qos=$(sudo -n nicctl show qos 2>/dev/null) || fail "nicctl show qos failed"
+classification=$(awk '/Classification type/ {print $NF; exit}' <<< "$qos")
+[[ "$classification" == "DSCP" ]] || fail "classification is ${classification:-unset}, expected DSCP"
+
+priorities=$(awk '/PFC no-drop priorities/ {print $NF; exit}' <<< "$qos")
+bitmap=$(awk '/PFC priority bitmap/ {print $NF; exit}' <<< "$qos")
+[[ -n "$priorities" ]] || fail "PFC no-drop priorities are missing"
+[[ -n "$bitmap" && "$bitmap" != "0x0" ]] || fail "PFC is disabled"
+IFS=',' read -r -a priority_values <<< "$priorities"
+for priority in "${priority_values[@]}"; do
+    priority="${priority//[^0-9]/}"
+    [[ -n "$priority" ]] || fail "invalid PFC priority list: ${priorities}"
+    (( bitmap & (1 << priority) )) || fail "PFC bitmap ${bitmap} does not cover priority ${priority}"
+done
+
+dcqcn=$(sudo -n nicctl show dcqcn 2>/dev/null) || fail "nicctl show dcqcn failed"
+device_count=$(grep -c 'ROCE device' <<< "$dcqcn" || true)
+(( device_count > 0 )) || fail "no RoCE devices reported by nicctl"
+if grep 'Status' <<< "$dcqcn" | grep -qv 'Enabled'; then
+    fail "DCQCN is disabled on at least one RoCE device"
+fi
+cnp_count=$(awk '/DSCP value used for CNP/ {print $NF}' <<< "$dcqcn" | sort -u | grep -c . || true)
+(( cnp_count == 1 )) || fail "CNP DSCP is inconsistent across NICs"
+
+log "RDMA QoS/DCQCN preflight passed"

runners/srt-slurm/hooks/mi355x-amds/setup.sh

@@ -0,0 +1,51 @@
+#!/usr/bin/env bash
+set -eo pipefail
+
+bash "$(dirname "${BASH_SOURCE[0]}")/check-rdma.sh"
+
+# Preserve the legacy bare-process GPU drain gate. Slurm owns the node, but a
+# process left outside the prior job's container can still retain VRAM and make
+# the next model load fail much later with a misleading OOM.
+# shellcheck source=runners/srt-slurm/hooks/common.sh
+source "$(dirname "${BASH_SOURCE[0]}")/../common.sh"
+wait_for_amd_gpu_clean
+
+# Some MI355X experiments reserve large 2 MiB HugeTLB pools and leave the
+# reservation behind after their Slurm allocation exits. Those free hugepages
+# remain unavailable to ordinary host allocations, which can make a later
+# unchanged SGLang HiCache recipe fail even on a 3 TiB node. Reclaim only free
+# pages; pages currently used or reserved by host services are preserved.
+meminfo=/proc/meminfo
+nr_hugepages=/proc/sys/vm/nr_hugepages
+
+read_hugepage_value() {
+    local key="$1"
+    awk -v key="${key}:" '$1 == key {print $2}' "$meminfo"
+}
+
+total=$(read_hugepage_value HugePages_Total)
+free=$(read_hugepage_value HugePages_Free)
+reserved=$(read_hugepage_value HugePages_Rsvd)
+used=$((total - free))
+target=$((used + reserved))
+
+echo "MI355X host memory before preparation:"
+grep -E '^(MemAvailable|HugePages_Total|HugePages_Free|HugePages_Rsvd|HugePages_Surp|Hugetlb):' "$meminfo"
+
+if (( target < total )); then
+    printf '%s\n' "$target" | sudo -n tee "$nr_hugepages" >/dev/null
+fi
+
+after_total=$(read_hugepage_value HugePages_Total)
+after_free=$(read_hugepage_value HugePages_Free)
+echo "MI355X host memory after preparation:"
+grep -E '^(MemAvailable|HugePages_Total|HugePages_Free|HugePages_Rsvd|HugePages_Surp|Hugetlb):' "$meminfo"
+
+if (( after_total - after_free < used )); then
+    echo "Host preparation released hugepages that were in use" >&2
+    exit 1
+fi
+if (( after_free > reserved )); then
+    echo "Host preparation could not reclaim all unused hugepages" >&2
+    exit 1
+fi

runners/srt-slurm/mi300x-amd.yaml

@@ -12,10 +12,7 @@ default_sbatch_directives:
 use_gpus_per_node_directive: true
 use_segment_sbatch_directive: false
 use_exclusive_sbatch_directive: true
-default_bash_preamble: |-
-  if [[ "$$MODEL_PREFIX" == dsr1 && "$$FRAMEWORK" == sglang ]]; then
-      firmware=$$(rocm-smi --showfw | awk '/MEC/ {print $$NF; exit}')
-      if [[ -z "$$firmware" || "$$firmware" -lt 177 ]]; then
-          export HSA_NO_SCRATCH_RECLAIM=1
-      fi
-  fi
+default_host_setup:
+  commands:
+    - bash "${GITHUB_WORKSPACE}/runners/srt-slurm/hooks/mi300x-amd/setup.sh"
+  nodes: all

runners/srt-slurm/mi355x-amds.yaml

@@ -12,3 +12,10 @@ default_sbatch_directives:
 use_gpus_per_node_directive: true
 use_segment_sbatch_directive: false
 use_exclusive_sbatch_directive: true
+default_host_setup:
+  # The GPU-drain check can take 15 minutes; leave room for the fabric check
+  # and host-memory preparation.
+  timeout_seconds: 1200
+  commands:
+    - IBDEVICES=rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7 bash "${GITHUB_WORKSPACE}/runners/srt-slurm/hooks/mi355x-amds/setup.sh"
+  nodes: all

utils/test_srt_fixed_sequence.py

@@ -172,6 +172,9 @@ def test_native_post_eval_preserves_results_topology_and_failure(client_environm
     scripts = workspace / "benchmarks/single_node"
     scripts.mkdir(parents=True)
     shutil.copyfile(ROOT / "benchmarks/benchmark_lib.sh", scripts.parent / "benchmark_lib.sh")
+    hooks = workspace / "runners/srt-slurm/hooks"
+    hooks.mkdir(parents=True)
+    shutil.copyfile(ROOT / "runners/srt-slurm/hooks/common.sh", hooks / "common.sh")
     shutil.copyfile(ROOT / "benchmarks/single_node/srt_eval.sh", scripts / "srt_eval.sh")
     python = tmp_path / "bin/python3"
     python.write_text(