From 6b2d84ed5fd05fa6df26456b7e0bd6e1b680811f Mon Sep 17 00:00:00 2001
From: David Huber <69919478+DavidHuber-NOAA@users.noreply.github.com>
Date: Mon, 2 Dec 2024 08:12:32 -0500
Subject: [PATCH] Allow users to override the default account at setup time
 (#3127)

# Description
This adds a new option to setup_expt.py, and adds a missing argument to
generate_workflows.sh, to allow users to specify which HPC account they
wish to use. It also allows users to specify the environmental variable
`HPC_ACCOUNT` at setup time to specify the account. Documentation is
updated to note the new feature.

Resolves #3124

# Type of change
- [x] Bug fix (fixes something broken)
- [x] New feature (adds functionality)

# Change characteristics
- Is this a breaking change (a change in existing functionality)? NO
- Does this change require a documentation update? YES
- Does this change require an update to any of the following submodules?
NO

# How has this been tested?
Setup experiments on Hera.

# Checklist
- [x] My code follows the style guidelines of this project
- [x] I have performed a self-review of my own code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have documented my code, including function, input, and output
descriptions
- [x] My changes generate no new warnings
- [ ] New and existing tests pass with my changes
- [x] This change is covered by an existing CI test or a new one has
been added
- [x] I have made corresponding changes to the system documentation if
necessary
---
 docs/source/setup.rst          | 3 ++-
 workflow/generate_workflows.sh | 1 +
 workflow/setup_expt.py         | 5 +++++
 3 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/docs/source/setup.rst b/docs/source/setup.rst
index 0916033cbd..1f94610f65 100644
--- a/docs/source/setup.rst
+++ b/docs/source/setup.rst
@@ -35,7 +35,7 @@ The following command examples include variables for reference but users should
 
    cd workflow
    ./setup_expt.py gfs forecast-only --idate $IDATE --edate $EDATE [--app $APP] [--start $START] [--interval $INTERVAL_GFS] [--resdetatmos $RESDETATMOS] [--resdetocean $RESDETOCEAN]
-     [--pslot $PSLOT] [--configdir $CONFIGDIR] [--comroot $COMROOT] [--expdir $EXPDIR]
+     [--pslot $PSLOT] [--configdir $CONFIGDIR] [--comroot $COMROOT] [--expdir $EXPDIR] [--account $ACCOUNT]
 
 where:
 
@@ -61,6 +61,7 @@ where:
    * ``$INTERVAL_GFS`` is the forecast interval in hours [default: 6]
    * ``$COMROOT`` is the path to your experiment output directory. Your ``ROTDIR`` (rotating com directory) will be created using ``COMROOT`` and ``PSLOT``. [default: $HOME (but do not use default due to limited space in home directories normally, provide a path to a larger scratch space)]
    * ``$EXPDIR`` is the path to your experiment directory where your configs will be placed and where you will find your workflow monitoring files (i.e. rocoto database and xml file). DO NOT include PSLOT folder at end of path, it will be built for you. [default: $HOME]
+   * ``$ACCOUNT`` is the HPC (i.e. slurm or PBS) account to use for the experiment. [default: $HPC_ACCOUNT; if $HPC_ACCOUNT is not set, then the default in the host file (workflow/hosts/<machine>.yaml) will be used]
 
 Examples:
 
diff --git a/workflow/generate_workflows.sh b/workflow/generate_workflows.sh
index 6a4cb9910a..c98fa3028a 100755
--- a/workflow/generate_workflows.sh
+++ b/workflow/generate_workflows.sh
@@ -145,6 +145,7 @@ while [[ $# -gt 0 && "$1" != "--" ]]; do
         t) _tag="_${OPTARG}" ;;
         v) _verbose=true ;;
         V) _very_verbose=true && _verbose=true && _verbose_flag="-v" ;;
+        A) _set_account=true && _hpc_account="${OPTARG}" ;;
         d) _debug=true && _very_verbose=true && _verbose=true && _verbose_flag="-v" && PS4='${LINENO}: ' ;;
         h) _usage && exit 0 ;;
         :)
diff --git a/workflow/setup_expt.py b/workflow/setup_expt.py
index 27da4943d3..574dc0d91a 100755
--- a/workflow/setup_expt.py
+++ b/workflow/setup_expt.py
@@ -235,6 +235,7 @@ def _common_args(parser):
         parser.add_argument('--idate', help='starting date of experiment, initial conditions must exist!',
                             required=True, type=lambda dd: to_datetime(dd))
         parser.add_argument('--edate', help='end date experiment', required=False, type=lambda dd: to_datetime(dd))
+        parser.add_argument('--account', help='HPC account to use; default is host-dependent', required=False, default=os.getenv('HPC_ACCOUNT'))
         parser.add_argument('--interval', help='frequency of forecast (in hours); must be a multiple of 6', type=_validate_interval, required=False, default=6)
         parser.add_argument('--icsdir', help='full path to user initial condition directory', type=str, required=False, default='')
         parser.add_argument('--overwrite', help='overwrite previously created experiment (if it exists)',
@@ -407,6 +408,10 @@ def main(*argv):
 
     validate_user_request(host, user_inputs)
 
+    # Update the default host account if the user supplied one
+    if user_inputs.account is not None:
+        host.info.ACCOUNT = user_inputs.account
+
     # Determine ocean resolution if not provided
     if user_inputs.resdetocean <= 0:
         user_inputs.resdetocean = get_ocean_resolution(user_inputs.resdetatmos)