Terminate all processes on error in parallel_run_from_cli

CHANGELOG.md

@@ -15,6 +15,7 @@ All notable changes to this project will be documented in this file.
 
 ### Fixed
 
+-   #624 : Terminate all MPI processes when one of them fails in `parallel_run_from_cli`, see #621
 -   #576 : Fix bug in `TensorFemSpace.eval_field` caused by round-off at MPI subdomain boundaries
 -   #576 : Require `sympde==0.19.3` which fixes a bug in the linearity checks
 -   #579 : Require `h5py>=3.16` which installs correctly with `setuptools>=81.0`

docs/source/examples.rst

@@ -46,3 +46,13 @@ These examples include:
 Performance
 -----------
 There are also some examples in the `performance directory <https://github.com/pyccel/psydac/tree/devel/examples/performance>`_ explaining the assembly algorithms used in PSYDAC.
+
+Running examples in parallel
+----------------------------
+We very strongly recommend that every parallel script runs its main function through `psydac.utilities.parallel_utils.parallel_run_from_cli`, as done by the examples in `psydac/feec/polar/examples`.
+
+Without it, an uncaught exception terminates only the process which raised it, while the other processes hang as soon as they wait for a message from it (e.g. in a collective operation).
+On a cluster, such a job keeps its allocation until it reaches the wall-time limit.
+With `parallel_run_from_cli`, the traceback is printed and all the processes are terminated by calling `MPI_Abort`, whichever MPI launcher is used (e.g. `mpirun`, `mpiexec`, or `srun` in a batch job).
+
+Note that `MPI_Abort` terminates the other processes without any cleanup (`finally` blocks, context managers, `atexit` handlers), hence an HDF5 file which is being written may be left incomplete.

psydac/utilities/parallel_utils.py

@@ -19,6 +19,12 @@ def parallel_run_from_cli(parser, main):
     CLI arguments. When not run (serially) from IPython, it also calls
     `plt.show()` to keep any Matplotlib windows open after `main` is run.
 
+    If `main` raises an uncaught exception on any process, all the processes
+    are terminated by calling `MPI_Abort` at exit, as done by the launcher
+    `python -m mpi4py`. Otherwise, the other processes would hang as soon as
+    they wait for a message from the failing one. The exception is re-raised,
+    hence its traceback is printed and the failing process is cleaned up.
+
     Parameters
     ----------
     parser : callable
@@ -83,8 +89,15 @@ def parallel_run_from_cli(parser, main):
     # Store local variables in a dictionary
     namespace = locals()
 
-    # Run simulation in parallel, and update local variables
-    namespace.update(main(**args_dict))
+    # Run simulation in parallel, and update local variables. In case of an
+    # uncaught exception, terminate all processes at exit (see #621)
+    try:
+        namespace.update(main(**args_dict))
+    except BaseException as exc:
+        from mpi4py.run import set_abort_status
+
+        set_abort_status(exc)
+        raise
 
     # Keep matplotlib windows open (not necessary in IPython sessions)
     try:

psydac/utilities/tests/__init__.py

@@ -0,0 +1,5 @@
+#---------------------------------------------------------------------------#
+# This file is part of PSYDAC which is released under MIT License. See the  #
+# LICENSE file or go to https://github.com/pyccel/psydac/blob/devel/LICENSE #
+# for full license details.                                                 #
+#---------------------------------------------------------------------------#

psydac/utilities/tests/test_parallel_utils.py

@@ -0,0 +1,47 @@
+#---------------------------------------------------------------------------#
+# This file is part of PSYDAC which is released under MIT License. See the  #
+# LICENSE file or go to https://github.com/pyccel/psydac/blob/devel/LICENSE #
+# for full license details.                                                 #
+#---------------------------------------------------------------------------#
+import argparse
+
+import pytest
+
+from psydac.utilities.parallel_utils import parallel_run_from_cli
+
+
+@pytest.fixture
+def abort_statuses(monkeypatch) -> list:
+    """Record the calls to `mpi4py.run.set_abort_status`, which would
+    otherwise terminate the test session at exit."""
+    statuses = []
+    monkeypatch.setattr('mpi4py.run.set_abort_status', statuses.append)
+    # Figures left open by other tests would block `plt.show()` with an interactive backend
+    monkeypatch.setattr('matplotlib.pyplot.show', lambda: None)
+    return statuses
+
+def parser() -> argparse.Namespace:
+    return argparse.Namespace(n=3)
+
+#==============================================================================
+def test_parallel_run_from_cli(abort_statuses: list) -> None:
+
+    def main(*, n, mpi_comm):
+        squares = [i**2 for i in range(n)]
+        return locals()
+
+    namespace = parallel_run_from_cli(parser, main)
+
+    assert namespace['squares'] == [0, 1, 4]
+    assert abort_statuses == []
+
+def test_parallel_run_from_cli_abort(abort_statuses: list) -> None:
+
+    def main(*, n, mpi_comm):
+        raise RuntimeError('failure in main')
+
+    # The exception is re-raised, and all processes will be aborted at exit
+    with pytest.raises(RuntimeError, match='failure in main') as excinfo:
+        parallel_run_from_cli(parser, main)
+
+    assert abort_statuses == [excinfo.value]