Update atlas-run runner: remove old aprun command and update slurm (#251

)
ecmwf · Dec 19, 2024 · 47b31e7 · 47b31e7
1 parent 85196ae
commit 47b31e7
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 24 deletions.
diff --git a/cmake/features/ATLAS_RUN.cmake b/cmake/features/ATLAS_RUN.cmake
@@ -6,6 +6,12 @@ ecbuild_add_option( FEATURE ATLAS_RUN
 if( HAVE_ATLAS_RUN )
   set( MPIEXEC_EXECUTABLE ${CMAKE_CURRENT_SOURCE_DIR}/tools/atlas-run )
   set( MPIEXEC_NUMPROC_FLAG='-n' )
-  set( CMAKE_CROSSCOMPILING_EMULATOR ${CMAKE_CURRENT_SOURCE_DIR}/tools/atlas-run )
+  set( MPIEXEC_NUMTHREAD_FLAG='-c' )
+  if( CMAKE_VERSION VERSION_LESS 3.29 )
+    set( CMAKE_CROSSCOMPILING_EMULATOR ${CMAKE_CURRENT_SOURCE_DIR}/tools/atlas-run )
+  else()
+    unset(CMAKE_CROSSCOMPILING_EMULATOR)
+    set( CMAKE_TEST_LAUNCHER ${CMAKE_CURRENT_SOURCE_DIR}/tools/atlas-run )
+  endif()
 endif()
 
diff --git a/tools/atlas-run b/tools/atlas-run
@@ -23,6 +23,11 @@ while test $# -gt 0; do
     shift # past argument
     shift # past value
     ;;
+    -c)
+    ATLAS_RUN_NTHREADS="$2"
+    shift # past argument
+    shift # past value
+    ;;
     *)    # unknown option
     break
     ;;
@@ -38,48 +43,33 @@ if [ -n "${ATLAS_RUN_NTHREADS}" ]; then
   export OMP_NUM_THREADS=${ATLAS_RUN_NTHREADS}
 fi
 
-if command_exists aprun            && \
-  [[ ${EC_queue:-unset} != "nf" ]] && \
-  [[ ${EC_queue:-unset} != "ns" ]]
-then
-  LAUNCH="aprun ${ATLAS_RUN_MPI_ARGS}"
+if command_exists srun ; then
+  LAUNCH="srun ${ATLAS_RUN_MPI_ARGS}"
   if [ -z "${ATLAS_RUN_NPROCS}" ]; then
-     LAUNCH="${LAUNCH} -n 1"
+    LAUNCH="${LAUNCH} -n 1"
+    if [ -n "${SLURM_GPUS}" ]; then
+      ATLAS_RUN_NGPUS=1
+    fi
   else
     LAUNCH="${LAUNCH} -n ${ATLAS_RUN_NPROCS}"
   fi
   if [ -n "${OMP_NUM_THREADS}" ]; then
-    LAUNCH="${LAUNCH} -d ${OMP_NUM_THREADS}"
+    LAUNCH="${LAUNCH} -c ${OMP_NUM_THREADS}"
   fi
-
-elif command_exists srun ; then
-  LAUNCH="srun ${ATLAS_RUN_MPI_ARGS}"
   if [ -n "${ATLAS_RUN_NGPUS}" ]; then
     LAUNCH="${LAUNCH} --gpus-per-task=${ATLAS_RUN_NGPUS}"
   fi
-  if [ -z "${ATLAS_RUN_NPROCS}" ]; then
-    LAUNCH="${LAUNCH} -n 1"
-    if [ -n "${SLURM_GPUS}" ]; then
-      LAUNCH="${LAUNCH} --gpus-per-task=1"
-    fi
-  else
-    LAUNCH="${LAUNCH} -n ${ATLAS_RUN_NPROCS}"
-  fi
   SLURM_EXPORT_ENV=ALL # Required to propagate environment variables to srun'd program
 
 else
-  if [[ ${ARCH:-unset} == "cray" ]]; then
-    echo + module load cray-snplauncher
-    module load cray-snplauncher
-  fi
   if [ -z "${ATLAS_RUN_NPROCS}" ]; then
     unset LAUNCH
   elif command_exists mpirun ; then
     LAUNCH="mpirun ${ATLAS_RUN_MPI_ARGS} -np ${ATLAS_RUN_NPROCS}"
   elif command_exists mpiexec; then
     LAUNCH="mpiexec ${ATLAS_RUN_MPI_ARGS} -n ${ATLAS_RUN_NPROCS}"
   else
-    echo "No MPI driver found (mpirun,mpiexec,aprun,srun)"
+    echo "No MPI driver found (mpirun,mpiexec,srun)"
     exit 1
   fi
 fi