changes to enable DA cycling on gaea

NOAA-EMC · Oct 23, 2024 · d11184e · d11184e
1 parent 5cc20ec
commit d11184e
Show file tree

Hide file tree

Showing 3 changed files with 219 additions and 24 deletions.
diff --git a/env/GAEA.env b/env/GAEA.env
@@ -26,22 +26,127 @@ if [[ -n "${ntasks:-}" && -n "${max_tasks_per_node:-}" && -n "${tasks_per_node:-
     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}
-    # This may be useful when Gaea is fully ported, so ignore SC warning
-    # shellcheck disable=SC2034
     APRUN_default="${launcher} -n ${ntasks}"
 else
     echo "ERROR config.resources must be sourced before sourcing GAEA.env"
     exit 2
 fi
 
-if [[ "${step}" = "prep" ]]; then
+case ${step} in
+ "prep" | "prepbufr")
 
     export POE="NO"
-    export BACK="NO"
+    export BACK=${BACK:-"YES"}
     export sys_tp="GAEA"
     export launcher_PREP="srun"
+ ;;
+ "prepsnowobs")
 
-elif [[ "${step}" = "anal" ]] || [[ "${step}" = "analcalc" ]]; then
+    export APRUN_CALCFIMS="${APRUN_default}"
+ ;;
+ "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 NTHREADS_SNOWANL=${NTHREADSmax}
+    export APRUN_SNOWANL="${APRUN_default} --cpus-per-task=${NTHREADS_SNOWANL}"
+
+    export APRUN_APPLY_INCR="${launcher} -n 6"
+ ;;
+ "esnowrecen")
+
+    export NTHREADS_ESNOWRECEN=${NTHREADSmax}
+    export APRUN_ESNOWRECEN="${APRUN_default} --cpus-per-task=${NTHREADS_ESNOWRECEN}"
+
+    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 NTHREADS_MARINEANLVAR=${threads_per_task_anal:-${max_threads_per_task}}
+    export APRUN_MARINEANLVAR="${APRUN_default} --cpus-per-task=${NTHREADS_MARINEANLVAR}"
+ ;;
+"ocnanalecen")
+
+    export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"
+
+    max_threads_per_task=$((max_tasks_per_node / tasks_per_node_ocnanalecen))
+
+    export NTHREADS_OCNANALECEN=${threads_per_task_ocnanalecen:-${max_threads_per_task}}
+    [[ ${NTHREADS_OCNANALECEN} -gt ${max_threads_per_task} ]] && export NTHREADS_OCNANALECEN=${max_threads_per_task}
+    export APRUN_OCNANALECEN="${launcher} -n ${ntasks_ocnanalecen} --cpus-per-task=${NTHREADS_OCNANALECEN}"
+;;
+ "marineanlchkpt")
+
+    export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"
+
+    export NTHREADS_OCNANAL=${NTHREADSmax}
+    export APRUN_MARINEANLCHKPT="${APRUN_default} --cpus-per-task=${NTHREADS_OCNANAL}"
+ ;;
+ "anal" | "analcalc")
 
     export MKL_NUM_THREADS=4
     export MKL_CBWR=AUTO
@@ -50,7 +155,8 @@ elif [[ "${step}" = "anal" ]] || [[ "${step}" = "analcalc" ]]; then
     export USE_CFP=${USE_CFP:-"YES"}
     export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"
 
-    export NTHREADS_GSI=${NTHREADSmax}
+
+    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}
@@ -65,37 +171,126 @@ elif [[ "${step}" = "anal" ]] || [[ "${step}" = "analcalc" ]]; then
     export NTHREADS_GAUSFCANL=1
     ntasks_gausfcanl=${ntasks_gausfcanl:-1}
     export APRUN_GAUSFCANL="${launcher} -n ${ntasks_gausfcanl} --cpus-per-task=${NTHREADS_GAUSFCANL}"
-
-elif [[ "${step}" = "sfcanl" ]]; then
+ ;;
+ "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")
 
-elif [[ "${step}" = "fcst" ]]; then
+    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
+ ;;
 
-elif [[ "${step}" = "upp" ]]; then
+ "upp")
 
     export NTHREADS_UPP=${NTHREADS1}
     export APRUN_UPP="${APRUN_default} --cpus-per-task=${NTHREADS_UPP}"
+ ;;
 
-elif [[ "${step}" = "atmos_products" ]]; then
+ "atmos_products")
 
-    export USE_CFP="YES"  # Use MPMD for downstream product generation on Gaea
+    export USE_CFP="YES"  # Use MPMD for downstream product generation
+ ;;
 
-elif [[ "${step}" = "oceanice_products" ]]; then
+"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}"
 
-elif [[ "${step}" = "fit2obs" ]]; then
+ ;;
+ "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}"
 
-fi
+ ;;
+ *)
+    # Some other job not yet defined here
+    echo "WARNING: The job step ${step} does not specify Hercules-specific resources"
+ ;;
+esac
diff --git a/ush/load_ufsda_modules.sh b/ush/load_ufsda_modules.sh
@@ -34,13 +34,13 @@ source "${HOMEgfs}/ush/module-setup.sh"
 module use "${HOMEgfs}/sorc/gdas.cd/modulefiles"
 
 case "${MACHINE_ID}" in
-  ("hera" | "orion" | "hercules" | "wcoss2")
+  ("hera" | "orion" | "gaea" | "hercules" | "wcoss2")
     module load "${MODS}/${MACHINE_ID}"
     ncdump=$( command -v ncdump )
     NETCDF=$( echo "${ncdump}" | cut -d " " -f 3 )
     export NETCDF
     ;;
-  ("jet" | "gaea" | "s4" | "acorn")
+  ("jet" | "s4" | "acorn")
     echo WARNING: UFSDA NOT SUPPORTED ON THIS PLATFORM
     ;;  
   *)

diff --git a/workflow/hosts/gaea.yaml b/workflow/hosts/gaea.yaml
@@ -5,18 +5,18 @@ BASE_IC: '/gpfs/f5/ufs-ard/world-shared/global/glopara/data/ICSDIR'
 PACKAGEROOT: '/gpfs/f5/ufs-ard/world-shared/global/glopara/data/nwpara'
 COMROOT: '/gpfs/f5/ufs-ard/world-shared/global/glopara/data/com'
 COMINsyn: '${COMROOT}/gfs/prod/syndat'
-HOMEDIR: '/gpfs/f5/ufs-ard/scratch/${USER}'
-STMP: '/gpfs/f5/ufs-ard/scratch/${USER}'
-PTMP: '/gpfs/f5/ufs-ard/scratch/${USER}'
-NOSCRUB: $HOMEDIR
 ACCOUNT: ufs-ard
+HOMEDIR: '/gpfs/f5/${ACCOUNT}/scratch/${USER}'
+STMP: '/gpfs/f5/${ACCOUNT}/scratch/${USER}'
+PTMP: '/gpfs/f5/${ACCOUNT}/scratch/${USER}'
+NOSCRUB: $HOMEDIR
 SCHEDULER: slurm
 QUEUE: normal
-QUEUE_SERVICE: normal
+QUEUE_SERVICE: hpss
 PARTITION_BATCH: batch
-PARTITION_SERVICE: batch
+PARTITION_SERVICE: dtn_f5_f6
 RESERVATION: ''
-CLUSTERS: 'c5'
+CLUSTERS: 'c5,es'
 CHGRP_RSTPROD: 'NO'
 CHGRP_CMD: 'chgrp rstprod'
 HPSSARCH: 'NO'