slurm: batch the queries and parallelize the writes in the GPU power/clock helpers

roles/slurm/templates/etc/slurm/shared/bin/set_gpu_clocks.sh

@@ -1,20 +1,70 @@
 #!/usr/bin/env bash
 set -e
 
-gpu_count="$(nvidia-smi -L | wc -l)"
-
 case "$1" in
     default)
         nvidia-smi -rac           # Reset application clocks
         nvidia-smi -acp 0         # Reset application clock permissions
         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]"

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"