Merge branch 'main' into prolog-epilog-skip

Author: ekouts

checks/apps/cp2k/cp2k_uenv.py

@@ -15,11 +15,11 @@
 
 cp2k_references = {
     'md': {
-        'gh200': {'time_run': (69, None, 0.05, 's')},
+        'gh200': {'time_run': (45, None, 0.05, 's')},
         'zen2': {'time_run': (91, None, 0.05, 's')}
     },
     'pbe': {
-        'gh200': {'time_run': (67, None, 0.05, 's')},
+        'gh200': {'time_run': (50, None, 0.05, 's')},
         'zen2': {'time_run': (68, None, 0.05, 's')}
     },
     'rpa': {

checks/apps/cp2k/src/mps-wrapper.sh

@@ -1,19 +1,33 @@
 #!/bin/bash
 
-set -u
+set -eu
 
-export CUDA_MPS_PIPE_DIRECTORY=/tmp/nvidia-mps
-export CUDA_MPS_LOG_DIRECTORY=/tmp/nvidia-log
-export CUDA_VISIBLE_DEVICES=$(( SLURM_LOCALID % 4 ))
+export CUDA_DEVICE_MAX_COPY_CONNECTIONS=8
+export CUDA_DEVICE_MAX_CONNECTIONS=8
 
-if [ "${SLURM_LOCALID}" -eq 0 ]; then
-    CUDA_VISIBLE_DEVICES=0,1,2,3 nvidia-cuda-mps-control -d
+export CUDA_MPS_PIPE_DIRECTORY=/tmp/nvidia-mps-$((SLURM_LOCALID % 4))
+export CUDA_MPS_LOG_DIRECTORY=/tmp/nvidia-log-$((SLURM_LOCALID % 4))-$(id -un)
+
+export HWLOC_KEEP_NVIDIA_GPU_NUMA_NODES=0
+numa_nodes=$(hwloc-calc --physical --intersect NUMAnode $(hwloc-bind --get --taskset)) # do not set CUDA_VISIBLE_DEVICES, enough to set it for the daemon
+
+# Launch MPS from a single rank per GPU
+if [[ $SLURM_LOCALID -lt 4 ]]; then
+    mkdir -p ${CUDA_MPS_PIPE_DIRECTORY}
+    mkdir -p ${CUDA_MPS_LOG_DIRECTORY}
+    CUDA_VISIBLE_DEVICES=$((SLURM_LOCALID % 4)) nvidia-cuda-mps-control -d
 fi
 
-sleep 5
+# Wait for MPS to start
+sleep 1
 
-exec "$@"
+# Run the command
+"$@"
+result=$?
 
-if [ "${SLURM_LOCALID}" -eq 0 ]; then
+# Quit MPS control daemon before exiting
+if [[ $SLURM_LOCALID -lt 4 ]]; then
     echo quit | nvidia-cuda-mps-control
 fi
+
+exit $result