HERCULES.env

#! /usr/bin/env bash

if [[ $# -ne 1 ]]; then

    echo "Must specify an input argument to set runtime environment variables!"
    exit 1

fi

step=$1

export launcher="srun -l --export=ALL --hint=nomultithread"
export mpmd_opt="--multi-prog --output=mpmd.%j.%t.out"

# Configure MPI environment
export MPI_BUFS_PER_PROC=2048
export MPI_BUFS_PER_HOST=2048
export MPI_GROUP_MAX=256
export MPI_MEMMAP_OFF=1
export MP_STDOUTMODE="ORDERED"
export KMP_AFFINITY=scatter
export OMP_STACKSIZE=2048000
export NTHSTACK=1024000000
#export LD_BIND_NOW=1
export I_MPI_EXTRA_FILESYSTEM=1
export I_MPI_EXTRA_FILESYSTEM_LIST=lustre

ulimit -s unlimited
ulimit -a

# Calculate common variables
# Check first if the dependent variables are set
if [[ -n "${ntasks:-}" && -n "${max_tasks_per_node:-}" && -n "${tasks_per_node:-}" ]]; then
    max_threads_per_task=$((max_tasks_per_node / tasks_per_node))
    NTHREADSmax=${threads_per_task:-${max_threads_per_task}}
    NTHREADS1=${threads_per_task:-1}
    [[ ${NTHREADSmax} -gt ${max_threads_per_task} ]] && NTHREADSmax=${max_threads_per_task}
    [[ ${NTHREADS1} -gt ${max_threads_per_task} ]] && NTHREADS1=${max_threads_per_task}
    APRUN_default="${launcher} -n ${ntasks}"
else
    echo "ERROR config.resources must be sourced before sourcing HERCULES.env"
    exit 2
fi

case ${step} in
 "prep" | "prepbufr")

    export POE="NO"
    export BACK=${BACK:-"YES"}
    export sys_tp="HERCULES"
    export launcher_PREP="srun --hint=nomultithread"
 ;;
 "prep_emissions")

    export APRUN="${APRUN_default}"
 ;;
 "waveinit" | "waveprep" | "wavepostsbs" | "wavepostbndpnt" | "wavepostpnt" | "wavepostbndpntbll")

    export CFP_MP="YES"
    [[ "${step}" = "waveprep" ]] && export MP_PULSE=0
    export wavempexec=${launcher}
    export wave_mpmd=${mpmd_opt}

 ;;
 "atmanlvar")

    export NTHREADS_ATMANLVAR=${NTHREADSmax}
    export APRUN_ATMANLVAR="${APRUN_default} --cpus-per-task=${NTHREADS_ATMANLVAR}"
 ;;
 "atmanlfv3inc")

    export NTHREADS_ATMANLFV3INC=${NTHREADSmax}
    export APRUN_ATMANLFV3INC="${APRUN_default} --cpus-per-task=${NTHREADS_ATMANLFV3INC}"
 ;;
 "atmensanlobs")

    export NTHREADS_ATMENSANLOBS=${NTHREADSmax}
    export APRUN_ATMENSANLOBS="${APRUN_default} --cpus-per-task=${NTHREADS_ATMENSANLOBS}"
 ;;
 "atmensanlsol")

    export NTHREADS_ATMENSANLSOL=${NTHREADSmax}
    export APRUN_ATMENSANLSOL="${APRUN_default} --cpus-per-task=${NTHREADS_ATMENSANLSOL}"
 ;;
 "atmensanlletkf")

    export NTHREADS_ATMENSANLLETKF=${NTHREADSmax}
    export APRUN_ATMENSANLLETKF="${APRUN_default} --cpus-per-task=${NTHREADS_ATMENSANLLETKF}"
 ;;
 "atmensanlfv3inc")

    export NTHREADS_ATMENSANLFV3INC=${NTHREADSmax}
    export APRUN_ATMENSANLFV3INC="${APRUN_default} --cpus-per-task=${NTHREADS_ATMENSANLFV3INC}"
 ;;
 "aeroanlvar")

    export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"

    export NTHREADS_AEROANL=${NTHREADSmax}
    export APRUN_AEROANL="${APRUN_default} --cpus-per-task=${NTHREADS_AEROANL}"
 ;;
  "aeroanlgenb")

    export NTHREADS_AEROANLGENB=${NTHREADSmax}
    export APRUN_AEROANLGENB="${APRUN_default} --cpus-per-task=${NTHREADS_AEROANLGENB}"
 ;;
 "prepobsaero")

    export NTHREADS_PREPOBSAERO=${NTHREADS1}
    export APRUN_PREPOBSAERO="${APRUN_default} --cpus-per-task=${NTHREADS_PREPOBSAERO}"
;;
 "snowanl")

    export APRUN_CALCFIMS="${launcher} -n 1"

    export NTHREADS_SNOWANL=${NTHREADSmax}
    export APRUN_SNOWANL="${APRUN_default} --cpus-per-task=${NTHREADS_SNOWANL}"

    export APRUN_APPLY_INCR="${launcher} -n 6"
 ;;
 "esnowanl")

    export APRUN_CALCFIMS="${launcher} -n 1"

    export NTHREADS_ESNOWANL=${NTHREADSmax}
    export APRUN_ESNOWANL="${APRUN_default} --cpus-per-task=${NTHREADS_ESNOWANL}"

    export APRUN_APPLY_INCR="${launcher} -n 6"
 ;;

 "marinebmat")

    export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"
    export APRUN_MARINEBMAT="${APRUN_default}"
 ;;
 "marineanlvar")

    export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"
    export APRUN_MARINEANLVAR="${APRUN_default}"
 ;;
"ocnanalecen")

   export APRUN_OCNANALECEN="${APRUN_default}"
;;

 "marineanlchkpt")

    export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"

    export NTHREADS_OCNANAL=${NTHREADSmax}
    export APRUN_MARINEANLCHKPT="${APRUN_default} --cpus-per-task=${NTHREADS_OCNANAL}"
 ;;
 "marineanlletkf")

    export NTHREADS_MARINEANLLETKF=${NTHREADSmax}
    export APRUN_MARINEANLLETKF="${APRUN_default}"
 ;;
 "anal" | "analcalc")

    export MKL_NUM_THREADS=4
    export MKL_CBWR=AUTO

    export CFP_MP=${CFP_MP:-"YES"}
    export USE_CFP=${USE_CFP:-"YES"}
    export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"


    export NTHREADS_GSI=${threads_per_task_anal:-${max_threads_per_task}}
    export APRUN_GSI="${APRUN_default} --cpus-per-task=${NTHREADS_GSI}"

    export NTHREADS_CALCINC=${threads_per_task_calcinc:-1}
    [[ ${NTHREADS_CALCINC} -gt ${max_threads_per_task} ]] && export NTHREADS_CALCINC=${max_threads_per_task}
    export APRUN_CALCINC="${launcher} \$ncmd --cpus-per-task=${NTHREADS_CALCINC}"

    export NTHREADS_CYCLE=${threads_per_task_cycle:-12}
    [[ ${NTHREADS_CYCLE} -gt ${max_tasks_per_node} ]] && export NTHREADS_CYCLE=${max_tasks_per_node}
    ntasks_cycle=${ntiles:-6}
    export APRUN_CYCLE="${launcher} -n ${ntasks_cycle} --cpus-per-task=${NTHREADS_CYCLE}"

    export NTHREADS_GAUSFCANL=1
    ntasks_gausfcanl=${ntasks_gausfcanl:-1}
    export APRUN_GAUSFCANL="${launcher} -n ${ntasks_gausfcanl} --cpus-per-task=${NTHREADS_GAUSFCANL}"
 ;;
 "sfcanl")

    export NTHREADS_CYCLE=${threads_per_task:-14}
    [[ ${NTHREADS_CYCLE} -gt ${max_tasks_per_node} ]] && export NTHREADS_CYCLE=${max_tasks_per_node}
    export APRUN_CYCLE="${APRUN_default} --cpus-per-task=${NTHREADS_CYCLE}"
 ;;
 "eobs")

    export MKL_NUM_THREADS=4
    export MKL_CBWR=AUTO

    export CFP_MP=${CFP_MP:-"YES"}
    export USE_CFP=${USE_CFP:-"YES"}
    export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"


    export NTHREADS_GSI=${NTHREADSmax}
    [[ ${NTHREADS_GSI} -gt ${max_threads_per_task} ]] && export NTHREADS_GSI=${max_threads_per_task}
    export APRUN_GSI="${APRUN_default} --cpus-per-task=${NTHREADS_GSI}"
 ;;
 "eupd")

    export CFP_MP=${CFP_MP:-"YES"}
    export USE_CFP=${USE_CFP:-"YES"}
    export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"


    export NTHREADS_ENKF=${NTHREADSmax}
    export APRUN_ENKF="${launcher} -n ${ntasks_enkf:-${ntasks}} --cpus-per-task=${NTHREADS_ENKF}"
 ;;
 "fcst" | "efcs")

    export OMP_STACKSIZE=512M

    (( nnodes = (ntasks+tasks_per_node-1)/tasks_per_node ))
    (( ufs_ntasks = nnodes*tasks_per_node ))
    # With ESMF threading, the model wants to use the full node
    export APRUN_UFS="${launcher} -n ${ufs_ntasks}"
    unset nnodes ufs_ntasks
 ;;

 "upp")

    export NTHREADS_UPP=${NTHREADS1}
    export APRUN_UPP="${APRUN_default} --cpus-per-task=${NTHREADS_UPP}"
 ;;

 "atmos_products")

    export USE_CFP="YES"  # Use MPMD for downstream product generation
 ;;

"oceanice_products")

    export NTHREADS_OCNICEPOST=${NTHREADS1}
    export APRUN_OCNICEPOST="${launcher} -n 1 --cpus-per-task=${NTHREADS_OCNICEPOST}"
;;

 "ecen")

    export NTHREADS_ECEN=${NTHREADSmax}
    export APRUN_ECEN="${APRUN_default} --cpus-per-task=${NTHREADS_ECEN}"

    export NTHREADS_CHGRES=${threads_per_task_chgres:-12}
    [[ ${NTHREADS_CHGRES} -gt ${max_tasks_per_node} ]] && export NTHREADS_CHGRES=${max_tasks_per_node}
    export APRUN_CHGRES="time"

    export NTHREADS_CALCINC=${threads_per_task_calcinc:-1}
    [[ ${NTHREADS_CALCINC} -gt ${max_threads_per_task} ]] && export NTHREADS_CALCINC=${max_threads_per_task}
    export APRUN_CALCINC="${APRUN_default} --cpus-per-task=${NTHREADS_CALCINC}"

 ;;
 "esfc")

    export NTHREADS_ESFC=${NTHREADSmax}
    export APRUN_ESFC="${APRUN_default} --cpus-per-task=${NTHREADS_ESFC}"

    export NTHREADS_CYCLE=${threads_per_task_cycle:-14}
    [[ ${NTHREADS_CYCLE} -gt ${max_tasks_per_node} ]] && export NTHREADS_CYCLE=${max_tasks_per_node}
    export APRUN_CYCLE="${APRUN_default} --cpus-per-task=${NTHREADS_CYCLE}"

 ;;
 "epos")

    export NTHREADS_EPOS=${NTHREADSmax}
    export APRUN_EPOS="${APRUN_default} --cpus-per-task=${NTHREADS_EPOS}"

 ;;
 "postsnd")

    export CFP_MP="YES"

    export NTHREADS_POSTSND=${NTHREADS1}
    export APRUN_POSTSND="${APRUN_default} --cpus-per-task=${NTHREADS_POSTSND}"

    export NTHREADS_POSTSNDCFP=${threads_per_task_postsndcfp:-1}
    [[ ${NTHREADS_POSTSNDCFP} -gt ${max_threads_per_task} ]] && export NTHREADS_POSTSNDCFP=${max_threads_per_task}
    export APRUN_POSTSNDCFP="${launcher} -n ${ntasks_postsndcfp} ${mpmd_opt}"

 ;;
 "awips")

    export NTHREADS_AWIPS=${NTHREADS1}
    export APRUN_AWIPSCFP="${APRUN_default} ${mpmd_opt}"

 ;;
 "gempak")

    echo "WARNING: ${step} is not enabled on ${machine}!"

 ;;
 "fit2obs")

    export NTHREADS_FIT2OBS=${NTHREADS1}
    export MPIRUN="${APRUN_default} --cpus-per-task=${NTHREADS_FIT2OBS}"

 ;;
 *)
    # Some other job not yet defined here
    echo "WARNING: The job step ${step} does not specify Hercules-specific resources"
 ;;
esac