run_proc.bash

#!/bin/bash -l

#=============================================================================================================================
#                         FUNCTIONS 
#=============================================================================================================================
moo_wait() {
# Limit the number of MASS retrievals that we submit in parallel. The recommendation is to keep
# this to around 30 max. 
let njobs=$(squeue --format '%.100j' --noheader --user $USER| grep "moo" | wc -l)
while (( ${njobs} > 25 ))
do
  let njobs=$(squeue --format '%.100j' --noheader --user $USER| grep "moo" | wc -l)
  echo "Number of MASS retrievals : $njobs. Sleeping."
  sleep 1m
done
}

slurm_wait() {
# Check that we aren't going to exceed user limit of number of jobs before we submit another one.
# limit for pending and running jobs is 1000, limit for running jobs is 500. Aim for under 500 in
# total to be on the safe side. Note some latency in the system (jobs take a while to show up in 
# queue) so set the limit a bit below 500. 
# Now that we have chunking of MASS restores, a lot of processing jobs can be submitted all at once
# after a chunk of files have been restored and might not all show up in the queue. So put in a short
# sleep between each job submission as well.  
sleep 1s
let njobs=$(squeue --format '%.100j' --noheader --user $USER  | wc -l)
while (( $njobs > 480 ))
do
  let njobs=$(squeue --format '%.100j' --noheader --user $USER | wc -l)
  echo "Number of slurm jobs : $njobs. Sleeping."
  sleep 2m
done
}

retrieve_data() {
   # $1 = $RUNID ; $2 = $FREQ ; $3 = $GRID ; $4+ = $TAGLIST
   exit_code=1
   let count=0
   while [[ "$exit_code" != "0" ]];do
      (( count += 1 ))
      slurm_wait
      sbatch --job-name=moo_${3}_${4} --output=${JOBOUT_PATH}/moo_${2}_${3}_${4}.out ${SCRPATH}/get_data.bash $1 $2 $3 ${@:4} | awk '{print $4}'
      exit_code=$?
      echo "Retrieval attempt $count. Exit code $exit_code." >> ${JOBOUT_PATH}/sbatch_moo_${2}_${3}_${4}.out
   done
}

run_tool() {
   # $1 = TOOL ; [possible flags]; $2 = $TAG ; $3 = $RUNID ; $4 = $FREQ ; $5+ = retrieval job IDs
   # global var njob
   local OPTARG OPTIND opt
   TOOL=$1;shift
   flags=""
   jobtag=""
   while getopts S:s:A:B opt ; do
     if [ "$opt" == "B" ];then  
        jobtag="${jobtag}_B" 
        flags="$flags -${opt}"
     else
        jobtag="${jobtag}_${OPTARG}" 
        flags="$flags -${opt} $OPTARG"
     fi     
   done
   shift `expr $OPTIND - 1`  
   # echo "run_tool running $TOOL $flags $1 $2 $3"
   sbatchschopt='--wait ' #--qos=long '  
   sbatchrunopt="--dependency=afterany:${tmasks_sbatch} --job-name=P$$_${TOOL}_${1}_${2} --output=${JOBOUT_PATH}/${TOOL}${jobtag}_${3}_${1}.out"
   if [[ -n "${@:4}" && "${@:4}" != ":" && "${@:4}" != "::" ]]; then
      sbatchrunopt=$(echo "$sbatchrunopt" | sed "s/--dependency=afterany:/--dependency=afterany:${@:4}:/")
   fi
   exit_code=1
   
   while [[ "$exit_code" != "0" ]];do
      first_tool_id="first_${TOOL}_ID"
      if ! declare -p "$first_tool_id" &>/dev/null; then
         declare -g "$first_tool_id"
         sbatch_output=$(sbatch ${sbatchschopt} ${sbatchrunopt} ${SCRPATH}/${TOOL}.bash ${flags} $2 $1 $3 | awk '{print $4}')
         exit_code=$?
         eval "$first_tool_id=$sbatch_output"
      else
         slurm_wait
         sbatchrunopt=$(echo "$sbatchrunopt" | sed "s/--dependency=afterany:/--dependency=afterany:${!first_tool_id}:/")
         sbatch ${sbatchschopt} ${sbatchrunopt} ${SCRPATH}/${TOOL}.bash ${flags} $2 $1 $3 > /dev/null 2>&1 &
         exit_code=$?
      fi
   done
   njob=$((njob+1))
}

progress_bar() {
   sleep 4
   echo''
   ijob=$njob
   eval "printf '|' ; printf '%0.s ' {0..100} ; printf '|\r' ;"
   while [[ $ijob -ne 0 ]] ; do
     ijob=$(squeue -u ${USER} | grep "P$$" | wc -l)
     icar=$(( ( (njob - ijob) * 100 ) / njob ))
     eval "printf '|' ; printf '%0.s=' {0..$icar} ; printf '\r' ; "
     sleep 1
   done
   eval "printf '|' ; printf '%0.s=' {0..100} ; printf '|\n' ;"
   echo ''
}
#=============================================================================================================================

while getopts B:C:V: opt ; do
  case $opt in
     B) BATHY=${OPTARG} ;;
     C) chunksize=${OPTARG} ;;
     V) ZCOMSHMSK=${OPTARG} ;;
  esac   
done
shift `expr $OPTIND - 1`  

if [[ -z "$chunksize" ]];then chunksize=1;fi
echo ""
echo "chunksize is $chunksize"
if [[ ! -z "$BATHY" ]]; then echo "BATHY is $BATHY"; fi
if [[ ! -z "$ZCOMSHMSK" ]];then
   echo "Vertical remapping is activated:" 
   echo "    the z-levels mesh_mask.nc is $ZCOMSHMSK"
fi
echo ""

if [ $# -le 4 ]; then echo 'run_proc.sh [-C chunksize] [-B BATHY] [-V ZCOMSHMSK] [MESHMASK] [YEARB] [YEARE] [FREQ] [RUNID list]'; exit 42; fi

MESHMASK=$1
YEARB=$2
YEARE=$3
FREQ=$4
RUNIDS=${@:5}

. param.bash

# clean ERROR.txt file
if [[ -f ERROR.txt ]]; then rm ERROR.txt ; fi

# Check if the mesh_mask.nc is using GVC
GVC=`ncdump -v ln_sco ${MSKPATH}/${MESHMASK} | sed -e "1,/data:/d" -e '$d' -e "s/ln_sco =//g" -e "s/;//g"`
GVC=$(($GVC + 0))

# Check which metrics we want to compute
runTRP=0
runBSF=0
runTS=0
runMOC=0
needBATHY=0
runVRMP=0
[[ $runACC == 1 || $runMargSea == 1 || $runITF == 1 || $runNAtlOverflows == 1  || $runArcTrans == 1 ]] && runTRP=1
[[ $runBSF_SO == 1 || $runBSF_NA == 1 ]] && runBSF=1
[[ $runDEEPTS == 1 || $runSSS_LabSea == 1 || $runSST_SO == 1 || $runSST_NWCorner == 1 ]] && runTS=1
[[ $runAMOC == 1 || $runRAPID == 1 || $runOSNAP == 1 ]] && runMOC=1
[[ $runAMOC == 1 && $GVC == 1 ]] && export runVRMP=1
[[ $runTRP == 1 || $runMOC == 1 || $runVRMP == 1 ]] && needBATHY=1

# Check that everything is in place
if [[ ! -d ${CDFPATH} || ! -n "$(ls -A "$CDFPATH")" ]] ; then  
   echo "E R R O R : CDFTOOLS/bin directory does not exist or is empty: ${CDFPATH}"
   exit 41
fi

if [[ ! -f ${NMLPATH} ]] ; then 
   echo "E R R O R : nam_cdf_names namelist file does not exist : ${NMLPATH}"
   exit 41
fi

if [[ ! -f ${MSKPATH}/${MESHMASK} ]] ; then 
   echo "E R R O R : mesh_mask file does not exist : ${MSKPATH}/${MESHMASK}"
   exit 41
fi

if [[ $needBATHY == 1 && ! -f ${MSKPATH}/${BATHY} ]] ; then
   echo "E R R O R : bathymetry file does not exist : ${MSKPATH}/${BATHY}"
   echo "Bathymetry file can be created from mesh_mask file using SCRIPT/bathy_from_dommesh.py"
   exit 41
fi

if [[ $runVRMP == 1 && ! -f ${ZCOMSHMSK} ]] ; then
   echo "E R R O R : the z-levels mesh_mask.nc file for vertical rempping does not exist : ${TARGETMSHMSK}"
   exit 41
fi

if [[ $needBATHY == 1 && $runVRMP == 1 ]] ; then
   LVMSK=0
   for var in "s2z_msk" "mask_loczgr"; do
       lvmsk=`ncdump -h ${MSKPATH}/${BATHY} | grep ${var}`
       if [[ ! -z "${lvmsk}" ]]; then LVMSK=$(($LVMSK + 1)); fi
   done
   if [[ $LVMSK == 0 ]]; then
      echo "E R R O R : you want to carry out vertical remapping but the "
      echo "            ${MSKPATH}/${BATHY} file"
      echo "            that you are passing in input does not include any of"
      echo "            the following two variables: s2z_msk or mask_loczgr, please check!!"
      exit 41 
   fi
fi

# loop over years
echo ''
for RUNID in `echo $RUNIDS`; do

   # set up jobout directory file
   JOBOUT_PATH=${EXEPATH}/JOBOUT/${RUNID}
   if [ ! -d ${JOBOUT_PATH} ]; then mkdir -p ${JOBOUT_PATH} ; fi

   # create working directory
   if [ ! -d ${DATPATH}/${RUNID} ]; then mkdir -p ${DATPATH}/${RUNID} ; fi
   cd ${DATPATH}/${RUNID}

   # Remove old links to nam_cdf_names and mesh_mask files and re-link
   # to be sure we are using up-to-date versions.

   # check nam_cdf_names namelist
   if [[ -L nam_cdf_names ]] ;then rm nam_cdf_names; fi
   ln -s ${NMLPATH} nam_cdf_names 

   # check geometry files
   if [[ -L mesh.nc     ]] ; then rm mesh.nc; fi
   ln -s ${MSKPATH}/${MESHMASK} mesh.nc
   if [[ -L mask.nc     ]] ; then rm mask.nc ; fi
   ln -s ${MSKPATH}/${MESHMASK} mask.nc 
   if [[ $needBATHY == 1 ]] ; then
       if [[ -L bathy.nc  ]] ; then rm bathy.nc ; fi
       ln -s ${MSKPATH}/${BATHY} bathy.nc 
   fi
   if [[ $runVRMP == 1 ]] ; then
      if [[ -L zco_mesh.nc ]] ; then rm zco_mesh.nc ; fi
      ln -s ${ZCOMSHMSK} zco_mesh.nc
   fi
         
   # subbasins file not currently used by any metrics
   #if [ -L subbasin.nc ] ; then rm subbasin.nc ; fi
   #ln -s ${MSKPATH}/subbasins_${CONFIG}-GO6.nc subbasin.nc 

   echo "================================================================="
   echo "Computing the requested diagnostics for the $RUNID simulation ..."
   echo "================================================================="
   echo ""

   # Conservative vertical remapping if needed
   if [[ $runVRMP == 1 ]]; then
      echo ""
      echo "Computing z-level domain for conservative vertical remapping ..."
      echo ""
      msh4vmap_sbatch=$(sbatch --output=${JOBOUT_PATH}/mk_msh4vmap.out ${SCRPATH}/mk_msh4vmap.bash $RUNID | awk '{print $4}')
   fi

   # Generate subbasin mask for the Atlantic if needed
   if [[ $runAMOC == 1 ]]; then
      subbasin_sbatch=$(sbatch --output=${JOBOUT_PATH}/mk_subbasin.out ${SCRPATH}/mk_subbasin.bash $RUNID | awk '{print $4}')
   fi

   njob=0
   LSTY=$(eval echo {${YEARB}..${YEARE}})
   if   [[ $FREQ == 1m ]]; then MONTHB=1  ; MONTHE=12 ; LSTM=$(eval echo {$MONTHB..$MONTHE}) ;
   elif [[ $FREQ == 1y ]]; then MONTHB=12 ; MONTHE=12 ; LSTM=$(eval echo {$MONTHB..$MONTHE}) ;
   else 
        echo "E R R O R : $FREQ not supported; exit 42"
        exit 42
   fi

   # generate tmasks
   tmasks_sbatch=$(sbatch --output=${JOBOUT_PATH}/mk_msks.out ${SCRPATH}/mk_msks.bash $RUNID | awk '{print $4}')

   let tagcount=0
   let tag2count=0
   TAG_LIST=""
   TAGDJF_LIST=""
   TAG09_LIST=""
   TAG02_LIST=""
   TAG03_LIST=""
   for YEAR in $(printf "%04d " $LSTY); do

      for MONTH in $(printf "%02d " $LSTM); do
         # define tags
         (( tagcount+=1 ))
         TAG_LIST="$TAG_LIST ${YEAR}${MONTH}01"

         if [[ $tagcount == $chunksize || ( $YEAR == $YEARE && $MONTH == $MONTHE ) ]]
         then

         echo "TAG_LIST : $TAG_LIST"

         if [[ $runRETRIEVE_DATA == 1 ]]; then
            # get data (retrieve_data function is defined in this script)
            moo_wait
            [[ $runTRP == 1 || $runBSF == 1 || $runMOC == 1 || $runMHT == 1  ]] && mooVyid=$(retrieve_data $RUNID $FREQ grid-V $TAG_LIST)
            moo_wait
            [[ $runTRP == 1 || $runBSF == 1 || $runMOC == 1 ]] && mooUyid=$(retrieve_data $RUNID $FREQ grid-U $TAG_LIST)
            moo_wait
            [[ $runTRP == 1 || $runQHF == 1 || $runTS == 1 || $runMOC == 1 || $runHTC == 1 || $runSTC == 1 || $runGSL_NAC == 1 || $runMHT == 1 || $runQHF == 1 || $runMedOVF == 1 || $runAABW == 1 ]] && mooTyid=$(retrieve_data $RUNID $FREQ grid-T $TAG_LIST)
         fi

         echo "mooTyid : $mooTyid"
         echo "mooUyid : $mooUyid"
         echo "mooVyid : $mooVyid"

         for TAG in $TAG_LIST; do
            # run cdftools

            if [[ $runVRMP == 1 ]]; then
               jdepe="--dependency=afterany:$mooVyid:$mooUyid:$mooTyid:$msh4vmap_sbatch"
               jname="--job-name=mk_vrmp_${TAG}_${RUNID}"
               joutp="--output=${JOBOUT_PATH}/mk_vrmp_${FREQ}_${TAG}.out"
               vrmp_sbatch=$(sbatch $jdepe $jname $joutp ${SCRPATH}/mk_vrmp.bash $RUNID $TAG $FREQ | awk '{print $4}') 
            fi

            [[ $runBSF_SO == 1 ]]  && run_tool mk_psi_SO                 $TAG $RUNID $FREQ $mooVyid:$mooUyid
            [[ $runDEEPTS == 1 ]]  && run_tool mk_deepTS -A AMU          $TAG $RUNID $FREQ $mooTyid
            [[ $runDEEPTS == 1 ]]  && run_tool mk_deepTS -A WROSS        $TAG $RUNID $FREQ $mooTyid
            [[ $runSST_SO == 1 ]]  && run_tool mk_sst_SO                 $TAG $RUNID $FREQ $mooTyid
            [[ $runACC == 1 ]]     && run_tool mk_trp  -S ACC            $TAG $RUNID $FREQ $mooVyid:$mooUyid:$mooTyid
            [[ $runACC == 1 ]]     && run_tool mk_trp  -S ACC -B         $TAG $RUNID $FREQ $mooVyid:$mooUyid:$mooTyid
            [[ $runACC == 1 ]]     && run_tool mk_trp  -s ACC-shelfbreak $TAG $RUNID $FREQ $mooVyid:$mooUyid:$mooTyid
            [[ $runNAtlOverflows == 1 ]] && run_tool mk_trp  -S DenmarkStrait      $TAG $RUNID $FREQ $mooVyid:$mooUyid:$mooTyid
            [[ $runNAtlOverflows == 1 ]] && run_tool mk_trp  -S FaroeBankChannel   $TAG $RUNID $FREQ $mooVyid:$mooUyid:$mooTyid
            [[ $runArcTrans == 1 ]]      && run_tool mk_trp  -S FramStrait         $TAG $RUNID $FREQ $mooVyid:$mooUyid:$mooTyid
            [[ $runArcTrans == 1 ]]      && run_tool mk_trp  -S BeringStrait       $TAG $RUNID $FREQ $mooVyid:$mooUyid:$mooTyid
            [[ $runArcTrans == 1 ]]      && run_tool mk_trp  -S DavisStrait        $TAG $RUNID $FREQ $mooVyid:$mooUyid:$mooTyid
            [[ $runArcTrans == 1 ]]      && run_tool mk_trp  -S BarentsSea         $TAG $RUNID $FREQ $mooVyid:$mooUyid:$mooTyid
            [[ $runArcTrans == 1 ]]      && run_tool mk_trp  -S WSC                $TAG $RUNID $FREQ $mooVyid:$mooUyid:$mooTyid
            [[ $runMargSea == 1 ]]       && run_tool mk_trp  -S GibraltarStrait    $TAG $RUNID $FREQ $mooVyid:$mooUyid:$mooTyid
            [[ $runMargSea == 1 ]]       && run_tool mk_trp  -S BabElMandeb        $TAG $RUNID $FREQ $mooVyid:$mooUyid:$mooTyid
            [[ $runMargSea == 1 ]]       && run_tool mk_trp  -S StraitOfHormuz     $TAG $RUNID $FREQ $mooVyid:$mooUyid:$mooTyid 
            [[ $runITF == 1 ]]           && run_tool mk_trp  -S LombokStrait       $TAG $RUNID $FREQ $mooVyid:$mooUyid:$mooTyid
            [[ $runITF == 1 ]]           && run_tool mk_trp  -S OmbaiStrait        $TAG $RUNID $FREQ $mooVyid:$mooUyid:$mooTyid
            [[ $runITF == 1 ]]           && run_tool mk_trp  -S TimorPassage       $TAG $RUNID $FREQ $mooVyid:$mooUyid:$mooTyid 
            [[ $runBSF_NA == 1 ]]        && run_tool mk_psi_NA    $TAG $RUNID $FREQ $mooVyid:$mooUyid
            if [[ $runVRMP == 1 ]]; then
               [[ $runAMOC == 1 ]] && run_tool mk_amoc $TAG $RUNID $FREQ $mooVyid:$mooUyid:$mooTyid:$vrmp_sbatch
            else
               [[ $runAMOC == 1 ]] && run_tool mk_amoc $TAG $RUNID $FREQ $mooVyid:$mooUyid:$mooTyid
            fi
            [[ $runRAPID == 1 ]]         && run_tool mk_rapid_moc $TAG $RUNID $FREQ $mooVyid:$mooUyid:$mooTyid
            [[ $runOSNAP == 1 ]]         && run_tool mk_osnap_moc $TAG $RUNID $FREQ $mooVyid:$mooUyid:$mooTyid
            [[ $runSST_NWCorner == 1 ]]  && run_tool mk_sst_NA    $TAG $RUNID $FREQ $mooTyid
            [[ $runSSS_LabSea == 1 ]]    && run_tool mk_sss       $TAG $RUNID $FREQ $mooTyid
            [[ $runHTC == 1 ]]           && run_tool mk_htc       $TAG $RUNID $FREQ $mooTyid
            [[ $runSTC == 1 ]]           && run_tool mk_stc       $TAG $RUNID $FREQ $mooTyid
            [[ $runGSL_NAC == 1 ]]       && run_tool mk_gsl_nac   $TAG $RUNID $FREQ $mooTyid
            #[[ $runOVF == 1 ]]           && run_tool mk_ovf       $TAG $RUNID $FREQ $mooTyid
            [[ $runMHT == 1 ]]     && run_tool mk_mht  $TAG $RUNID $FREQ $mooVyid:$mooTyid
            [[ $runQHF == 1 ]]     && run_tool mk_hfds $TAG $RUNID $FREQ $mooTyid 
            [[ $runMedOVF == 1 ]]   && run_tool mk_medovf $TAG $RUNID $FREQ $mooTyid
            [[ $runAABW == 1 ]] && run_tool mk_aabw $TAG $RUNID $FREQ $mooTyid
         done
         let tagcount=0
         TAG_LIST=""
         fi
      done

      # define tags      
      (( tag2count+=1 ))
      TAGDJF_LIST="$TAGDJF_LIST ${YEAR}1201"
      TAG09_LIST="$TAG09_LIST ${YEAR}0901"
      TAG02_LIST="$TAG02_LIST ${YEAR}0201"
      TAG03_LIST="$TAG03_LIST ${YEAR}0301"

      if [[ (( $tag2count = $chunksize || $YEAR = $YEARE )) ]]
      then

      echo "TAGDJF_LIST : $TAGDJF_LIST"
      echo "TAG09_LIST : $TAG09_LIST"
      echo "TAG03_LIST : $TAG03_LIST"

      # get data (retrieve_data function are defined in this script)
      if [[ $runRETRIEVE_DATA == 1 ]]; then
         moo_wait
         [[ $runDEEPTS == 1 ]]                                              && mooDJFsid=$(retrieve_data $RUNID 1s grid-T $TAGDJF_LIST)
         moo_wait
         [[ $runSIE == 1 || $runMLD_Weddell == 1 ]]                         && mooT09mid=$(retrieve_data $RUNID 1m grid-T $TAG09_LIST)
         moo_wait
         [[ $runSIE == 1 ]]                                                 && mooT02mid=$(retrieve_data $RUNID 1m grid-T $TAG02_LIST)
         moo_wait
         [[ $runSIE == 1 || $runMLD_LabSea == 1 ]]                          && mooT03mid=$(retrieve_data $RUNID 1m grid-T $TAG03_LIST)
      fi
      
      echo "mooDJFsid : $mooDJFsid"
      echo "mooT09mid : $mooT09mid"
      echo "mooT03mid : $mooT03mid"

      # run cdftools
      for TAG in $TAGDJF_LIST;do
         [[ $runDEEPTS == 1 ]]  && run_tool mk_deepTS -A WWED $TAG $RUNID 1s $mooDJFsid
         [[ $runDEEPTS == 1 ]]  && run_tool mk_deepTS -A EROSS $TAG $RUNID 1s $mooDJFsid
      done
      for TAG in $TAG09_LIST;do
         [[ $runMLD_Weddell == 1 ]] && run_tool mk_mxl_SO  $TAG $RUNID 1m    $mooT09mid
         [[ $runSIE == 1 ]]         && run_tool mk_sie  $TAG $RUNID 1m    $mooT09mid 
      done
      for TAG in $TAG02_LIST;do
         [[ $runSIE == 1 ]] && run_tool mk_sie  $TAG $RUNID 1m    $mooT02mid
      done
      for TAG in $TAG03_LIST;do
         [[ $runMLD_LabSea == 1 ]] && run_tool mk_mxl_NA  $TAG $RUNID 1m    $mooT03mid
         [[ $runSIE == 1 ]] && run_tool mk_sie  $TAG $RUNID 1m    $mooT03mid
      done
      let tag2count=0
      TAGDJF_LIST=""
      TAG09_LIST=""
      TAG02_LIST=""
      TAG03_LIST=""

      fi
   done

   # print task bar
   progress_bar  
 
   # wait it is all done (probably useless because of the progress bar loop)
   wait

done # end runids

rm -f ${SCRPATH}/.obs_done* # remove obs_done flag file if it exists

# print out
sleep 1
ls > /dev/null 2>&1 # without this the following command sometimes failed (maybe it force to flush all the file on disk)
if [ -f ERROR.txt ]; then
   echo ""
   echo "ERRORS are present :"
   cat ERROR.txt
   echo ""
   echo "if error expected (as missing data because data coverage larger than run coverage), diagnostics will be missing for these cases."
else
   echo ""
   echo "data processing is done for ${RUNIDS} between ${YEARB} and ${YEARE}"
fi
echo ""
echo "You can now run < ./run_plot_{PACKAGE}.bash [KEY] [FREQ] [RUNIDS] > if no more files to process (other resolution, other periods ...)"
echo ""
echo "by default ./run_plot_{PACKAGE}.bash will process all the file in the data directory, if you want some specific period, you need to tune the glob.glob pattern in the script"
echo ""