diff --git a/roles/slurm/templates/etc/slurm/shared/bin/set_gpu_clocks.sh b/roles/slurm/templates/etc/slurm/shared/bin/set_gpu_clocks.sh index 13918792e..d171abf2f 100755 --- a/roles/slurm/templates/etc/slurm/shared/bin/set_gpu_clocks.sh +++ b/roles/slurm/templates/etc/slurm/shared/bin/set_gpu_clocks.sh @@ -1,8 +1,6 @@ #!/usr/bin/env bash set -e -gpu_count="$(nvidia-smi -L | wc -l)" - case "$1" in default) nvidia-smi -rac # Reset application clocks @@ -10,11 +8,63 @@ case "$1" in nvidia-smi -c DEFAULT # Reset compute mode to default ;; max) - for i in $(seq 0 "$(( gpu_count - 1 ))" ) ; do - nextSM="$(nvidia-smi -i "$i" --query-gpu=clocks.max.sm --format=csv,noheader,nounits)" - nextMEM="$(nvidia-smi -i "$i" --query-gpu=clocks.max.mem --format=csv,noheader,nounits)" - nvidia-smi -i "${i}" -ac "${nextMEM}","${nextSM}" + # Query every GPU in a single call, then apply in parallel: + # "nvidia-smi -ac" also costs roughly a second per GPU. + # + # Index, memory clock and SM clock come from one query so the three + # values of a GPU cannot drift out of alignment. Reading them through + # "readarray -t maxSM < <(nvidia-smi ...)" hid the query's exit status + # from both readarray and "set -e": a failed query left the arrays + # empty and the helper exited 0 having set nothing, and two separate + # queries returning different row counts made "${maxMEM[$i]}" empty for + # the trailing GPUs, producing an "-ac ,1980" argument. + tmp=$(mktemp) + trap 'rm -f "$tmp"' EXIT + + if ! nvidia-smi --query-gpu=index,clocks.max.mem,clocks.max.sm \ + --format=csv,noheader,nounits > "$tmp"; then + echo "$0: querying the maximum clocks failed" >&2 + exit 1 + fi + + # Every GPU that nvidia-smi lists has to be present in the query result + # before anything is written. + if ! expected=$(nvidia-smi -L | grep -c '^GPU '); then + echo "$0: could not determine the number of GPUs" >&2 + exit 1 + fi + + indexes=() + maxMEM=() + maxSM=() + while IFS=', ' read -r index mem sm _; do + [ -z "$index" ] && continue + if ! [[ "$index" =~ ^[0-9]+$ ]] \ + || ! [[ "$mem" =~ ^[0-9]+$ ]] \ + || ! [[ "$sm" =~ ^[0-9]+$ ]]; then + echo "$0: unexpected row from nvidia-smi: '$index, $mem, $sm'" >&2 + exit 1 + fi + indexes+=("$index") + maxMEM+=("$mem") + maxSM+=("$sm") + done < "$tmp" + + if [ "${#indexes[@]}" -ne "$expected" ]; then + echo "$0: got ${#indexes[@]} usable clock rows, expected $expected" >&2 + exit 1 + fi + + pids=() + for i in "${!indexes[@]}" ; do + nvidia-smi -i "${indexes[$i]}" -ac "${maxMEM[$i]}","${maxSM[$i]}" >/dev/null & + pids+=("$!") + done + rc=0 + for pid in "${pids[@]}" ; do + wait "$pid" || rc=1 done + exit "$rc" ;; *) echo "Usage: $0 [default|max]" diff --git a/roles/slurm/templates/etc/slurm/shared/bin/set_gpu_power_levels.sh b/roles/slurm/templates/etc/slurm/shared/bin/set_gpu_power_levels.sh index 4b3cb54e5..da140c160 100755 --- a/roles/slurm/templates/etc/slurm/shared/bin/set_gpu_power_levels.sh +++ b/roles/slurm/templates/etc/slurm/shared/bin/set_gpu_power_levels.sh @@ -1,24 +1,79 @@ #!/usr/bin/env bash set -e -gpu_count="$(nvidia-smi -L | wc -l)" +case "$1" in + max) + query=power.max_limit + ;; + default) + query=power.default_limit + ;; + min) + query=power.min_limit + ;; + *) + echo "Usage: $0 [max,default,min]" + exit 1 + ;; +esac -for i in $(seq 0 "$(( gpu_count - 1 ))" ) +# Query every GPU in a single call instead of one call per GPU. +# +# The result is written to a file rather than read through +# "readarray -t limits < <(nvidia-smi ...)": inside a process substitution the +# query's exit status is invisible to both readarray and "set -e", so a failed +# query left the array empty, the write loop ran zero times, and the helper +# still exited 0 having configured nothing. A truncated result configured only +# some of the GPUs and also returned success. +# +# The index is queried alongside the value so each write targets the GPU that +# nvidia-smi actually reported, rather than assuming the array subscript equals +# the GPU index. +tmp=$(mktemp) +trap 'rm -f "$tmp"' EXIT + +if ! nvidia-smi --query-gpu=index,"$query" --format=csv,noheader,nounits > "$tmp"; then + echo "$0: querying $query failed" >&2 + exit 1 +fi + +# Every GPU that nvidia-smi lists has to be present in the query result before +# anything is written. +if ! expected=$(nvidia-smi -L | grep -c '^GPU '); then + echo "$0: could not determine the number of GPUs" >&2 + exit 1 +fi + +indexes=() +limits=() +while IFS=', ' read -r index limit _; do + [ -z "$index" ] && continue + if ! [[ "$index" =~ ^[0-9]+$ ]] || ! [[ "$limit" =~ ^[0-9]+(\.[0-9]+)?$ ]]; then + echo "$0: unexpected row from nvidia-smi: '$index, $limit'" >&2 + exit 1 + fi + indexes+=("$index") + limits+=("$limit") +done < "$tmp" + +if [ "${#indexes[@]}" -ne "$expected" ]; then + echo "$0: got ${#indexes[@]} usable rows for $query, expected $expected" >&2 + exit 1 +fi + +# "nvidia-smi -pl" takes roughly a second per GPU, so applying the limits +# serially adds ~8 s to the prolog of every full-node job on an 8-GPU node. +# Apply them in parallel and collect the exit status of each child. +pids=() +for i in "${!indexes[@]}" +do + nvidia-smi -i "${indexes[$i]}" -pl "${limits[$i]}" >/dev/null & + pids+=("$!") +done + +rc=0 +for pid in "${pids[@]}" do - case "$1" in - max) - next="$(nvidia-smi -i "$i" --query-gpu=power.max_limit --format=csv,noheader,nounits)" - ;; - default) - next="$(nvidia-smi -i "$i" --query-gpu=power.default_limit --format=csv,noheader,nounits)" - ;; - min) - next="$(nvidia-smi -i "$i" --query-gpu=power.min_limit --format=csv,noheader,nounits)" - ;; - *) - echo "Usage: $0 [max,default,min]" - exit 1 - ;; - esac - nvidia-smi -i "$i" -pl "$next" + wait "$pid" || rc=1 done +exit "$rc"