slurm: make run-parts.sh exclusive detection work with custom prefix and recent Slurm

[100milliongold]

Problem

The exclusive-job detection in roles/slurm/templates/etc/slurm/shared/bin/run-parts.sh fails in two independent ways:

1. PATH. It calls scontrol/squeue through PATH. slurmd's environment does not include a custom slurm_install_prefix, so both commands return nothing, numcpus_sys and numcpus_job are both empty, "" == "" is true, and every job runs the *-exclusive-* scripts. On a shared node that resets power limits and application clocks on all GPUs and drops page caches for each job (prolog took 6–8 s; srun: Prolog hung on node).

2. Parsing. grep -Eio "TRES=cpu=[0-9]+" matches both ReqTRES= and AllocTRES= lines on recent Slurm, so numcpus_job becomes a multi-line value ('1\n56' in a bash -x trace) that never compares equal. With scontrol on PATH, exclusive jobs are therefore never detected.

Reproduced on DGX OS 7.5.0, Slurm 26.05.1, slurm_install_prefix: /raid/slurm/usr/local.

Fix

• Call {{ slurm_install_prefix }}/bin/squeue by absolute path (the script is deployed with the template module, so the variable is available).
• Read allocated CPUs and node count with squeue -o %C / -o %D instead of parsing scontrol show job.
• Guard against an empty result so a lookup failure means "not exclusive".

Verification

On the system above, before the fix every srun --gres=gpu:1 job logged Running .../50-exclusive-gpu; after symlinking the binaries onto PATH (which exercises failure 2) a bash -x run showed numcpus_job='1\n56'. The patched logic yields numcpus_job=56, numcpus_sys=256, exclusive=0 for that job.

[dholt]

run-parts.sh has set -e, so the direct numcpus_job=$(squeue ...) and numnodes_job=$(squeue ...) assignments terminate the script when squeue returns nonzero; redirecting stderr does not suppress that exit status. Please perform both lookups inside a conditional or otherwise explicitly handle failure so exclusive remains 0 and normal non-exclusive scripts still run. Add the required decision-matrix proof for failed squeue as well as empty output.

---

Automated triage review (agent-generated on the maintainer's behalf; a human maintainer decides merges).

[100milliongold]

Thanks — the set -e point is correct and I reproduced it before changing anything:

$ bash -c 'set -e; v=$(false); echo "reached"' ; echo "rc=$?"
rc=1

The echo never runs, so a transient squeue failure took down the whole prolog/epilog run rather than just the exclusive detection.

What changed (commit ebb66b8):

• Both allocation fields now come from a single -o "%C %D" call placed inside an if condition, so the failure is visible and handled. A failed, empty or non-numeric lookup leaves exclusive=0 and logs a warning.
• The same reasoning is applied to the last-user-job count, which you did not ask about but has the same shape with a worse failure direction: piping squeue into wc -l hides its exit status, and a failed lookup would be read as "no other jobs" and run the *-lastuserjob-* cleanup scripts while another job of the same user is still on the node. A failed lookup now leaves last_user_job=0.

Decision matrix, verified on a DGX B300 (Ubuntu 24.04, bash 5.2, 256 CPUs) with a squeue stub that honours the -o format and the -j flag:

case	before	after
exclusive	rc=0, all three parts ran	rc=0, all three parts ran
shared	rc=0, exclusive part skipped	rc=0, exclusive part skipped
other jobs present	rc=0, lastuserjob part skipped	rc=0, lastuserjob part skipped
empty output	rc=0, silently non-exclusive	rc=0, non-exclusive + 1 warning
non-numeric output	rc=0, silently non-exclusive	rc=0, non-exclusive + 1 warning
nonzero exit	rc=1, no part ran at all	rc=0, normal part ran + 2 warnings

The three normal cases are unchanged, so there is no regression; the failure cases stop taking the whole run down and stop deciding silently.

One note on the test itself: my first stub ignored -o and returned both fields to every query, which does not exercise the old code fairly (it asks for %C and %D separately). The table above is from the corrected stub.

[100milliongold]

Thanks for the review — you were right, and I've now verified both failure modes
explicitly rather than reasoning about them.

set -e really does abort on the assignment

Two minimal cases, run as-is:

$ cat a.sh
#!/usr/bin/env bash
set -e
echo "before"
out=$(false)
echo "after  out='$out'"
$ ./a.sh ; echo "rc=$?"
before
rc=1                      # "after" is never printed

$ cat b.sh
#!/usr/bin/env bash
set -e
echo "before"
out=$(bash -c 'echo boom >&2; exit 3' 2>/dev/null)
echo "after  out='$out'"
$ ./b.sh ; echo "rc=$?"
before
rc=3                      # redirecting stderr does not suppress the status

Decision table

Run on a 256-CPU DGX B300 (nproc=256). The squeue lookup is replaced by a stub
that reads its arguments; both the pre-fix and post-fix logic are copied
verbatim from run-parts.sh, set -e included.

case	squeue rc	output	after: exclusive	after: script continues	before: exclusive	before: script continues
1 failure, empty output	1	''	0	yes	–	no (aborts)
2 failure, non-empty output	1	'56 1'	0	yes	–	no (aborts)
3 success, empty output	0	''	0	yes	0	yes
4 success, non-numeric	0	'N/A N/A'	0	yes	0	yes
5 success, partial allocation	0	'1 1'	0	yes	0	yes
6 success, whole node	0	'256 1'	1	yes	1	yes

Both cases you asked about end with exclusive=0 and the script running on:
a non-zero squeue exit (rows 1–2) and an empty result (row 3). Rows 3–6 are
unchanged from the previous behaviour, so exclusive detection itself is not
regressed — row 6 is the one that shows a whole-node job is still detected.

Row 4 is not something squeue should produce; it is there because the guard is
[[ ... =~ ^[0-9]+$ ]] rather than a non-empty test, and I wanted the table to
cover that path too.

Harness

Included so the table can be reproduced rather than taken on trust.

#!/usr/bin/env bash
cd "$(dirname "$0")"
NPROC=$(grep -c ^processor /proc/cpuinfo)

make_fake () {   # $1=exit status  $2=%C value  $3=%D value
  cat > ./fake_squeue <<EOS
#!/usr/bin/env bash
rc=$1
case "\$*" in
  *"%C %D"*) printf '%s' "$2 $3" ;;
  *%C*)      printf '%s' "$2"     ;;
  *%D*)      printf '%s' "$3"     ;;
  *-u*)      printf '%s' ""       ;;
  *)         echo "unexpected args: \$*" >&2; exit 99 ;;
esac
exit \$rc
EOS
  chmod +x ./fake_squeue
}

The last *) branch matters: an earlier version of this harness ignored the
arguments and returned the same value for %C and %D, which silently produced
256 * 256 and made row 6 look like a regression that was not there. A stub that
answers questions it was not asked is worse than one that fails.

Full harness

#!/usr/bin/env bash
# Decision-table harness for PR #1391.
# The squeue stub reads its arguments. Unexpected args fail loudly instead of returning a value.
cd "$(dirname "$0")"
NPROC=$(grep -c ^processor /proc/cpuinfo)

make_fake () {   # $1=exit status  $2=%C value  $3=%D value
  cat > ./fake_squeue <<EOS
#!/usr/bin/env bash
rc=$1
case "\$*" in
  *"%C %D"*) printf '%s' "$2 $3" ;;     # after the fix: both at once
  *%C*)      printf '%s' "$2"     ;;     # before the fix: %C alone
  *%D*)      printf '%s' "$3"     ;;     # before the fix: %D alone
  *-u*)      printf '%s' ""       ;;     # user-jobs lookup
  *)         echo "unexpected args: \$*" >&2; exit 99 ;;
esac
exit \$rc
EOS
  chmod +x ./fake_squeue
}

AFTER='
set -e
log () { :; }
squeue_bin=./fake_squeue
SLURM_JOBID=123
exclusive=0
if job_alloc=$("$squeue_bin" -h -j "$SLURM_JOBID" -o "%C %D" 2>/dev/null); then
    read -r numcpus_job numnodes_job <<<"$job_alloc" || true
    if [[ "$numcpus_job" =~ ^[0-9]+$ ]] && [[ "$numnodes_job" =~ ^[0-9]+$ ]]; then
        numcpus_sys=$(( $(grep -c ^processor /proc/cpuinfo) * numnodes_job ))
        if [ "$numcpus_sys" -eq "$numcpus_job" ]; then exclusive=1; fi
    else
        log "[WARN] no usable allocation"
    fi
else
    log "[WARN] squeue failed"
fi
echo "$exclusive"
'

BEFORE='
set -e
squeue_bin=./fake_squeue
SLURM_JOBID=123
exclusive=0
numcpus_job=$("$squeue_bin" -h -j "$SLURM_JOBID" -o %C 2>/dev/null)
numnodes_job=$("$squeue_bin" -h -j "$SLURM_JOBID" -o %D 2>/dev/null)
numcpus_sys=$(( $(grep -c ^processor /proc/cpuinfo) * ${numnodes_job:-1} ))
if [ -n "$numcpus_job" ] && [ "$numcpus_sys" -eq "$numcpus_job" ] 2>/dev/null ; then
    exclusive=1
fi
echo "$exclusive"
'

one () {   # $1=snippet  -> "exclusive|continued"
  local out rc
  out=$(bash -c "$1" 2>/dev/null); rc=$?
  if [ $rc -ne 0 ]; then echo "-|no (aborts)"; else echo "$out|yes"; fi
}

row () {   # $1=label $2=rc $3=%C $4=%D
  make_fake "$2" "$3" "$4"
  local a b
  a=$(one "$AFTER"); b=$(one "$BEFORE")
  printf '%-20s | %-2s | %-7s | %-9s | %-12s | %-9s | %s\n' \
    "$1" "$2" "'$3 $4'" "${a%|*}" "${a#*|}" "${b%|*}" "${b#*|}"
}

echo "nproc=$NPROC"
printf '%-20s | %-2s | %-7s | %-9s | %-12s | %-9s | %s\n' \
  "case" "rc" "output" "after" "continues" "before" "continues"
echo "---------------------|----|---------|-----------|--------------|-----------|--------"
row "1 fail, empty"    1 ""     ""
row "2 fail, output"   1 "56"   "1"
row "3 ok, empty"      0 ""     ""
row "4 ok, non-numeric" 0 "N/A" "N/A"
row "5 ok, partial"    0 "1"    "1"
row "6 ok, whole node" 0 "$NPROC" "1"