From a20594b197a46ec630c75e87ac29a8a919cfa9a0 Mon Sep 17 00:00:00 2001
From: Nicolas Castet <ncastet@nvidia.com>
Date: Thu, 8 Aug 2024 09:46:35 -0500
Subject: [PATCH] Replace costly call to 'scontrol show config' with pmix env
 vars

---
 conf/hooks/extra/50-slurm-pmi.sh | 20 ++++++++------------
 1 file changed, 8 insertions(+), 12 deletions(-)

diff --git a/conf/hooks/extra/50-slurm-pmi.sh b/conf/hooks/extra/50-slurm-pmi.sh
index 72c5ab7..2651f6f 100755
--- a/conf/hooks/extra/50-slurm-pmi.sh
+++ b/conf/hooks/extra/50-slurm-pmi.sh
@@ -29,14 +29,10 @@ done
 if [[ -z "${SLURM_MPI_TYPE-}" || "${SLURM_MPI_TYPE}" == pmix* ]] && compgen -e "PMIX_" > /dev/null; then
     source "${ENROOT_LIBRARY_PATH}/common.sh"
 
-    common::checkcmd scontrol awk
-
-    scontrol show config | awk '/^SlurmdSpoolDir|^TmpFS/ {print $3}' \
-      | { read -r slurm_spool; read -r slurm_tmpfs; } || :
-
-    if [ -z "${slurm_spool}" ] || [ -z "${slurm_tmpfs}" ]; then
-        common::err "Could not read SLURM configuration"
-    fi
+    # Precedence of vars documented at https://docs.openpmix.org/en/latest/how-things-work/session_dirs.html
+    default_tmpdir=${TMPDIR:-${TEMP:-${TMP:-/tmp}}}
+    system_tmpdir=${PMIX_SYSTEM_TMPDIR:-${default_tmpdir}}
+    server_tmpdir=${PMIX_SERVER_TMPDIR:-${default_tmpdir}}
 
     for var in $(compgen -e "PMIX_"); do
         printf "%s=%s\n" "${var}" "${!var}" >> "${ENROOT_ENVIRON}"
@@ -51,12 +47,12 @@ if [[ -z "${SLURM_MPI_TYPE-}" || "${SLURM_MPI_TYPE}" == pmix* ]] && compgen -e "
         printf "PMIX_MCA_gds=%s\n" ${PMIX_GDS_MODULE} >> "${ENROOT_ENVIRON}"
     fi
 
-    if [ -e "${slurm_tmpfs}/spmix_appdir_${SLURM_JOB_UID}_${SLURM_JOB_ID}.${SLURM_STEP_ID}" ]; then
-        printf "%s x-create=dir,bind,rw,nosuid,noexec,nodev,private,nofail\n" "${slurm_tmpfs}/spmix_appdir_${SLURM_JOB_UID}_${SLURM_JOB_ID}.${SLURM_STEP_ID}" >> "${ENROOT_MOUNTS}"
+    if [ -e "${system_tmpdir}/spmix_appdir_${SLURM_JOB_UID}_${SLURM_JOB_ID}.${SLURM_STEP_ID}" ]; then
+        printf "%s x-create=dir,bind,rw,nosuid,noexec,nodev,private,nofail\n" "${system_tmpdir}/spmix_appdir_${SLURM_JOB_UID}_${SLURM_JOB_ID}.${SLURM_STEP_ID}" >> "${ENROOT_MOUNTS}"
     else
-        printf "%s x-create=dir,bind,rw,nosuid,noexec,nodev,private,nofail\n" "${slurm_tmpfs}/spmix_appdir_${SLURM_JOB_ID}.${SLURM_STEP_ID}" >> "${ENROOT_MOUNTS}"
+        printf "%s x-create=dir,bind,rw,nosuid,noexec,nodev,private,nofail\n" "${system_tmpdir}/spmix_appdir_${SLURM_JOB_ID}.${SLURM_STEP_ID}" >> "${ENROOT_MOUNTS}"
     fi
-    printf "%s x-create=dir,bind,rw,nosuid,noexec,nodev,private\n" "${slurm_spool}/pmix.${SLURM_JOB_ID}.${SLURM_STEP_ID}" >> "${ENROOT_MOUNTS}"
+    printf "%s x-create=dir,bind,rw,nosuid,noexec,nodev,private\n" "${server_tmpdir}" >> "${ENROOT_MOUNTS}"
 fi
 
 # Check for PMI/PMI2 support.