Support MPI parallelism in R2SManager (#3632)

Author: paulromano

openmc/deplete/microxs.py

@@ -170,13 +170,14 @@ def get_microxs_and_flux(
         # Reinitialize with tallies
         openmc.lib.init(intracomm=comm)
 
-    # create temporary run
     with TemporaryDirectory() as temp_dir:
-        if run_kwargs is None:
-            run_kwargs = {}
-        else:
-            run_kwargs = dict(run_kwargs)
-        run_kwargs.setdefault('cwd', temp_dir)
+        # Indicate to run in temporary directory unless being executed through
+        # openmc.lib, in which case we don't need to specify the cwd
+        run_kwargs = dict(run_kwargs) if run_kwargs else {}
+        if not openmc.lib.is_initialized:
+            run_kwargs.setdefault('cwd', temp_dir)
+
+        # Run transport simulation
         statepoint_path = model.run(**run_kwargs)
 
         if comm.rank == 0:
@@ -189,15 +190,18 @@ def get_microxs_and_flux(
             if path_input is not None:
                 model.export_to_model_xml(path_input)
 
-            with StatePoint(statepoint_path) as sp:
-                if reaction_rate_mode == 'direct':
-                    rr_tally = sp.tallies[rr_tally.id]
-                    rr_tally._read_results()
-                flux_tally = sp.tallies[flux_tally.id]
-                flux_tally._read_results()
+        # Broadcast updated statepoint path to all ranks
+        statepoint_path = comm.bcast(statepoint_path)
+
+        # Read in tally results (on all ranks)
+        with StatePoint(statepoint_path) as sp:
+            if reaction_rate_mode == 'direct':
+                rr_tally = sp.tallies[rr_tally.id]
+                rr_tally._read_results()
+            flux_tally = sp.tallies[flux_tally.id]
+            flux_tally._read_results()
 
     # Get flux values and make energy groups last dimension
-    flux_tally = comm.bcast(flux_tally)
     flux = flux_tally.get_reshaped_data()  # (domains, groups, 1, 1)
     flux = np.moveaxis(flux, 1, -1)  # (domains, 1, 1, groups)
 
@@ -206,7 +210,6 @@ def get_microxs_and_flux(
 
     if reaction_rate_mode == 'direct':
         # Get reaction rates
-        rr_tally = comm.bcast(rr_tally)
         reaction_rates = rr_tally.get_reshaped_data()  # (domains, groups, nuclides, reactions)
 
         # Make energy groups last dimension

openmc/deplete/r2s.py

@@ -11,6 +11,9 @@
 from .microxs import get_microxs_and_flux, write_microxs_hdf5, read_microxs_hdf5
 from .results import Results
 from ..checkvalue import PathLike
+from ..mpi import comm
+from openmc.lib import TemporarySession
+from openmc.utility_funcs import change_directory
 
 
 def get_activation_materials(
@@ -199,8 +202,10 @@ def run(
         """
 
         if output_dir is None:
+            # Create timestamped output directory and broadcast to all ranks for
+            # consistency (different ranks may have slightly different times)
             stamp = datetime.now().strftime('%Y-%m-%dT%H-%M-%S')
-            output_dir = Path(f'r2s_{stamp}')
+            output_dir = Path(comm.bcast(f'r2s_{stamp}'))
 
         # Set run_kwargs for the neutron transport step
         if micro_kwargs is None:
@@ -257,18 +262,19 @@ def step1_neutron_transport(
 
         """
 
-        output_dir = Path(output_dir)
+        output_dir = Path(output_dir).resolve()
         output_dir.mkdir(parents=True, exist_ok=True)
 
         if self.method == 'mesh-based':
             # Compute material volume fractions on the mesh
             if mat_vol_kwargs is None:
                 mat_vol_kwargs = {}
-            self.results['mesh_material_volumes'] = mmv = \
-                self.domains.material_volumes(self.neutron_model, **mat_vol_kwargs)
+            self.results['mesh_material_volumes'] = mmv = comm.bcast(
+                self.domains.material_volumes(self.neutron_model, **mat_vol_kwargs))
 
             # Save results to file
-            mmv.save(output_dir / 'mesh_material_volumes.npz')
+            if comm.rank == 0:
+                mmv.save(output_dir / 'mesh_material_volumes.npz')
 
             # Create mesh-material filter based on what combos were found
             domains = openmc.MeshMaterialFilter.from_volumes(self.domains, mmv)
@@ -299,13 +305,16 @@ def step1_neutron_transport(
         micro_kwargs.setdefault('path_statepoint', output_dir / 'statepoint.h5')
         micro_kwargs.setdefault('path_input', output_dir / 'model.xml')
 
-        # Run neutron transport and get fluxes and micros
-        self.results['fluxes'], self.results['micros'] = get_microxs_and_flux(
-            self.neutron_model, domains, **micro_kwargs)
+        # Run neutron transport and get fluxes and micros. Run via openmc.lib to
+        # maintain a consistent parallelism strategy with the activation step.
+        with TemporarySession():
+            self.results['fluxes'], self.results['micros'] = get_microxs_and_flux(
+                self.neutron_model, domains, **micro_kwargs)
 
         # Save flux and micros to file
-        np.save(output_dir / 'fluxes.npy', self.results['fluxes'])
-        write_microxs_hdf5(self.results['micros'], output_dir / 'micros.h5')
+        if comm.rank == 0:
+            np.save(output_dir / 'fluxes.npy', self.results['fluxes'])
+            write_microxs_hdf5(self.results['micros'], output_dir / 'micros.h5')
 
     def step2_activation(
         self,
@@ -457,15 +466,17 @@ def step3_photon_transport(
         # photon model if it is different from the neutron model to account for
         # potential material changes
         if self.method == 'mesh-based' and different_photon_model:
-            self.results['mesh_material_volumes_photon'] = photon_mmv = \
-                self.domains.material_volumes(self.photon_model, **mat_vol_kwargs)
+            self.results['mesh_material_volumes_photon'] = photon_mmv = comm.bcast(
+                self.domains.material_volumes(self.photon_model, **mat_vol_kwargs))
 
             # Save photon MMV results to file
-            photon_mmv.save(output_dir / 'mesh_material_volumes.npz')
+            if comm.rank == 0:
+                photon_mmv.save(output_dir / 'mesh_material_volumes.npz')
 
-        tally_ids = [tally.id for tally in self.photon_model.tallies]
-        with open(output_dir / 'tally_ids.json', 'w') as f:
-            json.dump(tally_ids, f)
+        if comm.rank == 0:
+            tally_ids = [tally.id for tally in self.photon_model.tallies]
+            with open(output_dir / 'tally_ids.json', 'w') as f:
+                json.dump(tally_ids, f)
 
         self.results['photon_tallies'] = {}
 
@@ -514,8 +525,9 @@ def step3_photon_transport(
                 time_index = len(self.results['depletion_results']) + time_index
 
             # Run photon transport calculation
-            run_kwargs['cwd'] = Path(output_dir) / f'time_{time_index}'
-            statepoint_path = self.photon_model.run(**run_kwargs)
+            photon_dir = Path(output_dir) / f'time_{time_index}'
+            with TemporarySession(self.photon_model, cwd=photon_dir):
+                statepoint_path = self.photon_model.run(**run_kwargs)
 
             # Store tally results
             with openmc.StatePoint(statepoint_path) as sp:

openmc/lib/core.py

@@ -13,6 +13,7 @@
 
 from . import _dll
 from .error import _error_handler
+from ..mpi import comm
 from openmc.checkvalue import PathLike
 import openmc.lib
 import openmc
@@ -632,17 +633,23 @@ class TemporarySession:
     model : openmc.Model, optional
         OpenMC model to use for the session. If None, a minimal working model is
         created.
+    cwd : PathLike, optional
+        Working directory in which to run OpenMC. If None, a temporary directory
+        is created and deleted automatically.
     **init_kwargs
         Keyword arguments to pass to :func:`openmc.lib.init`.
 
     Attributes
     ----------
     model : openmc.Model
         The OpenMC model used for the session.
+    comm : mpi4py.MPI.Intracomm
+        The MPI intracommunicator used for the session.
 
     """
-    def __init__(self, model=None, **init_kwargs):
-        self.init_kwargs = init_kwargs
+    def __init__(self, model=None, cwd=None, **init_kwargs):
+        self.init_kwargs = dict(init_kwargs)
+        self.cwd = cwd
         if model is None:
             surf = openmc.Sphere(boundary_type="vacuum")
             cell = openmc.Cell(region=-surf)
@@ -652,6 +659,10 @@ def __init__(self, model=None, **init_kwargs):
                 particles=1, batches=1, output={'summary': False})
         self.model = model
 
+        # Determine MPI intercommunicator
+        self.init_kwargs.setdefault('intracomm', comm)
+        self.comm = self.init_kwargs['intracomm']
+
     def __enter__(self):
         """Initialize the OpenMC library in a temporary directory."""
         # If already initialized, the context manager is a no-op
@@ -662,14 +673,24 @@ def __enter__(self):
         # Store original working directory
         self.orig_dir = Path.cwd()
 
-        # Set up temporary directory
-        self.tmp_dir = TemporaryDirectory()
-        working_dir = Path(self.tmp_dir.name)
-        working_dir.mkdir(parents=True, exist_ok=True)
-        os.chdir(working_dir)
+        if self.cwd is None:
+            # Set up temporary directory on rank 0
+            if self.comm.rank == 0:
+                self._tmp_dir = TemporaryDirectory()
+                self.cwd = self._tmp_dir.name
+
+            # Broadcast the path so that all ranks use the same directory
+            self.cwd = self.comm.bcast(self.cwd)
 
-        # Export model and initialize OpenMC
-        self.model.export_to_model_xml()
+        # Create and change to specified directory
+        self.cwd = Path(self.cwd)
+        self.cwd.mkdir(parents=True, exist_ok=True)
+        os.chdir(self.cwd)
+
+        # Export model on first rank and initialize OpenMC
+        if self.comm.rank == 0:
+            self.model.export_to_model_xml()
+        self.comm.barrier()
         openmc.lib.init(**self.init_kwargs)
 
         return self
@@ -683,7 +704,11 @@ def __exit__(self, exc_type, exc_value, traceback):
             finalize()
         finally:
             os.chdir(self.orig_dir)
-            self.tmp_dir.cleanup()
+
+            # Make sure all ranks have finalized before deleting temporary dir
+            self.comm.barrier()
+            if hasattr(self, '_tmp_dir'):
+                self._tmp_dir.cleanup()
 
 
 class _DLLGlobal: