Refactor parallel.py (NGWPC-10583)

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:

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}")

NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/mpi_utils.py

@@ -1,24 +1,21 @@
 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)
 
-    rand_uint64 = MPI.COMM_WORLD.bcast(rand_uint64, root=0)
+    rand_uint64 = comm.bcast(rand_uint64, root=0)
 
     # uuid.UUID expects a built-in Python int. Convert the NumPy uint64
     uid_64bit_hex = uuid.UUID(int=int(rand_uint64)).hex