Slurm Job: Cluster Sampler & Diagnostics (One-Click)

admin | October 9, 2025

This job collects GPU/CPU, memory, NUMA, PCIe/NVLink, NIC/IB, and optional Nsight/NCCL/iperf3 telemetry across all allocated nodes while your workload runs, then bundles everything into a single .tgz.

Usage: Save as profile_env.slurm and submit:

sbatch --export=ALL,WORKLOAD="torchrun --nproc_per_node=8 train.py --cfg config.yaml",ENABLE_NSYS=1,RUN_NCCL_TESTS=1,DURATION=1800 profile_env.slurm

#!/usr/bin/env bash
#
# profile_env.slurm — cluster-wide performance sampler & diagnostics
#
#SBATCH -J prof-playbook
#SBATCH -o prof-%x-%j.out
#SBATCH -e prof-%x-%j.err
#SBATCH –time=01:00:00
#SBATCH –nodes=1
#SBATCH –ntasks-per-node=1
#SBATCH –gres=gpu:1
#SBATCH –cpus-per-task=8
## Uncomment/adjust for your site:
## #SBATCH –partition=gpu
## #SBATCH –qos=normal

set -euo pipefail

#############################
# Tunables (override via: sbatch –export=ALL,WORKLOAD=”python train.py”,DURATION=900 …)
#############################
WORKLOAD=”${WORKLOAD:-}”                 # e.g., “python train.py”; if empty, uses a tiny CUDA sample
DURATION=”${DURATION:-600}”              # seconds to sample (upper bound)
SAMPLE_INT=”${SAMPLE_INT:-1}”            # sampler interval (seconds)
ENABLE_NSYS=”${ENABLE_NSYS:-0}”          # 1 to record short Nsight Systems traces per node
NSYS_SECONDS=”${NSYS_SECONDS:-45}”       # Nsight trace duration
RUN_NCCL_TESTS=”${RUN_NCCL_TESTS:-0}”    # 1 to run nccl-tests all_reduce_perf
NCCL_TEST_BIN=”${NCCL_TEST_BIN:-all_reduce_perf}”  # path or in $PATH
IPERF_SERVER=”${IPERF_SERVER:-}”         # host/IP to test NIC TCP throughput (iperf3 server required)
OUTROOT=”${OUTROOT:-$PWD}”
TAG=”${TAG:-$(date +%Y%m%d-%H%M%S)}”
OUTDIR=”$OUTROOT/prof-${SLURM_JOB_ID:-nojob}-$TAG”

mkdir -p “$OUTDIR”

echo “==[JOB]==========================================================”
echo ” JOBID        : ${SLURM_JOB_ID:-local}”
echo ” NNODES/NTASKS: ${SLURM_NNODES:-1} / ${SLURM_NTASKS:-1}”
echo ” GPUS         : ${SLURM_GPUS:-unknown} (CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES)”
echo ” WORKLOAD     : ${WORKLOAD:-<built-in demo>}”
echo ” DURATION     : $DURATION  s”
echo ” SAMPLE_INT   : $SAMPLE_INT s”
echo ” ENABLE_NSYS  : $ENABLE_NSYS  (for $NSYS_SECONDS s)”
echo ” RUN_NCCL     : $RUN_NCCL_TESTS”
echo ” IPERF_SERVER : ${IPERF_SERVER:-<none>}”
echo ” OUTDIR       : $OUTDIR”
echo “==================================================================”

# Helper to run a command on every allocated node
node_run() { local cmd=”$1″; srun –ntasks-per-node=1 –nodes=”${SLURM_NNODES:-1}” –label bash -lc “$cmd”; }

# ——– Per-node init: inventory & topology ——–
node_run ‘bash -lc ”
set -euo pipefail
NDIR=\”‘$OUTDIR’/\${HOSTNAME}\”
mkdir -p \”\$NDIR\”
{
  echo \”# Slurm/Env\”
  env | egrep \”SLURM|CUDA|NCCL\” || true
  echo
  echo \”# System\”
  uname -a
  lsb_release -a 2>/dev/null || cat /etc/os-release || true
  date -Iseconds
  echo
  echo \”# CPU/NUMA\”
  lscpu || true
  numactl –hardware || true
} > \”\$NDIR/env.txt\”

if command -v nvidia-smi >/dev/null 2>&1; then
  nvidia-smi -L > \”\$NDIR/gpu_list.txt\” || true
  nvidia-smi topo -m > \”\$NDIR/gpu_topo.txt\” || true
  nvidia-smi -q -x > \”\$NDIR/nvidia_smi.xml\” || nvidia-smi -q > \”\$NDIR/nvidia_smi.txt\” || true
  nvidia-smi pmon -c 1 -s um > \”\$NDIR/pmon_header.txt\” || true
fi

(lspci -nn | egrep -i \”nvidia|mellanox|ethernet|infiniband|network\” || true) > \”\$NDIR/lspci.txt\”
if command -v ibstat >/dev/null 2>&1; then ibstat > \”\$NDIR/ibstat.txt\” || true; fi
if command -v ibv_devinfo >/dev/null 2>&1; then ibv_devinfo > \”\$NDIR/ibv_devinfo.txt\” || true; fi

ip -br link show up | awk ‘\”$1!=\”lo\”{print $1}’\” > \”\$NDIR/ifaces.txt\” || true
while read -r IFACE; do
  (ethtool \”\$IFACE\” && ethtool -k \”\$IFACE\” && ethtool -S \”\$IFACE\” | egrep -i \”err|drop|disc|pause|fcs|crc\” || true) > \”\$NDIR/ethtool_\${IFACE}.txt\” 2>&1 || true
done < \"\$NDIR/ifaces.txt\"
"'

# -------- Start background samplers on each node --------
node_run 'bash -lc "
set -euo pipefail
NDIR=\"'$OUTDIR'/\${HOSTNAME}\"
mkdir -p \"\$NDIR\"
echo $$ > \”\$NDIR/sampler_parent.pid\”

if command -v nvidia-smi >/dev/null 2>&1; then
  (nvidia-smi dmon -s pucvmet -d ‘$SAMPLE_INT’ > \”\$NDIR/gpu_dmon.log\”) &
  echo $! > \”\$NDIR/gpu_dmon.pid\”
  ( (nvidia-smi nvlink -s; while true; do nvidia-smi nvlink -s; sleep 10; done) > \”\$NDIR/nvlink_watch.log\” 2>&1 ) &
  echo $! > \”\$NDIR/nvlink_watch.pid\”
fi

(command -v mpstat >/dev/null 2>&1 && mpstat -P ALL ‘$SAMPLE_INT’ > \”\$NDIR/mpstat.log\”) & echo $! > \”\$NDIR/mpstat.pid\” || true
(command -v pidstat >/dev/null 2>&1 && pidstat -u -r -d ‘$SAMPLE_INT’ > \”\$NDIR/pidstat.log\”) & echo $! > \”\$NDIR/pidstat.pid\” || true
(command -v vmstat  >/dev/null 2>&1 && vmstat ‘$SAMPLE_INT’ > \”\$NDIR/vmstat.log\”)  & echo $! > \”\$NDIR/vmstat.pid\”  || true
(command -v iostat  >/dev/null 2>&1 && iostat -xy ‘$SAMPLE_INT’ > \”\$NDIR/iostat.log\”)  & echo $! > \”\$NDIR/iostat.pid\”  || true

if command -v dcgmi >/dev/null 2>&1; then
  (dcgmi dmon -e 100 -d ‘$SAMPLE_INT’ > \”\$NDIR/dcgm_dmon.log\”) &
  echo $! > \”\$NDIR/dcgm_dmon.pid\”
fi

if command -v numastat >/dev/null 2>&1; then
  (while true; do
     for P in $(pgrep -f -n python || true); do numastat -p \$P; done
     sleep 10
   done > \”\$NDIR/numastat_watch.log\” 2>&1) &
  echo $! > \”\$NDIR/numastat_watch.pid\”
fi
“‘

cleanup() {
  node_run ‘bash -lc ”
    set -euo pipefail
    NDIR=\”‘$OUTDIR’/\${HOSTNAME}\”
    for f in gpu_dmon pidstat mpstat vmstat iostat dcgm_dmon nvlink_watch numastat_watch; do
      if [[ -f \”\$NDIR/\${f}.pid\” ]]; then kill \$(cat \”\$NDIR/\${f}.pid\”) 2>/dev/null || true; fi
    done
  “‘ || true
}
trap cleanup EXIT

# ——– Optional Nsight Systems short trace ——–
if [[ “$ENABLE_NSYS” -eq 1 ]] && command -v nsys >/dev/null 2>&1; then
  node_run ‘bash -lc ”
    set -euo pipefail
    NDIR=\”‘$OUTDIR’/\${HOSTNAME}\”
    mkdir -p \”\$NDIR\”
    nsys profile -t cuda,osrt,nvtx -o \”\$NDIR/nsys_\${HOSTNAME}\” –duration ‘$NSYS_SECONDS’ –stop-on-exit true –capture-range=none sleep ‘$NSYS_SECONDS’
  “‘
fi

# ——– Run workload (or a tiny CUDA demo) ——–
echo “== Running workload ==”
if [[ -z “$WORKLOAD” ]]; then
  node_run ‘bash -lc ”
    set -euo pipefail
    NDIR=\”‘$OUTDIR’/\${HOSTNAME}\”
    echo \”No WORKLOAD provided; running a small CUDA loop…\” | tee -a \”\$NDIR/workload.log\”
    python – <<'PY' || sleep 30
import torch, time
if torch.cuda.is_available():
    a=torch.randn((8192,8192),device='cuda'); b=torch.randn((8192,8192),device='cuda')
    for i in range(50):
        c=a@b; torch.cuda.synchronize()
        time.sleep(0.1)
else:
    time.sleep(30)
PY
  "'
else
  RUNWRK=$(printf '%q ' $WORKLOAD)
  node_run "bash -lc 'set -euo pipefail; NDIR=$OUTDIR/\${HOSTNAME}; echo Running: $RUNWRK | tee -a \"\$NDIR/workload.log\"; $RUNWRK |& tee -a \"\$NDIR/workload.log\"'"
fi
echo "== Workload section complete =="

# -------- Optional network & NCCL checks --------
if [[ -n "$IPERF_SERVER" ]] && command -v iperf3 >/dev/null 2>&1; then
  node_run “iperf3 -c $IPERF_SERVER -P 8 -t 30 | tee ‘$OUTDIR/\${HOSTNAME}/iperf3_\${HOSTNAME}.log'”
fi

if [[ “$RUN_NCCL_TESTS” -eq 1 ]]; then
  node_run “bash -lc ‘set -euo pipefail; NDIR=$OUTDIR/\${HOSTNAME}; if command -v $NCCL_TEST_BIN >/dev/null 2>&1; then $NCCL_TEST_BIN -b 8M -e 512M -f 2 -g \${SLURM_GPUS_PER_NODE:-1} | tee \”\$NDIR/nccl_all_reduce.log\”; else echo \”$NCCL_TEST_BIN not found\” | tee \”\$NDIR/nccl_all_reduce.log\”; fi'”
fi

# ——– Final snapshots & packaging ——–
node_run ‘bash -lc ”
set -euo pipefail
NDIR=\”‘$OUTDIR’/\${HOSTNAME}\”
if command -v nvidia-smi >/dev/null 2>&1; then
  nvidia-smi –query-gpu=index,utilization.gpu,utilization.memory,clocks.sm,clocks.mem,power.draw,temperature.gpu –format=csv -l 1 -f \”\$NDIR/nvidia_smi_final.csv\” -c 3 || true
fi
free -h > \”\$NDIR/free.txt\” || true
df -h   > \”\$NDIR/df.txt\” || true
{
  echo \”=== tail gpu_dmon ===\”; tail -n 30 \”\$NDIR/gpu_dmon.log\” 2>/dev/null || true
  echo
  echo \”=== tail pidstat ===\”; tail -n 30 \”\$NDIR/pidstat.log\” 2>/dev/null || true
} > \”\$NDIR/quick_summary.txt\” || true
“‘

cleanup

tar -C “$OUTROOT” -czf “$OUTDIR.tgz” “$(basename “$OUTDIR”)”
echo “Artifacts packaged at: $OUTDIR.tgz”
echo “Per-node logs under   : $OUTDIR/<hostname>/”
echo “Done.”
    </hostname></none></built-in>

Prefer a direct file? You can also grab the ready-made script: Download profile_env.slurm