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/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/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/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..fa7d7c44 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py @@ -1,18 +1,14 @@ -import atexit +from __future__ import annotations from functools import partial import os -import uuid -import signal -import sys - +import typing 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 from . import mpi_utils @@ -22,178 +18,100 @@ mpi4py.rc.initialize = False mpi4py.rc.finalize = False +if typing.TYPE_CHECKING: + 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 = 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) - self.__register_exit_handlers() - 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 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"): 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 = getattr(self, "comm", None) - if comm is None: + if self.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") + self.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 +140,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): @@ -244,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""" @@ -284,68 +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 __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 +210,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: T) -> T: """Broadcast a single parameter value to all processors. Generic function for sending a parameter value out to the processors. @@ -365,90 +218,30 @@ def broadcast_parameter(self, value_broadcast, config_options, param_type): :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) - - def scatter_array_logan(self, geoMeta, array_broadcast, ConfigOptions): - """Scatter an array based on the input dataset type. - - Generic function for calling scatter functons based on + 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 from rank 0 to all other ranks. + + Generic function for calling scatter functions based on the input dataset type. - :param geoMeta: - :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: - :return: + :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: 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. @@ -469,35 +262,30 @@ def scatter_array_scatterv_no_cache(self, geoMeta, src_array, ConfigOptions): try: self.comm.Bcast(data_type_buffer, root=0) - except: - ConfigOptions.errMsg = ( - "Unable to broadcast numpy datatype value from rank 0" - ) - err_handler.err_out(ConfigOptions) - return None + except Exception as 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 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( - self.rank - ) - err_handler.err_out(ConfigOptions) - 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] @@ -544,10 +332,12 @@ 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) - 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, @@ -558,14 +348,17 @@ 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: 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`. - def merge_slabs_gatherv(self, local_slab, options, 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 + 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 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: @@ -581,13 +374,10 @@ def merge_slabs_gatherv(self, local_slab, options, allgather: bool = False): try: self.comm.Allgather([shapes, MPI.INTEGER], [global_shapes, MPI.INTEGER]) - except: - 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) + except Exception: + 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: # check that all slabes are the same width and sum the number of rows @@ -596,19 +386,12 @@ def merge_slabs_gatherv(self, local_slab, options, allgather: bool = False): 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) - # 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) + err_handler.log_critical(self.config_options, self) + return None # generate counts counts = [ @@ -621,9 +404,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 +418,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 +445,9 @@ def merge_slabs_gatherv(self, local_slab, options, allgather: bool = False): recvbuf=[recvbuf, counts, offsets, data_type], root=0, ) - except: - options.errMsg = "Failed to Gatherv to rank 0 from rank " + str(self.rank) - err_handler.log_critical(options, self) + except Exception: + 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 - # options.errMsg = "Gatherv complete" - # err_handler.log_msg(options,self) - return recvbuf 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() 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