Console scripts: abort all ranks on an uncaught exception (issue #807)

.github/workflows/test_pr_and_main.yml

@@ -103,6 +103,10 @@ jobs:
         run: |
           coverage run $COV_ARGS -m pytest mpisppy/tests/test_cvar.py -v
 
+      - name: Test entry points
+        run: |
+          coverage run $COV_ARGS -m pytest mpisppy/tests/test_entry_points.py -v
+
       - name: Test outer_bound_only
         run: |
           coverage run $COV_ARGS -m pytest mpisppy/tests/test_outer_bound_only.py -v

doc/src/debug_utils.rst

@@ -89,7 +89,7 @@ printed when the report is not ok:
 
 .. code-block:: bash
 
-   mpiexec -np N python my_driver.py --inspect-buffers-on-shutdown
+   mpiexec -np N python -m mpi4py my_driver.py --inspect-buffers-on-shutdown
 
 When the flag is unset (the default), the inspector is never called
 and the shutdown-poll cost is unchanged.

doc/src/generic_cylinders.rst

@@ -18,12 +18,12 @@ other features without requiring you to write a driver program.
 
        mpi-sppy-generic-cylinders --module-name farmer --num-scens 3 --EF --EF-solver-name gurobi
 
-   For multi-rank parallel runs, prefer the ``python -m mpi4py`` module
-   form shown below (``mpiexec -np 3 python -m mpi4py -m
-   mpisppy.generic_cylinders ...``): mpi4py's runner aborts all ranks if
-   one raises an exception, whereas a bare
-   ``mpiexec -np 3 mpi-sppy-generic-cylinders ...`` runs but may leave
-   the other ranks hanging when one fails.
+   The console script installs the same abort-on-exception protection
+   as mpi4py's runner, so for multi-rank parallel runs
+   ``mpiexec -np 3 mpi-sppy-generic-cylinders ...`` and the
+   ``mpiexec -np 3 python -m mpi4py -m mpisppy.generic_cylinders ...``
+   module form shown below are equally safe: if one rank raises an
+   exception, all ranks are aborted rather than leaving the job hung.
 
 Your Model File (Module)
 ------------------------
@@ -154,7 +154,7 @@ two-stage problems:
 
 .. code-block:: bash
 
-    mpiexec -np 2 python -m mpisppy.generic_cylinders \
+    mpiexec -np 2 python -m mpi4py -m mpisppy.generic_cylinders \
         --module-name farmer --num-scens 3 \
         --solver-name gurobi --lshaped-hub --xhatlshaped \
         --max-iterations 100 --rel-gap 1e-4

doc/src/seqsamp.rst

@@ -24,9 +24,10 @@ a ``--module-name`` argument pointing to your model module.
    ``pip install -e .`` from a clone) puts the console script
    ``mpi-sppy-mrp-generic`` on your ``PATH``; it is equivalent to
    ``python -m mpisppy.mrp_generic`` and can be used in place of that
-   prefix in the examples below. For the ``--xhat-method cylinders``
-   (multi-rank) case, prefer the ``python -m mpi4py`` module form so
-   that mpi4py's abort-on-exception handler is installed.
+   prefix in the examples below; the console script installs the same
+   abort-on-exception handler as the ``python -m mpi4py`` module form,
+   so either is safe for the ``--xhat-method cylinders`` (multi-rank)
+   case.
 
 Your model module must provide the same functions required by
 ``generic_cylinders``: ``scenario_creator``, ``scenario_names_creator``,

mpisppy/entry_points.py

@@ -0,0 +1,53 @@
+###############################################################################
+# mpi-sppy: MPI-based Stochastic Programming in PYthon
+#
+# Copyright (c) 2024, Lawrence Livermore National Security, LLC, Alliance for
+# Sustainable Energy, LLC, The Regents of the University of California, et al.
+# All rights reserved. Please see the files COPYRIGHT.md and LICENSE.md for
+# full copyright and license information.
+###############################################################################
+"""Console-script wrappers that add mpi4py-style abort on uncaught exceptions.
+
+pip's console entry points bypass ``python -m mpi4py``, whose runner prints
+the traceback and calls MPI_Abort when a rank dies; without that, the
+surviving ranks block forever in a collective and the whole mpiexec job
+hangs.  These wrappers give the entry points the same protection, so
+``mpiexec -np 3 mpi-sppy-generic-cylinders ...`` is as safe as the
+``python -m mpi4py -m mpisppy.generic_cylinders`` form.  Serial runs
+(and the no-mpi4py mock in mpisppy.MPI) re-raise normally so tracebacks
+and exit codes are unchanged.
+"""
+
+import traceback
+
+from mpisppy import MPI
+
+
+def _run_with_mpi_abort(real_main):
+    try:
+        real_main()
+    except SystemExit:
+        # argparse exits (--help, bad flags) happen identically on every
+        # rank, so a plain exit cannot strand the others.
+        raise
+    except BaseException:
+        comm = MPI.COMM_WORLD
+        if comm.Get_size() > 1 and hasattr(comm, "Abort"):
+            traceback.print_exc()
+            comm.Abort(1)
+        raise
+
+
+def generic_cylinders_main():
+    from mpisppy.generic_cylinders import main
+    _run_with_mpi_abort(main)
+
+
+def mrp_generic_main():
+    from mpisppy.mrp_generic import main
+    _run_with_mpi_abort(main)
+
+
+def one_sided_test_main():
+    from mpi_one_sided_test import main
+    _run_with_mpi_abort(main)

mpisppy/tests/test_entry_points.py

@@ -0,0 +1,139 @@
+###############################################################################
+# mpi-sppy: MPI-based Stochastic Programming in PYthon
+#
+# Copyright (c) 2024, Lawrence Livermore National Security, LLC, Alliance for
+# Sustainable Energy, LLC, The Regents of the University of California, et al.
+# All rights reserved. Please see the files COPYRIGHT.md and LICENSE.md for
+# full copyright and license information.
+###############################################################################
+# Serial tests for the console-script wrappers in mpisppy.entry_points.
+
+import io
+import sys
+import contextlib
+import unittest
+
+import mpisppy.MPI as MPI
+import mpisppy.entry_points as entry_points
+
+try:
+    import mpi4py  # noqa: F401
+    have_mpi4py = True
+except ImportError:
+    have_mpi4py = False
+
+
+class _FakeComm:
+    """Stands in for MPI.COMM_WORLD so no test ever really aborts."""
+    def __init__(self, size):
+        self._size = size
+        self.abort_code = None
+
+    def Get_size(self):
+        return self._size
+
+    def Abort(self, errorcode=0):
+        self.abort_code = errorcode
+
+
+class _FakeCommNoAbort:
+    """Mimics the no-mpi4py mock comm, which has no Abort method."""
+    def __init__(self, size):
+        self._size = size
+
+    def Get_size(self):
+        return self._size
+
+
+class TestEntryPoints(unittest.TestCase):
+
+    def setUp(self):
+        self._saved_comm = MPI.COMM_WORLD
+
+    def tearDown(self):
+        MPI.COMM_WORLD = self._saved_comm
+
+    def _run_failing_main(self, comm, exc=ValueError):
+        MPI.COMM_WORLD = comm
+        def failing_main():
+            raise exc("boom")
+        with contextlib.redirect_stderr(io.StringIO()):
+            with self.assertRaises(exc):
+                entry_points._run_with_mpi_abort(failing_main)
+
+    def test_serial_reraises_without_abort(self):
+        comm = _FakeComm(1)
+        self._run_failing_main(comm)
+        self.assertIsNone(comm.abort_code)
+
+    def test_multirank_aborts(self):
+        comm = _FakeComm(3)
+        self._run_failing_main(comm)
+        self.assertEqual(comm.abort_code, 1)
+
+    def test_mock_comm_without_abort_reraises(self):
+        self._run_failing_main(_FakeCommNoAbort(3))
+
+    def test_system_exit_passes_through(self):
+        comm = _FakeComm(3)
+        self._run_failing_main(comm, exc=SystemExit)
+        self.assertIsNone(comm.abort_code)
+
+    def test_successful_main_runs_once(self):
+        MPI.COMM_WORLD = _FakeComm(3)
+        calls = []
+        entry_points._run_with_mpi_abort(lambda: calls.append(1))
+        self.assertEqual(calls, [1])
+
+    def test_console_script_targets_exist(self):
+        # the callables named in pyproject.toml [project.scripts]
+        for name in ("generic_cylinders_main",
+                     "mrp_generic_main",
+                     "one_sided_test_main"):
+            self.assertTrue(callable(getattr(entry_points, name)))
+
+
+class TestConsoleScriptWrappers(unittest.TestCase):
+    """Run each real console-script wrapper end-to-end on its cheap serial
+    path: with no arguments (and one rank) every wrapped main prints a usage
+    message and raises SystemExit, which the wrapper must pass through. This
+    executes the wrappers' real imports, so a broken import target (e.g. the
+    repo-root ``mpi_one_sided_test`` py-module going missing from the install)
+    fails here instead of on a user's command line."""
+
+    def _run_wrapper(self, wrapper, prog):
+        saved_argv = sys.argv
+        sys.argv = [prog]
+        try:
+            with contextlib.redirect_stdout(io.StringIO()), \
+                    contextlib.redirect_stderr(io.StringIO()):
+                with self.assertRaises(SystemExit) as caught:
+                    wrapper()
+            return caught.exception.code
+        finally:
+            sys.argv = saved_argv
+
+    def test_generic_cylinders_usage_exit(self):
+        # no args -> generic_cylinders.main prints usage and quit()s
+        code = self._run_wrapper(entry_points.generic_cylinders_main,
+                                 "mpi-sppy-generic-cylinders")
+        self.assertIn(code, (None, 0))
+
+    def test_mrp_generic_usage_exit(self):
+        # no args -> mrp_generic.main prints usage and quit()s
+        code = self._run_wrapper(entry_points.mrp_generic_main,
+                                 "mpi-sppy-mrp-generic")
+        self.assertIn(code, (None, 0))
+
+    @unittest.skipUnless(have_mpi4py, "mpi_one_sided_test imports mpi4py directly")
+    @unittest.skipUnless(MPI.COMM_WORLD.Get_size() == 1,
+                         "at more than one rank the one-sided test really runs")
+    def test_one_sided_test_single_rank_exit(self):
+        # at one rank the script demands an mpiexec launch and exits 2
+        code = self._run_wrapper(entry_points.one_sided_test_main,
+                                 "mpi-sppy-one-sided-test")
+        self.assertEqual(code, 2)
+
+
+if __name__ == "__main__":
+    unittest.main()

pyproject.toml

@@ -37,13 +37,15 @@ tracker = "https://github.com/Pyomo/mpi-sppy/issues"
 # Console entry points: pip creates a launcher on the user's PATH for each,
 # so pip-installed users can run these drivers without knowing where the
 # files landed in site-packages.  Each target is an importable module:callable.
-# For real multi-rank parallel runs the module form is still recommended so
-# the mpi4py runner installs its abort-on-exception handler, e.g.:
-#   mpiexec -np 3 python -m mpi4py -m mpisppy.generic_cylinders --module-name farmer ...
+# The mpisppy.entry_points wrappers replicate the mpi4py runner's
+# abort-on-exception behavior, so multi-rank runs of these scripts
+# (e.g. mpiexec -np 3 mpi-sppy-generic-cylinders ...) cannot hang on an
+# uncaught exception; the python -m mpi4py module form remains equivalent
+# for running from a checkout.
 [project.scripts]
-mpi-sppy-generic-cylinders = "mpisppy.generic_cylinders:main"
-mpi-sppy-mrp-generic = "mpisppy.mrp_generic:main"
-mpi-sppy-one-sided-test = "mpi_one_sided_test:main"
+mpi-sppy-generic-cylinders = "mpisppy.entry_points:generic_cylinders_main"
+mpi-sppy-mrp-generic = "mpisppy.entry_points:mrp_generic_main"
+mpi-sppy-one-sided-test = "mpisppy.entry_points:one_sided_test_main"
 
 [project.optional-dependencies]
 doc = [

run_coverage.bash

@@ -71,6 +71,9 @@ run_phase "test_maximization (serial)" \
 run_phase "test_cvar (serial)" \
     coverage run --rcfile=.coveragerc -m pytest mpisppy/tests/test_cvar.py -v
 
+run_phase "test_entry_points (serial)" \
+    coverage run --rcfile=.coveragerc -m pytest mpisppy/tests/test_entry_points.py -v
+
 run_phase "test_outer_bound_only (serial)" \
     coverage run --rcfile=.coveragerc -m pytest mpisppy/tests/test_outer_bound_only.py -v