Merge pull request #346 from gafusion/user_config_mem

Introduced user configurable memory specification to qsub
gafusion · Feb 6, 2024 · 5c697f4 · 5c697f4
2 parents 225949e + c5fa449
commit 5c697f4
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 2 deletions.
diff --git a/platform/qsub/qsub.OMEGA b/platform/qsub/qsub.OMEGA
@@ -6,7 +6,14 @@ echo "#SBATCH -o $SIMDIR/batch.out" >> $bfile
 echo "#SBATCH -e $SIMDIR/batch.err" >> $bfile
 echo "#SBATCH -t $WALLTIME" >> $bfile
 echo "#SBATCH -n $cores_used" >> $bfile
-echo "#SBATCH --mem 45G" >> $bfile
+
+if [ -n $MEMNODE ] 
+then
+  echo "#SBATCH --mem $MEMNODE" >> $bfile
+elif [ -n $MEMPERCPU ]
+  echo "#SBATCH --mem-per-cpu $MEMPERCPU" >> $bfile
+fi
+
 if [ "$QUEUE" = "null_queue" ]
 then
   echo "#SBATCH -p medium" >> $bfile

diff --git a/platform/qsub/qsub.PPPL b/platform/qsub/qsub.PPPL
@@ -6,11 +6,19 @@ echo "#SBATCH -o $SIMDIR/batch.out" >> $bfile
 echo "#SBATCH -e $SIMDIR/batch.err" >> $bfile
 echo "#SBATCH -t $WALLTIME" >> $bfile
 echo "#SBATCH -n $cores_used" >> $bfile
-echo "#SBATCH --mem=2GB" >> $bfile
+
+if [ -n $MEMNODE ] 
+then
+  echo "#SBATCH --mem $MEMNODE" >> $bfile
+elif [ -n $MEMPERCPU ]
+  echo "#SBATCH --mem-per-cpu $MEMPERCPU" >> $bfile
+fi
+
 if [ "$QUEUE" = "null_queue" ]
 then
   echo "#SBATCH -p general" >> $bfile
 else
   echo "#SBATCH -p $QUEUE" >> $bfile
 fi
+
 echo "$CODE -e $LOCDIR -n $nmpi -nomp $nomp -numa $numa -mpinuma $mpinuma -p $SIMROOT" >> $bfile
diff --git a/shared/bin/gacode_qsub b/shared/bin/gacode_qsub
@@ -44,6 +44,14 @@ then
   echo "         -mpinuma <n>" 
   echo "         MPI tasks per active NUMA."
   echo
+  echo "         -mem <n>[unit]" 
+  echo "         -mem-per-cpu <n>[unit]" 
+  echo "         Memory required per node or per cpu, resp."
+  echo "         Default unit MB. Slurm recognizes [K|M|G|T]"
+  echo "         No space beween number and unit."
+  echo "         Precedence is User then -mem"
+  echo "         (default --mem=16GB)."
+  echo
   echo "         -queue"
   echo "         Queue name"
   echo
@@ -95,6 +103,11 @@ nomp=1
 numa=0
 mpinuma=0
 
+#Default Memory Per Node
+MEMPERNODE=16G
+MEMPERCPU=
+USERMPN=false
+
 #=============================================================
 # Parse command line options
 #
@@ -117,6 +130,10 @@ while [[ $# -gt 0 ]] ; do
 
   -numa) shift ; numa=$1 ;;
 
+  -mem) shift ; MEMPERNODE=$1 ; USERMPN=true;;
+
+  -mem-per-cpu) shift ; MEMPERCPU=$1 ;;
+
   -mpinuma) shift ; mpinuma=$1 ;;
 
   -queue) shift ; QUEUE=$1 ;;
@@ -130,7 +147,18 @@ while [[ $# -gt 0 ]] ; do
   esac
   shift
 done
+
 #==============================================================
+# Resolve double specification of memory
+#
+if [[ $USERMPN == true && -n $MEMPERCPU ]]; then
+  echo "WARNING: Both -mem and -mem-per-cpu specified."
+  echo "Applying Precedence -mem=$MEMPERNODE."
+  MEMPERCPU=
+elif [[ $USERMPN == false ]]; then
+  MEMPERNODE=
+fi
+export MEMP
 
 #==============================================================
 # Manage paths