mpi: Add raise_mpi utility for collective exceptions

Author: FabioLuporini

devito/mpi/distributed.py

@@ -70,6 +70,7 @@ def __getattr__(self, name):
     'SubDistributor',
     'devito_mpi_finalize',
     'devito_mpi_init',
+    'mpi_raise',
 ]
 
 
@@ -101,6 +102,21 @@ def devito_mpi_finalize():
         MPI.Finalize()
 
 
+def mpi_raise(error, exception=ValueError, comm=None):
+    """
+    Raise `exception` with the first non-None error message in rank order.
+
+    All ranks in `comm` must call this routine, including those with no local
+    error (`error=None`). This prevents a rank-local exception from stranding
+    peers in subsequent MPI calls. With no communicator or `MPI.COMM_NULL`,
+    only the local error is checked.
+    """
+    if comm is not None and comm is not MPI.COMM_NULL:
+        error = next((i for i in comm.allgather(error) if i is not None), None)
+    if error is not None:
+        raise exception(error)
+
+
 class AbstractDistributor(ABC):
 
     """

tests/test_mpi.py

@@ -13,13 +13,14 @@
 )
 from devito.arch.compiler import OneapiCompiler
 from devito.data import LEFT, RIGHT
+from devito.exceptions import InvalidArgument
 from devito.ir import Cluster, Interval, IterationSpace
 from devito.ir.clusters.algorithms import check_halo_writes
 from devito.ir.iet import (
     Call, Conditional, FindNodes, FindSymbols, Iteration, retrieve_iteration_tree
 )
 from devito.ir.support.space import Backward, Forward
-from devito.mpi import MPI
+from devito.mpi import MPI, mpi_raise
 from devito.mpi.distributed import CustomTopology
 from devito.mpi.routines import ComputeCall, HaloUpdateCall, HaloUpdateList, MPICall
 from devito.tools import Bunch
@@ -28,6 +29,32 @@
 from examples.seismic.acoustic import acoustic_setup
 
 
+class TestMPIUtils:
+
+    @pytest.mark.parametrize('comm', [None, MPI.COMM_NULL], ids=['none', 'null'])
+    @pytest.mark.parametrize('error', [None, '', 'invalid argument'])
+    @pytest.mark.parametrize('exception', [ValueError, InvalidArgument])
+    def test_raise_serial(self, comm, error, exception):
+        if error is None:
+            mpi_raise(error, exception, comm=comm)
+        else:
+            with pytest.raises(exception, match=f'^{error}$'):
+                mpi_raise(error, exception, comm=comm)
+
+    @pytest.mark.parametrize('failing_ranks', [(), (0,), (1,), (0, 1)])
+    @pytest.mark.parallel(mode=[(2, 'basic')])
+    def test_raise_collective(self, failing_ranks, mode):
+        grid = Grid(shape=(16, 16))
+        rank = grid.distributor.myrank
+        error = f'rank {rank}' if rank in failing_ranks else None
+
+        if failing_ranks:
+            with pytest.raises(InvalidArgument, match=f'^rank {failing_ranks[0]}$'):
+                mpi_raise(error, InvalidArgument, comm=grid.comm)
+        else:
+            mpi_raise(error, InvalidArgument, comm=grid.comm)
+
+
 class TestDistributor:
 
     @pytest.mark.parallel(mode=[2, 4])