From 3e87dad6974204e6bea471d609e4bdb482beeed3 Mon Sep 17 00:00:00 2001 From: Ian Todd Date: Wed, 10 Jun 2026 11:59:25 -0400 Subject: [PATCH 1/5] Explicit cleanup call on BMI finalize() --- .../NextGen_Forcings_Engine/bmi_model.py | 3 + .../NextGen_Forcings_Engine/core/mpi_utils.py | 17 +- .../NextGen_Forcings_Engine/core/parallel.py | 253 ++++-------------- .../NextGen_Forcings_Engine/nc_utils.py | 37 --- 4 files changed, 59 insertions(+), 251 deletions(-) delete mode 100644 NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/nc_utils.py diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py index 3c27eacb..f5cc7564 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py @@ -584,6 +584,9 @@ def finalize(self): # process exits gc.collect() # make sure objects are deleted from memory LOG.info(Pld(St.COMPLETE, msg="Finishing BMI finalize()", modnm=MODNM)) + mpi_meta: MpiConfig = getattr(self, "_mpi_meta", None) + if mpi_meta is not None: + mpi_meta.cleanup() # ------------------------------------------------------------------- # ------------------------------------------------------------------- diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/mpi_utils.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/mpi_utils.py index 7f37b920..b4a89ad4 100644 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/mpi_utils.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/mpi_utils.py @@ -1,26 +1,23 @@ import uuid - +import numpy as np import mpi4py mpi4py.rc.threads = False -from mpi4py import MPI - -import numpy as np +from mpi4py import MPI # noqa: E402 -# def get_new_broadcasted_uid(comm: MPI.Comm) -> str: -def get_new_broadcasted_uid() -> str: +def get_new_broadcasted_uid(comm: MPI.Comm | None = None) -> str: """Broadcast a random uint64 then return the hash of that. Used for generating a random string shared among all ranks.""" - # if not isinstance(comm, MPI.Comm): - # raise TypeError(f"Expected comm to be type MPI.Comm, got: {type(comm)}") + if comm is None: + comm = MPI.COMM_WORLD rand_uint64 = None - if MPI.COMM_WORLD.rank == 0: + if comm.rank == 0: rng = np.random.default_rng() rand_uint64 = rng.integers(0, 2**64, dtype=np.uint64) else: rand_uint64 = None - rand_uint64 = MPI.COMM_WORLD.bcast(rand_uint64, root=0) + rand_uint64 = comm.bcast(rand_uint64, root=0) # Convert the NumPy uint64 to a built-in Python int. Python 3.14's uuid # implementation expects a native int internally, while this remains fully diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py index e0e48891..72a13a0f 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py @@ -1,16 +1,20 @@ -import atexit +from __future__ import annotations from functools import partial import os -import uuid import signal import sys +import typing +if typing.TYPE_CHECKING: + from typing import Any, Type, TypeVar + from .geoMod import GriddedGeoMeta + T = TypeVar("T") import mpi4py import numpy as np mpi4py.rc.threads = False -from mpi4py import MPI +from mpi4py import MPI # noqa: E402 from .config import ConfigOptions from . import err_handler @@ -46,7 +50,6 @@ def __init__(self, config_options: ConfigOptions): self.log_debug = partial(err_handler.log_msg, self.config_options, self, True) self.log_info = partial(err_handler.log_msg, self.config_options, self, False) self.log_warning = partial(err_handler.log_warning, self.config_options, self) - self.__register_exit_handlers() def initialize_comm(self, comm=None): """Initialize MPI communication. @@ -57,25 +60,29 @@ def initialize_comm(self, comm=None): try: self.comm = comm if comm is not None else MPI.COMM_WORLD self.comm.Set_errhandler(MPI.ERRORS_ARE_FATAL) - except AttributeError as ae: + except AttributeError: self.config_options.errMsg = ( "Unable to initialize the MPI Communicator object" ) - raise ae + raise try: self.size = self.comm.Get_size() - except MPI.Exception as mpi_exception: + except MPI.Exception: self.config_options.errMsg = "Unable to retrieve the MPI size." - raise mpi_exception + raise try: self.rank = self.comm.Get_rank() - except MPI.Exception as mpi_exception: + except MPI.Exception: self.config_options.errMsg = "Unable to retrieve the MPI processor rank." - raise mpi_exception + raise - self.__broadcast_new_64bit_uid() + try: + self.uid64 = mpi_utils.get_new_broadcasted_uid() + except Exception: + self.config_options.errMsg = "Unable to generate a global unique ID." + raise wait_for_debug = os.getenv("WAIT_FOR_DEBUGPY", "") if wait_for_debug.lower() in ("true", "1"): @@ -107,93 +114,18 @@ def abort_with_cleanup(self, errorcode: int) -> None: """Call cleanup methods, before calling MPI Abort. Do not make direct calls to MPI Abort without this method. Use this method for all MPI abort needs.""" - comm = getattr(self, "comm", None) + comm: MPI.Comm = getattr(self, "comm", None) if comm is None: raise RuntimeError("comm is not initialized") - # if self.rank == 0: - if True: - self._cleanup() - err_handler.log_msg( - self.config_options, self, debug=True, msg="About to MPI Abort" - ) - comm.Abort(errorcode) - comm.Barrier() # For testing case of only rank 0 aborting. - raise RuntimeError("At bottom of abort_with_cleanup, should not get here.") - - def __signals_handled(self) -> tuple[int]: - """Return a tuple of signals to be handled by cleanup routine.""" - ### signal.valid_signals() contains many that are unrelated to stoppage / interruption / error. - # sigs = [s for s in signal.valid_signals() if s not in (signal.SIGKILL, signal.SIGSTOP)] - sigs = ( - signal.SIGINT, - signal.SIGTERM, - signal.SIGHUP, - signal.SIGQUIT, - signal.SIGSEGV, - signal.SIGABRT, - signal.SIGFPE, - signal.SIGBUS, - signal.SIGILL, - ) - return sigs - - def __register_exit_handlers(self) -> None: - """Register exit handlers for unhandled exceptions, signals, and regular exits. - TODO: consider WCOSS gating. - TODO: note that when non-0 ranks call Abort directly, the rank 0 exit handler is still invoked, at least in some cases, - so there may be opportunities to streamline this further to have only rank 0 perform the cleanup. Would need to test - against potential deadlock conditions to be sure (would need to confirm that a non-0 rank initiating an abort would cause - rank 0 break out of a collective call if it happens to be waiting at one).""" - # Exceptions - sys.exepthook = self.__excepthook - # Regular exits - atexit.register(self._cleanup) - # Signals - for sig in self.__signals_handled(): - signal.signal(sig, self.__signal_handler) - - def __excepthook(self, ex_type, value, tb) -> None: - """Custom excepthook which follows these steps: - 1. Call Python's built-in excepthook. - 2. Log .errMsg as CRITICAL (unless it is None). - 3. Cleanup. - 4. MPI Abort. - - To apply, set `sys.excepthook` to this method.""" - sys.__excepthook__(ex_type, value, tb) - if self.config_options.errMsg is not None: - err_handler.log_critical( - self.config_options, - self, - msg=f"In excepthook, found errMsg = {repr(self.config_options.errMsg)}", - ) - self.abort_with_cleanup(1) - - def __signal_handler(self, signum, frame) -> None: - """Handle termination signals by cleaning up before exit.""" - ### Unregister the signal handler - for s in self.__signals_handled(): - signal.signal(s, signal.SIG_DFL) - ### Cleanup and re-send the original signal to itself - # self._cleanup() - # os.kill(os.getpid(), signum) - ### Cleanup and abort directly - self.abort_with_cleanup(signum) - - def _cleanup(self) -> None: - """High-level cleanup routine called by exit handlers.""" - # if self.rank != 0: - # return - err_handler.log_msg( - self.config_options, self, debug=True, msg="About to clean up" - ) + self.cleanup() + self.log_debug("About to MPI Abort") + comm.Abort(errorcode) + + def cleanup(self) -> None: + """High-level cleanup routine called during BMI finalization.""" + self.log_debug("About to clean up") self._cleanup_scratch_dir() self._cleanup_geogrid() - # TODO: Consider if this can be gated for non-wcoss only - try: - atexit.unregister(self._cleanup) - except Exception: - pass def _cleanup_scratch_dir(self) -> None: """Remove contents of scratch dir. @@ -222,7 +154,7 @@ def _cleanup_scratch_dir(self) -> None: # # Only delete files that don't start with either of these skip_starts = (".nfs", "NextGen_Forcings_Engine") - to_delete = [_ for _ in contents if not _.startswith(skip_starts)] + to_delete = [n for n in contents if not n.startswith(skip_starts)] for fn in to_delete: fp = os.path.join(self.config_options.scratch_dir, fn) if os.path.isfile(fp): @@ -342,10 +274,6 @@ def __test_exit(self, mode: str, rank: int) -> None: self.log_debug(msg) raise RuntimeError(msg) - def __broadcast_new_64bit_uid(self): - """Broadcast a random uint64 then save the hash of that to self.uid64, which effectively broadcasts the same unique string to all ranks.""" - self.uid64 = mpi_utils.get_new_broadcasted_uid() - def wait_for_debugpy_client(self): """Block until the debugpy clients have attached to cppdbg/gdb. @@ -357,7 +285,7 @@ def wait_for_debugpy_client(self): debugpy.listen(("localhost", 5678 + self.rank)) debugpy.wait_for_client() - def broadcast_parameter(self, value_broadcast, config_options, param_type): + def broadcast_parameter(self, value_broadcast: Any, config_options: ConfigOptions, param_type: Type[T]) -> T: """Broadcast a single parameter value to all processors. Generic function for sending a parameter value out to the processors. @@ -380,74 +308,14 @@ def broadcast_parameter(self, value_broadcast, config_options, param_type): return None return param.item(0) - def scatter_array_logan(self, geoMeta, array_broadcast, ConfigOptions): + def scatter_array(self, geo_meta: GriddedGeoMeta, src_array: np.ndarray, config_options: ConfigOptions): """Scatter an array based on the input dataset type. Generic function for calling scatter functons based on the input dataset type. - :param geoMeta: + :param geo_meta: :param array_broadcast: - :param ConfigOptions: - :return: - """ - # Determine which type of input array we have based on the - # type of numpy array. - data_type_flag = -1 - if self.rank == 0: - if array_broadcast.dtype == np.float32: - data_type_flag = 1 - if array_broadcast.dtype == np.float64: - data_type_flag = 2 - - # Broadcast the numpy datatype to the other processors. - if self.rank == 0: - tmpDict = {"varTmp": data_type_flag} - else: - tmpDict = None - try: - tmpDict = self.comm.bcast(tmpDict, root=0) - except Exception: - ConfigOptions.errMsg = ( - "Unable to broadcast numpy datatype value from rank 0" - ) - err_handler.log_critical(ConfigOptions, self) - return None - data_type_flag = tmpDict["varTmp"] - - # Broadcast the global array to the child processors, then - if self.rank == 0: - arrayGlobalTmp = array_broadcast - else: - if data_type_flag == 1: - arrayGlobalTmp = np.empty( - [geoMeta.ny_global, geoMeta.nx_global], np.float32 - ) - else: # data_type_flag == 2: - arrayGlobalTmp = np.empty( - [geoMeta.ny_global, geoMeta.nx_global], np.float64 - ) - try: - self.comm.Bcast(arrayGlobalTmp, root=0) - except Exception: - ConfigOptions.errMsg = ( - "Unable to broadcast a global numpy array from rank 0" - ) - err_handler.log_critical(ConfigOptions, self) - return None - arraySub = arrayGlobalTmp[ - geoMeta.y_lower_bound : geoMeta.y_upper_bound, - geoMeta.x_lower_bound : geoMeta.x_upper_bound, - ] - return arraySub - - def scatter_array_scatterv_no_cache(self, geoMeta, src_array, ConfigOptions): - """Scatter an array based on the input dataset type. - - Generic function for calling scatter functons based on - the input dataset type. - :param geoMeta: - :param array_broadcast: - :param ConfigOptions: + :param config_options: :return: """ # Determine which type of input array we have based on the @@ -469,34 +337,33 @@ def scatter_array_scatterv_no_cache(self, geoMeta, src_array, ConfigOptions): try: self.comm.Bcast(data_type_buffer, root=0) - except: - ConfigOptions.errMsg = ( + except Exception: + config_options.errMsg = ( "Unable to broadcast numpy datatype value from rank 0" ) - err_handler.err_out(ConfigOptions) + err_handler.err_out(config_options) return None data_type_flag = data_type_buffer[0] - data_type_buffer = None # gather buffer offsets and bounds to rank 0 bounds = np.array( [ - np.int32(geoMeta.x_lower_bound), - np.int32(geoMeta.y_lower_bound), - np.int32(geoMeta.x_upper_bound), - np.int32(geoMeta.y_upper_bound), + np.int32(geo_meta.x_lower_bound), + np.int32(geo_meta.y_lower_bound), + np.int32(geo_meta.x_upper_bound), + np.int32(geo_meta.y_upper_bound), ] ) global_bounds = np.zeros((self.size * 4), np.int32) try: self.comm.Allgather([bounds, MPI.INTEGER], [global_bounds, MPI.INTEGER]) - except: - ConfigOptions.errMsg = "Failed all gathering global bounds at rank" + str( + except Exception: + config_options.errMsg = "Failed all gathering global bounds at rank" + str( self.rank ) - err_handler.err_out(ConfigOptions) + err_handler.err_out(config_options) return None # create slices for x and y bounds arrays @@ -544,9 +411,9 @@ def scatter_array_scatterv_no_cache(self, geoMeta, src_array, ConfigOptions): # scatter the data try: self.comm.Scatterv([sendbuf, counts, offsets, data_type], recvbuf, root=0) - except: - ConfigOptions.errMsg = "Failed Scatterv from rank 0" - err_handler.error_out(ConfigOptions) + except Exception: + config_options.errMsg = "Failed Scatterv from rank 0" + err_handler.err_out(config_options) return None subarray = np.reshape( @@ -558,10 +425,7 @@ def scatter_array_scatterv_no_cache(self, geoMeta, src_array, ConfigOptions): ).copy() return subarray - # use scatterv based scatter_array - scatter_array = scatter_array_scatterv_no_cache - - def merge_slabs_gatherv(self, local_slab, options, allgather: bool = False): + def merge_slabs_gatherv(self, local_slab: np.ndarray, options: ConfigOptions, allgather: bool = False): """If allgather is True, then Allgatherv will be used instead of Gatherv, which causes all ranks to be distributed to all other ranks. This is necessary for the hydrofabric case, to handle how ngen's hydrologic @@ -581,13 +445,10 @@ def merge_slabs_gatherv(self, local_slab, options, allgather: bool = False): try: self.comm.Allgather([shapes, MPI.INTEGER], [global_shapes, MPI.INTEGER]) - except: + except Exception: options.errMsg = "Failed all gathering slab shapes at rank" + str(self.rank) err_handler.log_critical(options, self) - global_bounds = None - - # options.errMsg = "All gather for global shapes complete" - # err_handler.log_msg(options,self) + return None if len(local_slab.shape) == 2: # check that all slabes are the same width and sum the number of rows @@ -601,14 +462,7 @@ def merge_slabs_gatherv(self, local_slab, options, allgather: bool = False): + str(i) ) err_handler.log_critical(options, self) - # TODO why was there an abort here? - # Switched it to a new wrapped/cleanup abort, - # but would like to remove that call too if - # there is no reason to keep it. - self.abort_with_cleanup(1) - - # options.errMsg = "Checking of Rows and Columns complete" - # err_handler.log_msg(options,self) + return None # generate counts counts = [ @@ -621,9 +475,6 @@ def merge_slabs_gatherv(self, local_slab, options, allgather: bool = False): for i in range(0, len(counts) - 1): offsets.append(offsets[i] + counts[i]) - # options.errMsg = "Counts and Offsets generated" - # err_handler.log_msg(options,self) - # create the receive buffer if allgather or self.rank == 0: recvbuf = np.empty([total_rows, width], local_slab.dtype) @@ -638,9 +489,6 @@ def merge_slabs_gatherv(self, local_slab, options, allgather: bool = False): for i in range(0, len(counts) - 1): offsets.append(offsets[i] + counts[i]) - # options.errMsg = "Counts and Offsets generated" - # err_handler.log_msg(options,self) - # create the receive buffer if allgather or self.rank == 0: recvbuf = np.empty([sum(global_shapes)], local_slab.dtype) @@ -668,12 +516,9 @@ def merge_slabs_gatherv(self, local_slab, options, allgather: bool = False): recvbuf=[recvbuf, counts, offsets, data_type], root=0, ) - except: + except Exception: options.errMsg = "Failed to Gatherv to rank 0 from rank " + str(self.rank) err_handler.log_critical(options, self) return None - # options.errMsg = "Gatherv complete" - # err_handler.log_msg(options,self) - return recvbuf diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/nc_utils.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/nc_utils.py deleted file mode 100644 index 1148f24f..00000000 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/nc_utils.py +++ /dev/null @@ -1,37 +0,0 @@ -import types - -import netCDF4 - -from . import retry_utils -from .core.config import ConfigOptions -from .core.parallel import MpiConfig - - -@retry_utils.retry_w_mpi_context( - abort=True, num_retries=3, sleep_start=1, sleep_factor=3 -) -def nc_Dataset_retry( - mpi_config: MpiConfig, - config_options: ConfigOptions, - err_handler: types.ModuleType, - *args, - **kwargs, -): - """netCDF4 Dataset open with MPI retry logic.""" - return netCDF4.Dataset(*args, **kwargs) - - -@retry_utils.retry_w_mpi_context( - abort=True, num_retries=3, sleep_start=1, sleep_factor=3 -) -def nc_read_var_retry( - mpi_config: MpiConfig, - config_options: ConfigOptions, - err_handler: types.ModuleType, - nc_var: netCDF4.Variable, - slices=None, -): - """Read NetCDF variable data with MPI retry logic.""" - if slices is None: - return nc_var[:].data - return nc_var[slices].data From 5eae5380ce60c9581f0483df74ba3104f7e06c7c Mon Sep 17 00:00:00 2001 From: Ian Todd Date: Tue, 14 Jul 2026 11:17:41 -0400 Subject: [PATCH 2/5] Simplify broadcast_parameter to use bcast --- .../core/disaggregateMod.py | 8 ++--- .../NextGen_Forcings_Engine/core/parallel.py | 33 ++++++++----------- .../NextGen_Forcings_Engine/core/regrid.py | 28 +++++++--------- 3 files changed, 26 insertions(+), 43 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/disaggregateMod.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/disaggregateMod.py index d7a3d473..c4d8b9af 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/disaggregateMod.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/disaggregateMod.py @@ -150,9 +150,7 @@ def ak_ext_ana_disaggregate( date_iter += timedelta(hours=1) - found_target_hh = mpi_config.broadcast_parameter( - found_target_hh, config_options, param_type=bool - ) + found_target_hh = mpi_config.broadcast_parameter(found_target_hh) err_handler.check_program_status(config_options, mpi_config) if not found_target_hh: if mpi_config.rank == 0: @@ -167,9 +165,7 @@ def ak_ext_ana_disaggregate( supplemental_precip.regridded_precip2[:] = config_options.globalNdv return - read_hours = mpi_config.broadcast_parameter( - read_hours, config_options, param_type=int - ) + read_hours = mpi_config.broadcast_parameter(read_hours) err_handler.check_program_status(config_options, mpi_config) if read_hours != 6: if mpi_config.rank == 0: diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py index 72a13a0f..749586c1 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py @@ -4,11 +4,6 @@ import signal import sys import typing -if typing.TYPE_CHECKING: - from typing import Any, Type, TypeVar - from .geoMod import GriddedGeoMeta - T = TypeVar("T") - import mpi4py import numpy as np @@ -26,6 +21,10 @@ mpi4py.rc.initialize = False mpi4py.rc.finalize = False +if typing.TYPE_CHECKING: + from typing import Any, Type, TypeVar + from .geoMod import GriddedGeoMeta + T = TypeVar("T") class MpiConfig: """MPI config class. @@ -40,7 +39,7 @@ def __init__(self, config_options: ConfigOptions): NOTE: this class overrides the system excepthook so that cleanup steps and MPI abort can be triggered on unhandled exceptions. """ - self.comm = None + self.comm: MPI.Intracomm = None self.rank = None self.size = None self.uid64: str | None = ( @@ -285,7 +284,7 @@ def wait_for_debugpy_client(self): debugpy.listen(("localhost", 5678 + self.rank)) debugpy.wait_for_client() - def broadcast_parameter(self, value_broadcast: Any, config_options: ConfigOptions, param_type: Type[T]) -> T: + def broadcast_parameter(self, value_broadcast: T) -> T: """Broadcast a single parameter value to all processors. Generic function for sending a parameter value out to the processors. @@ -293,20 +292,14 @@ def broadcast_parameter(self, value_broadcast: Any, config_options: ConfigOption :param config_options: :return: """ - dtype = np.dtype(param_type) - - if self.rank == 0: - param = np.asarray(value_broadcast, dtype=dtype) - else: - param = np.empty(dtype=dtype, shape=()) - + if self.size == 1: + return value_broadcast try: - self.comm.Bcast(param, root=0) - except MPI.Exception: - config_options.errMsg = "Unable to broadcast single value from rank 0." - err_handler.log_critical(config_options, self) - return None - return param.item(0) + return self.comm.bcast(value_broadcast, root=0) + except Exception as e: + self.config_options.errMsg = f"Unable to broadcst single value {value_broadcast} from rank 0: {e.__class__.__name__} -- {e}" + err_handler.log_critical(self.config_options, self) + raise def scatter_array(self, geo_meta: GriddedGeoMeta, src_array: np.ndarray, config_options: ConfigOptions): """Scatter an array based on the input dataset type. diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py index 427791ce..0e0e34f6 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py @@ -213,11 +213,11 @@ def regrid_ak_ext_ana(input_forcings, config_options, wrf_hydro_geo_meta, mpi_co input_forcings.nx_global = ds.dimensions["x"].size input_forcings.ny_global = mpi_config.broadcast_parameter( - input_forcings.ny_global, config_options, param_type=int + input_forcings.ny_global ) err_handler.check_program_status(config_options, mpi_config) input_forcings.nx_global = mpi_config.broadcast_parameter( - input_forcings.nx_global, config_options, param_type=int + input_forcings.nx_global ) err_handler.check_program_status(config_options, mpi_config) @@ -10499,7 +10499,7 @@ def regrid_ndfd(input_forcings, config_options, wrf_hydro_geo_meta, mpi_config): ) # look to see if current time is in file: - skip_file = np.ubyte(0) + skip_file = False if mpi_config.rank == 0: times = [datetime.utcfromtimestamp(t) for t in id_tmp["time"][:]] if ndfd_var != "qpf": @@ -10523,16 +10523,14 @@ def regrid_ndfd(input_forcings, config_options, wrf_hydro_geo_meta, mpi_config): # TODO: qpf special handling if forecast_time > times[-1] - timedelta(hours=6): pt.log_debug("Forecast time beyond NDFD precip range, skipping") - skip_file = 1 + skip_file = True else: time_index = int(hour // 6) pt.log_debug( f"Forecast hour {forecast_time} will use precip from {times[time_index] - timedelta(hours=6)} to {times[time_index]}" ) - skip_file = mpi_config.broadcast_parameter( - skip_file, config_options, param_type=np.ubyte - ) + skip_file = mpi_config.broadcast_parameter(skip_file) err_handler.check_program_status(config_options, mpi_config) if skip_file: @@ -11399,9 +11397,7 @@ def check_regrid_status( # mpi_config.comm.barrier() # Broadcast the flag to the other processors. - calc_regrid_flag = mpi_config.broadcast_parameter( - calc_regrid_flag, config_options, param_type=bool - ) + calc_regrid_flag = mpi_config.broadcast_parameter(calc_regrid_flag) err_handler.check_program_status(config_options, mpi_config) return calc_regrid_flag @@ -11611,9 +11607,7 @@ def check_supp_pcp_regrid_status( # mpi_config.comm.barrier() # Broadcast the flag to the other processors. - calc_regrid_flag = mpi_config.broadcast_parameter( - calc_regrid_flag, config_options, param_type=bool - ) + calc_regrid_flag = mpi_config.broadcast_parameter(calc_regrid_flag) mpi_config.comm.barrier() return calc_regrid_flag @@ -11862,11 +11856,11 @@ def calculate_weights( # Broadcast the forcing nx/ny values input_forcings.ny_global = mpi_config.broadcast_parameter( - input_forcings.ny_global, config_options, param_type=int + input_forcings.ny_global ) err_handler.check_program_status(config_options, mpi_config) input_forcings.nx_global = mpi_config.broadcast_parameter( - input_forcings.nx_global, config_options, param_type=int + input_forcings.nx_global ) err_handler.check_program_status(config_options, mpi_config) @@ -12245,10 +12239,10 @@ def calculate_supp_pcp_weights( # Broadcast the forcing nx/ny values supplemental_precip.ny_global = mpi_config.broadcast_parameter( - supplemental_precip.ny_global, config_options, param_type=int + supplemental_precip.ny_global ) supplemental_precip.nx_global = mpi_config.broadcast_parameter( - supplemental_precip.nx_global, config_options, param_type=int + supplemental_precip.nx_global ) # mpi_config.comm.barrier() From f5b3cf457c7e1af7ec5e92ef8a33d8934bee3616 Mon Sep 17 00:00:00 2001 From: Ian Todd Date: Tue, 14 Jul 2026 12:30:40 -0400 Subject: [PATCH 3/5] Remove use of err_handler.err_out --- .../NextGen_Forcings_Engine/core/parallel.py | 28 +++++++++---------- 1 file changed, 13 insertions(+), 15 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py index 749586c1..f30492be 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py @@ -304,7 +304,7 @@ def broadcast_parameter(self, value_broadcast: T) -> T: def scatter_array(self, geo_meta: GriddedGeoMeta, src_array: np.ndarray, config_options: ConfigOptions): """Scatter an array based on the input dataset type. - Generic function for calling scatter functons based on + Generic function for calling scatter functions based on the input dataset type. :param geo_meta: :param array_broadcast: @@ -330,12 +330,12 @@ def scatter_array(self, geo_meta: GriddedGeoMeta, src_array: np.ndarray, config_ try: self.comm.Bcast(data_type_buffer, root=0) - except Exception: + except Exception as e: config_options.errMsg = ( - "Unable to broadcast numpy datatype value from rank 0" + f"Unable to broadcast numpy datatype value from rank 0: {e.__class__.__name__} -- {e}" ) - err_handler.err_out(config_options) - return None + err_handler.log_critical(config_options, self) + raise data_type_flag = data_type_buffer[0] @@ -352,12 +352,10 @@ def scatter_array(self, geo_meta: GriddedGeoMeta, src_array: np.ndarray, config_ try: self.comm.Allgather([bounds, MPI.INTEGER], [global_bounds, MPI.INTEGER]) - except Exception: - config_options.errMsg = "Failed all gathering global bounds at rank" + str( - self.rank - ) - err_handler.err_out(config_options) - return None + except Exception as e: + config_options.errMsg = f"Failed all gathering global bounds at rank {self.rank}: {e.__class__.__name__} -- {e}" + err_handler.log_critical(config_options, self) + raise # create slices for x and y bounds arrays x_lower = global_bounds[0 : (self.size * 4) + 0 : 4] @@ -404,10 +402,10 @@ def scatter_array(self, geo_meta: GriddedGeoMeta, src_array: np.ndarray, config_ # scatter the data try: self.comm.Scatterv([sendbuf, counts, offsets, data_type], recvbuf, root=0) - except Exception: - config_options.errMsg = "Failed Scatterv from rank 0" - err_handler.err_out(config_options) - return None + except Exception as e: + config_options.errMsg = f"Failed Scatterv from rank 0: {e.__class__.__name__} -- {e}" + err_handler.log_critical(config_options, self) + raise subarray = np.reshape( recvbuf, From f971ab5865fe9b910b91ffa285b3b7449f284669 Mon Sep 17 00:00:00 2001 From: Ian Todd Date: Tue, 14 Jul 2026 14:51:05 -0400 Subject: [PATCH 4/5] Documentation updates --- .../NextGen_Forcings_Engine/core/parallel.py | 167 ++++++------------ 1 file changed, 53 insertions(+), 114 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py index f30492be..ee0032d4 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py @@ -1,8 +1,6 @@ from __future__ import annotations from functools import partial import os -import signal -import sys import typing import mpi4py import numpy as np @@ -11,7 +9,6 @@ from mpi4py import MPI # noqa: E402 -from .config import ConfigOptions from . import err_handler from . import mpi_utils @@ -22,39 +19,52 @@ mpi4py.rc.finalize = False if typing.TYPE_CHECKING: - from typing import Any, Type, TypeVar + from typing import TypeVar from .geoMod import GriddedGeoMeta + from .config import ConfigOptions + T = TypeVar("T") + class MpiConfig: """MPI config class. - Abstract class for defining the MPI parameters, + Class for defining the MPI parameters, along with initialization of the MPI communication handle from mpi4py. """ + comm: MPI.Intercomm + rank: int + """MPI rank of the process of this instance.""" + size: int + """The number of MPI processes on this run.""" + uid64: str + """Random 16 chars based on random uint64 shared between all processes in a run.""" + config_options: ConfigOptions + """Forcing Engine configurations options. Messaging will use the instance of ConfigOptions passed to it in the constructor.""" + def __init__(self, config_options: ConfigOptions): """Initialize the MPI abstract class that will contain basic information and communication handles. - NOTE: this class overrides the system excepthook so that - cleanup steps and MPI abort can be triggered on unhandled exceptions. + + NOTE: Temporary files that are created during a normal forcing run should be cleaned up using the `cleanup` method. """ - self.comm: MPI.Intracomm = None + self.comm = None self.rank = None self.size = None - self.uid64: str | None = ( - None # broadcasted random 16 chars based on random uint64 - ) + self.uid64 = None self.config_options = config_options self.log_debug = partial(err_handler.log_msg, self.config_options, self, True) self.log_info = partial(err_handler.log_msg, self.config_options, self, False) self.log_warning = partial(err_handler.log_warning, self.config_options, self) - def initialize_comm(self, comm=None): - """Initialize MPI communication. + def initialize_comm(self, comm: MPI.Intercomm | None = None) -> None: + """Initialize MPI communication, including getting MPI rank and size. + Also generates the UID for the run. - Initial function to initialize MPI. - :return: + Usage note: if an exception is thrown, an error messsage is added to the + `config_options` object with the expectation the caller will handle + logging the error message generated. The error will be reraised. """ try: self.comm = comm if comm is not None else MPI.COMM_WORLD @@ -87,38 +97,15 @@ def initialize_comm(self, comm=None): if wait_for_debug.lower() in ("true", "1"): self.wait_for_debugpy_client() - # self._test_exit() - - # ------------------------------------------------------ - # Exit handling, exception handling, cleanup, and abort. - # ------------------------------------------------------ - - def _test_exit(self) -> None: - """Various methods for testing potential exit conditions""" - self.__test_exit("exception", 0) - # self.__test_exit("exception", 1) - # self.__test_exit("signal", 0) - ### Signal on rank 1 causes a deadlock iff abort_with_cleanup only allows rank 0 to abort, so all ranks need to be able to abort. - # self.__test_exit("signal", 1) - # self.__test_exit("sysexit1", 0) - # self.__test_exit("sysexit1", 1) - # self.__test_exit("check_program_status", 0) - # self.__test_exit("check_program_status", 1) - # self.__test_exit("err_out_screen", 0) - # self.__test_exit("err_out_screen", 1) - # self.__test_exit("err_out_screen_para", 0) - # self.__test_exit("err_out_screen_para", 1) - def abort_with_cleanup(self, errorcode: int) -> None: """Call cleanup methods, before calling MPI Abort. Do not make direct calls to MPI Abort without this method. Use this method for all MPI abort needs.""" - comm: MPI.Comm = getattr(self, "comm", None) - if comm is None: + if self.comm is None: raise RuntimeError("comm is not initialized") self.cleanup() self.log_debug("About to MPI Abort") - comm.Abort(errorcode) + self.comm.Abort(errorcode) def cleanup(self) -> None: """High-level cleanup routine called during BMI finalization.""" @@ -175,21 +162,18 @@ def _cleanup_geogrid(self) -> None: self.try_delete_file_no_reraise(geogrid) else: self.log_debug("Cleanup: config_options.geogrid is not set") - return def try_list_dir_no_reraise(self, dir_path: str) -> list[str]: """Try to list the directory and return a list of its contents. Do not reraise an exception if it fails due to FileNotFoundError or NotADirectoryError""" self.log_debug(f"Trying to list directory: {dir_path}") try: - contents = os.listdir(dir_path) + return os.listdir(dir_path) except (FileNotFoundError, NotADirectoryError) as e: self.log_debug( f"Could not list (it may have already been deleted): {dir_path}: {e}" ) return [] - else: - return contents def try_delete_file_no_reraise(self, file_path: str) -> None: """Try to delete the file, do not reraise an exception if it fails due to OSError""" @@ -215,64 +199,6 @@ def try_remove_empty_dir_no_reraise(self, dir_path: str) -> None: else: self.log_info(f"Removed directory: {dir_path}") - def __test_exit(self, mode: str, rank: int) -> None: - """Intentionally exit in a particular way, for testing exit/cleanup behavior. - `mode` : str. Mode of exit. See match/case block below for accepted values. - `rank` : int. Rank to perform the mode of exit. Can be 0 or 1. They have different""" - self.log_debug(f"__test_exit(): provided: mode={repr(mode)}, rank={repr(rank)}") - if rank not in (0, 1): - raise ValueError(f"__test_exit(): unsupported value for rank: {repr(rank)}") - - if self.rank == rank: - match mode: - case "exception": - msg = "__test_exit(): raising intentional RuntimeError" - self.log_debug(msg) - self.config_options.errMsg = "TEST" - raise RuntimeError(msg) - - case "signal": - # msg = f"__test_exit(): sending signal.SIGHUP ({signal.SIGHUP})" - msg = f"__test_exit(): sending signal.SIGTERM ({signal.SIGTERM})" - self.log_debug(msg) - # os.kill(os.getpid(), signal.SIGHUP) - os.kill(os.getpid(), signal.SIGTERM) - - case "sysexit1": - msg = "__test_exit(): calling sys.exit(1)" - self.log_debug(msg) - sys.exit(1) - - case "check_program_status": - msg = "__test_exit(): setting critical msg before calling check_program_status()" - self.log_debug(msg) - err_handler.log_critical( - self.config_options, self, msg="TESTING EXIT HANDLING" - ) - - case "err_out_screen": - msg = "__test_exit(): calling err_out_screen()" - self.log_debug(msg) - err_handler.err_out_screen(msg) - - case "err_out_screen_para": - msg = "__test_exit(): calling err_out_screen_para()" - self.log_debug(msg) - err_handler.err_out_screen_para(msg, self) - - case _: - raise ValueError(f"Unsupported mode={repr(mode)} for __test_exit()") - - self.log_debug("__test_exit(): reaching check_program_status()") - err_handler.check_program_status(self.config_options, self) - - self.log_debug("__test_exit(): reaching MPI Barrier") - self.comm.Barrier() - - msg = "__test_exit(): got past MPI Barrier (should not get here)" - self.log_debug(msg) - raise RuntimeError(msg) - def wait_for_debugpy_client(self): """Block until the debugpy clients have attached to cppdbg/gdb. @@ -301,15 +227,21 @@ def broadcast_parameter(self, value_broadcast: T) -> T: err_handler.log_critical(self.config_options, self) raise - def scatter_array(self, geo_meta: GriddedGeoMeta, src_array: np.ndarray, config_options: ConfigOptions): - """Scatter an array based on the input dataset type. + def scatter_array( + self, + geo_meta: GriddedGeoMeta, + src_array: np.ndarray, + config_options: ConfigOptions, + ): + """Scatter an array based on the input dataset type from rank 0 to all other ranks. Generic function for calling scatter functions based on the input dataset type. - :param geo_meta: - :param array_broadcast: + + :param geo_meta: GriddedGeoMeta instance used to determine the extent of the data received. + :param src_array: Data to be shared with other MPI ranks. :param config_options: - :return: + :return: The results of the scattered data filtered to the extent of `geo_meta` """ # Determine which type of input array we have based on the # type of numpy array. @@ -331,9 +263,7 @@ def scatter_array(self, geo_meta: GriddedGeoMeta, src_array: np.ndarray, config_ try: self.comm.Bcast(data_type_buffer, root=0) except Exception as e: - config_options.errMsg = ( - f"Unable to broadcast numpy datatype value from rank 0: {e.__class__.__name__} -- {e}" - ) + config_options.errMsg = f"Unable to broadcast numpy datatype value from rank 0: {e.__class__.__name__} -- {e}" err_handler.log_critical(config_options, self) raise @@ -403,7 +333,9 @@ def scatter_array(self, geo_meta: GriddedGeoMeta, src_array: np.ndarray, config_ try: self.comm.Scatterv([sendbuf, counts, offsets, data_type], recvbuf, root=0) except Exception as e: - config_options.errMsg = f"Failed Scatterv from rank 0: {e.__class__.__name__} -- {e}" + config_options.errMsg = ( + f"Failed Scatterv from rank 0: {e.__class__.__name__} -- {e}" + ) err_handler.log_critical(config_options, self) raise @@ -416,11 +348,18 @@ def scatter_array(self, geo_meta: GriddedGeoMeta, src_array: np.ndarray, config_ ).copy() return subarray - def merge_slabs_gatherv(self, local_slab: np.ndarray, options: ConfigOptions, allgather: bool = False): - """If allgather is True, then Allgatherv will be used instead of Gatherv, which causes all ranks to be distributed to all other ranks. + def merge_slabs_gatherv( + self, local_slab: np.ndarray, options: ConfigOptions, allgather: bool = False + ) -> np.ndarray: + """Gather arrays from all processes. The returned array will have the gathered data if `self.rank == 0` or `allgather` is `True`. - This is necessary for the hydrofabric case, to handle how ngen's hydrologic + The use of `allgather` is necessary for the hydrofabric case, to handle how ngen's hydrologic catchment partitionining differs from ESMF's arbitrary partitioning. + + :param local_slab: Data that will be gathered from all processes. + :param options: + :param allgather: Boolean on whether the gathered array should be broadcasted to all processes instead of just rank 0. + :return: Numpy array of the data gathered from all processes. """ # Filter based on dimensionality of array if len(local_slab.shape) == 2: From 37c4ca9ea85557fa6fdfd517d1bea7ad33479e39 Mon Sep 17 00:00:00 2001 From: Ian Todd Date: Tue, 14 Jul 2026 15:21:16 -0400 Subject: [PATCH 5/5] Always use default ConfigOptions for messaging --- .../NextGen_Forcings_Engine/core/ioMod.py | 4 ---- .../NextGen_Forcings_Engine/core/parallel.py | 15 +++++++-------- 2 files changed, 7 insertions(+), 12 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/ioMod.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/ioMod.py index 75bce435..96bf60b6 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/ioMod.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/ioMod.py @@ -990,12 +990,10 @@ def gather_global_outputs(self, ConfigOptions, geoMetaWrfHydro, MpiConfig): self.output_local[ output_variable_attribute_dict[varTmp][0], :, : ], - ConfigOptions, ) elif ConfigOptions.grid_type == "hydrofabric": dataOutTmp = MpiConfig.merge_slabs_gatherv( self.output_local[output_variable_attribute_dict[varTmp][0], :], - ConfigOptions, allgather=True, ) # NOTE this assumes that the var order here matches var order elsewhere. @@ -1006,14 +1004,12 @@ def gather_global_outputs(self, ConfigOptions, geoMetaWrfHydro, MpiConfig): self.output_local_elem[ output_variable_attribute_dict[varTmp][0], : ], - ConfigOptions, ) else: dataOutTmp = MpiConfig.merge_slabs_gatherv( self.output_local[ output_variable_attribute_dict[varTmp][0], : ], - ConfigOptions, ) else: raise ValueError(f"Invalid grid_type: {ConfigOptions.grid_type}") diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py index ee0032d4..fa7d7c44 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py @@ -349,7 +349,7 @@ def scatter_array( return subarray def merge_slabs_gatherv( - self, local_slab: np.ndarray, options: ConfigOptions, allgather: bool = False + self, local_slab: np.ndarray, allgather: bool = False ) -> np.ndarray: """Gather arrays from all processes. The returned array will have the gathered data if `self.rank == 0` or `allgather` is `True`. @@ -357,7 +357,6 @@ def merge_slabs_gatherv( catchment partitionining differs from ESMF's arbitrary partitioning. :param local_slab: Data that will be gathered from all processes. - :param options: :param allgather: Boolean on whether the gathered array should be broadcasted to all processes instead of just rank 0. :return: Numpy array of the data gathered from all processes. """ @@ -376,8 +375,8 @@ def merge_slabs_gatherv( try: self.comm.Allgather([shapes, MPI.INTEGER], [global_shapes, MPI.INTEGER]) except Exception: - options.errMsg = "Failed all gathering slab shapes at rank" + str(self.rank) - err_handler.log_critical(options, self) + self.config_options.errMsg = "Failed all gathering slab shapes at rank" + str(self.rank) + err_handler.log_critical(self.config_options, self) return None if len(local_slab.shape) == 2: @@ -387,11 +386,11 @@ def merge_slabs_gatherv( for i in range(0, self.size): total_rows += global_shapes[2 * i] if global_shapes[(2 * i) + 1] != width: - options.errMsg = ( + self.config_options.errMsg = ( "Error: slabs with differing widths detected on slab for rank" + str(i) ) - err_handler.log_critical(options, self) + err_handler.log_critical(self.config_options, self) return None # generate counts @@ -447,8 +446,8 @@ def merge_slabs_gatherv( root=0, ) except Exception: - options.errMsg = "Failed to Gatherv to rank 0 from rank " + str(self.rank) - err_handler.log_critical(options, self) + self.config_options.errMsg = "Failed to Gatherv to rank 0 from rank " + str(self.rank) + err_handler.log_critical(self.config_options, self) return None return recvbuf