Support alternate ngen log path by persisting OS env var NGEN_RESULTS_DIR if already set

Author: mxkpp

python/nwm_fcst_mgr/forecast.py

@@ -19,6 +19,7 @@
 from nwm_fcst_mgr.log_level import log_level_set
 from nwm_fcst_mgr.git_util import print_git_info_all
 from nwm_fcst_mgr.exceptions import NgenCalledProcessError, NgenIntentionallyStoppedError
+from nwm_fcst_mgr.utils import set_os_env_key, OS_ENV_KEY_RESULTS_DIR
 
 # setup the logger
 log_level_set()
@@ -183,8 +184,9 @@ def preprocess(self) -> None:
         """Preprocess an ngen run, validate some inputs, and set the execution status."""
 
         # set environment variable for ngencerf backend
-        os.environ["NGEN_RESULTS_DIR"] = str(Path(self.real_path).parent)
-        logging.info(f"Set environment variable NGEN_RESULTS_DIR to: {os.environ['NGEN_RESULTS_DIR']}")
+        set_os_env_key(
+            OS_ENV_KEY_RESULTS_DIR, str(Path(self.real_path).parent), override=False
+        )
 
         # Read validation yaml file
         self.valid_config = load_yaml(self.valid_yaml)

python/nwm_fcst_mgr/utils.py

@@ -0,0 +1,43 @@
+"""Utilities"""
+
+import logging
+from os import environ
+
+
+OS_ENV_KEY_RESULTS_DIR = "NGEN_RESULTS_DIR"
+
+LOG = logging.getLogger(__name__)
+
+
+def set_os_env_key(key: str, val: str, override: bool = True) -> None:
+    """Set the value of the OS environment key.
+    Optionally, keep the existing value for that key without overriding, if it already exists.
+
+    Parameters:
+        key : str
+            OS environment key whose value will be modified.
+        val : str
+            New value to set to.
+        override : bool (default True)
+            If True, then do replace the existing value of that key if it already exists.
+            If False, then do not replace the value.
+    """
+    errors: list[Exception] = []
+    if not isinstance(key, str):
+        errors.append(TypeError(f"For key {key}, expected type {str}, got {type(key)}"))
+    if not isinstance(val, str):
+        errors.append(
+            TypeError(f"For value {val}, expected type {str}, got {type(val)}")
+        )
+    if errors:
+        raise RuntimeError(errors)
+
+    if key in environ:
+        msg_suffix = f"OS env key {repr(key)} already exists with value {repr(environ[key])}, override={override}"
+        if not override:
+            LOG.info("Will not override: " + msg_suffix)
+            return
+        LOG.info("Will override: " + msg_suffix)
+
+    LOG.info(f"Setting OS env key {repr(key)} to value {repr(val)}.")
+    environ[key] = val