Add two MPI barriers in R2S workflow (#3646)

openmc/deplete/microxs.py

@@ -177,8 +177,9 @@ def get_microxs_and_flux(
         if not openmc.lib.is_initialized:
             run_kwargs.setdefault('cwd', temp_dir)
 
-        # Run transport simulation
+        # Run transport simulation and synchronize
         statepoint_path = model.run(**run_kwargs)
+        comm.barrier()
 
         if comm.rank == 0:
             # Move the statepoint file if it is being saved to a specific path

openmc/deplete/r2s.py

@@ -392,6 +392,7 @@ class R2SManager:
         )
         output_path = output_dir / 'depletion_results.h5'
         integrator.integrate(final_step=False, path=output_path)
+        comm.barrier()
 
         # Get depletion results
         self.results['depletion_results'] = Results(output_path)

openmc/lib/core.py

@@ -7,6 +7,7 @@ import os
 from pathlib import Path
 from random import getrandbits
 from tempfile import TemporaryDirectory
+import traceback as tb
 
 import numpy as np
 from numpy.ctypeslib import as_array
@@ -700,6 +701,15 @@ class TemporarySession:
         if self.already_initialized:
             return
 
+        # If an exception occurred, abort all ranks immediately
+        if exc_type is not None:
+            # Print exception info on the rank that failed
+            tb.print_exception(exc_type, exc_value, traceback)
+            sys.stdout.flush()
+
+            # Abort all MPI processes
+            self.comm.Abort(1)
+
         try:
             finalize()
         finally: