Add deadlock safe function for calling fallible pieces of code on rank 0 (#4972)

Author: connorjward

pyop2/compilation.py

@@ -565,78 +565,63 @@ def make_so(compiler, code, extension, comm):
         exe = compiler.cc
         compiler_flags = compiler.cflags
 
-    # Compile on compilation communicator (ccomm) rank 0
-    if ccomm.rank == 0:
-        # Track exceptions as values so that they may be raised collectively
+    def compile_single_rank():
+        # Adding random 2-digit hexnum avoids using excessive filesystem inodes
+        tempdir = MEM_TMP_DIR.joinpath(f"{randint(0, 255):02x}")
+        tempdir.mkdir(parents=True, exist_ok=True)
+        # This path + filename should be unique
+        descriptor, filename = mkstemp(suffix=f".{extension}", dir=tempdir, text=True)
+        filename = Path(filename)
+
+        cname = filename
+        oname = filename.with_suffix(".o")
+        soname = filename.with_suffix(".so")
+        logfile = filename.with_suffix(".log")
+        errfile = filename.with_suffix(".err")
+
         try:
-            # Adding random 2-digit hexnum avoids using excessive filesystem inodes
-            tempdir = MEM_TMP_DIR.joinpath(f"{randint(0, 255):02x}")
-            tempdir.mkdir(parents=True, exist_ok=True)
-            # This path + filename should be unique
-            descriptor, filename = mkstemp(suffix=f".{extension}", dir=tempdir, text=True)
-            filename = Path(filename)
-
-            cname = filename
-            oname = filename.with_suffix(".o")
-            soname = filename.with_suffix(".so")
-            logfile = filename.with_suffix(".log")
-            errfile = filename.with_suffix(".err")
-        except BaseException as e:
-            result = e
-        else:
-            try:
-                with progress(INFO, 'Compiling wrapper'):
-                    # Write source code to disk
-                    with open(cname, "w") as fh:
-                        fh.write(code)
-                    os.close(descriptor)
-
-                    if not compiler.ld:
-                        # Compile and link
-                        cc = (exe,) + compiler_flags + ('-o', str(soname), str(cname)) + compiler.ldflags
-                        _run(cc, logfile, errfile)
-                    else:
-                        # Compile
-                        cc = (exe,) + compiler_flags + ('-c', '-o', str(oname), str(cname))
-                        _run(cc, logfile, errfile)
-                        # Extract linker specific "cflags" from ldflags and link
-                        ld = tuple(shlex.split(compiler.ld)) + ('-o', str(soname), str(oname)) + tuple(expandWl(compiler.ldflags))
-                        _run(ld, logfile, errfile, step="Linker", filemode="a")
-
-                result = soname
-            except subprocess.CalledProcessError as e:
-                msg = dedent(f"""
-                    Command "{e.cmd}" return error status {e.returncode}.
-                    Unable to compile code
-                """)
-                if os.environ.get("FIREDRAKE_CI", False):
-                    msg += dedent(f"""
-                        Code is:
-                        {code}
-                    """)
-                    with open(errfile) as err:
-                        msg += dedent(f"""
-                            Compiler output is:
-                            {''.join(err.readlines())}
-                        """)
+            with progress(INFO, 'Compiling wrapper'):
+                # Write source code to disk
+                with open(cname, "w") as fh:
+                    fh.write(code)
+                os.close(descriptor)
+
+                if not compiler.ld:
+                    # Compile and link
+                    cc = (exe,) + compiler_flags + ('-o', str(soname), str(cname)) + compiler.ldflags
+                    _run(cc, logfile, errfile)
                 else:
+                    # Compile
+                    cc = (exe,) + compiler_flags + ('-c', '-o', str(oname), str(cname))
+                    _run(cc, logfile, errfile)
+                    # Extract linker specific "cflags" from ldflags and link
+                    ld = tuple(shlex.split(compiler.ld)) + ('-o', str(soname), str(oname)) + tuple(expandWl(compiler.ldflags))
+                    _run(ld, logfile, errfile, step="Linker", filemode="a")
+        except subprocess.CalledProcessError as e:
+            msg = dedent(f"""
+                Command "{e.cmd}" return error status {e.returncode}.
+                Unable to compile code
+            """)
+            if os.environ.get("FIREDRAKE_CI", False):
+                msg += dedent(f"""
+                    Code is:
+                    {code}
+                """)
+                with open(errfile) as err:
                     msg += dedent(f"""
-                        Compile log in {logfile!s}
-                        Compile errors in {errfile!s}
+                        Compiler output is:
+                        {''.join(err.readlines())}
                     """)
-                result = CompilationError(msg)
-                result.__cause__ = e  # equivalent to 'raise XXX from e'
-            except BaseException as e:
-                # catch and broadcast all exceptions to prevent deadlocks
-                result = e
-    else:
-        result = None
+            else:
+                msg += dedent(f"""
+                    Compile log in {logfile!s}
+                    Compile errors in {errfile!s}
+                """)
+            raise CompilationError(msg) from e
+        else:
+            return soname
 
-    result = ccomm.bcast(result)
-    if isinstance(result, BaseException):
-        raise result
-    else:
-        return result
+    return mpi.safe_noncollective(ccomm, compile_single_rank, root=0)
 
 
 def _run(cc, logfile, errfile, step="Compilation", filemode="w"):

pyop2/mpi.py

@@ -34,6 +34,7 @@
 """PyOP2 MPI communicator."""
 
 
+from typing import Any, Callable
 from petsc4py import PETSc
 from mpi4py import MPI  # noqa
 from itertools import count
@@ -565,6 +566,44 @@ def finalize_safe_debug():
     return debug
 
 
+def safe_noncollective(comm: MPI.Comm, func: Callable[[], Any], *, root: int) -> Any:
+    """Run a function on a single rank of ``comm`` in a deadlock safe way.
+
+    If an exception is raised on the active rank then this is caught and
+    raised collectively.
+
+    Parameters
+    ----------
+    comm
+        The communicator.
+    func
+        The operation to be performed on a single rank. This should be a
+        callable that takes no arguments.
+    root
+        The rank performing the operation.
+
+    Returns
+    -------
+    Any
+        The result of ``func``, broadcasted to all ranks.
+
+    """
+    if comm.rank == root:
+        try:
+            result = func()
+        except BaseException as e:
+            result = e
+    else:
+        result = None
+
+    with temp_internal_comm(comm) as icomm:
+        result = icomm.bcast(result, root=root)
+    if isinstance(result, BaseException):
+        raise result
+    else:
+        return result
+
+
 @atexit.register
 def _free_comms():
     """Free all outstanding communicators."""

tests/pyop2/test_mpi.py

@@ -0,0 +1,26 @@
+import pytest
+import pytest_mpi
+from mpi4py import MPI
+
+import pyop2.mpi
+
+
+def passing_test():
+    return "pass"
+
+
+def failing_test():
+    raise RuntimeError("This test has failed")
+
+
+@pytest.mark.parallel(2)
+@pytest.mark.parametrize("root", [0, 1])
+def test_branches_on_rank_do_not_deadlock(root):
+    result = pyop2.mpi.safe_noncollective(MPI.COMM_WORLD, passing_test, root=root)
+    pytest_mpi.parallel_assert(result == "pass")
+
+    try:
+        result = pyop2.mpi.safe_noncollective(MPI.COMM_WORLD, failing_test, root=root)
+    except BaseException as e:
+        result = e
+    pytest_mpi.parallel_assert(isinstance(result, RuntimeError) and str(result) == "This test has failed")