fix: NUMA check

Author: milsanore

scripts/pin_cpus.sh

@@ -32,19 +32,43 @@ if ! [[ $CPU_SET_RANGE =~ ^[0-9,-]+$ ]]; then
     exit 1
 fi
 
-# --- minimal NUMA check ---
-# nodes_used=$(for cpu in $(echo "$CPU_SET_RANGE" | sed 's/,/ /g'); do
-#     for node_dir in /sys/devices/system/node/node*; do
-#         grep -qE "(^|,)$cpu($|,|-)" "$node_dir/cpulist" && echo $(basename $node_dir | sed 's/node//')
-#     done
-# done | sort -u)
-#
-# if [[ $(wc -w <<< "$nodes_used") -gt 1 ]]; then
-#     echo "error: requested CPUs span multiple NUMA nodes. nodes [$nodes_used]"
-#     exit 1
-# fi
-# echo "requested CPUs are within NUMA node. node [$nodes_used]"
-# --- end minimal NUMA check ---
+# --- NUMA check ---
+expand_cpu_range() {
+    local range=$1
+    local cpus=()
+    IFS=',' read -ra parts <<< "$range"
+    for part in "${parts[@]}"; do
+        if [[ $part == *-* ]]; then
+            IFS='-' read -r start end <<< "$part"
+            for ((i=start; i<=end; i++)); do
+                cpus+=("$i")
+            done
+        else
+            cpus+=("$part")
+        fi
+    done
+    echo "${cpus[@]}"
+}
+# determine NUMA nodes used by requested CPUs
+nodes_used=$(for cpu in $(expand_cpu_range "$CPU_SET_RANGE"); do
+    for node_dir in /sys/devices/system/node/node*; do
+        cpulist_file="$node_dir/cpulist"
+        [[ -f $cpulist_file ]] || continue
+        if grep -qE "(^|,)$cpu($|,|-)" "$cpulist_file"; then
+            echo "$(basename "$node_dir" | sed 's/node//')"
+        fi
+    done
+done | sort -u)
+if [[ -z "$nodes_used" ]]; then
+    echo "error: could not determine NUMA nodes. CPUs [$CPU_SET_RANGE]"
+    exit 1
+fi
+if [[ $(wc -w <<< "$nodes_used") -gt 1 ]]; then
+    echo "error: requested CPUs span multiple NUMA nodes. CPUs [$CPU_SET_RANGE] nodes [$nodes_used]"
+    exit 1
+fi
+echo "requested CPUs are within NUMA node. CPUs [$CPU_SET_RANGE] node [$nodes_used]"
+# --- end NUMA check ---
 
 # build path based on cgroup version
 if [[ -f /sys/fs/cgroup/cgroup.controllers ]]; then