Add MPI world-size preflight check: fail loudly on "rank 0 of N singletons"

When the host MPI launcher and the container's MPI/PMIx stack are incompatible,
every process initialises standalone (COMM_WORLD size 1, rank 0). ShapePipe
then treats each as master, hands each the full job list, and runs N
uncoordinated copies of the pipeline into the same output directory — silently,
with exit 0. This is the failure the OpenMPI-5 image fix prevents on candide,
but nothing guarded against a future recurrence on another cluster.

check_mpi_world() compares the size that actually wired up (COMM_WORLD) against
the size the launcher intended (OMPI_COMM_WORLD_SIZE, which is set per process
even when the world fails to form) and aborts on a mismatch. Empirically
verified on candide: SLURM_NTASKS is NOT reliable for this (reads 1 on
remote-node ranks even in a healthy run) — OMPI_COMM_WORLD_SIZE is. Tested both
ways on a real allocation: healthy OMPI-5 run passes and completes; OMPI-4
image under the OMPI-5 host launcher fires the check and exits non-zero
(together with the exit-code fix). Also catches partial wire-up (N-1 of N).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: cailmdaley

src/shapepipe/pipeline/mpi_run.py

@@ -6,9 +6,52 @@
 
 """
 
+import os
+
 from shapepipe.pipeline.worker_handler import WorkerHandler
 
 
+def check_mpi_world(comm):
+    """Check MPI World.
+
+    Verify that the MPI world formed at the size the launcher requested, and
+    abort loudly otherwise.
+
+    This guards against the "N rank-0 singletons" failure: when the host MPI
+    launcher and the container's MPI / PMIx stack are incompatible, each
+    process initialises standalone (``COMM_WORLD`` size 1, rank 0). ShapePipe
+    would then treat every process as master, hand each the full job list, and
+    run that many uncoordinated copies of the pipeline into the same output
+    directory -- silently, with exit code 0. Comparing the size that actually
+    wired up against the size the launcher intended (``OMPI_COMM_WORLD_SIZE``,
+    which is set per process even when the world fails to form) turns that
+    silent corruption into an immediate, legible error.
+
+    Parameters
+    ----------
+    comm : MPI.Comm
+        MPI communicator instance (``MPI.COMM_WORLD``)
+
+    Raises
+    ------
+    RuntimeError
+        if the launcher requested more ranks than actually wired up
+
+    """
+    intended = os.environ.get("OMPI_COMM_WORLD_SIZE")
+    actual = comm.Get_size()
+
+    if intended is not None and int(intended) != actual:
+        raise RuntimeError(
+            f"MPI world mismatch: the launcher requested {intended} ranks but "
+            f"only {actual} wired up (MPI_COMM_WORLD size {actual}). This is "
+            f"the 'rank 0 of N singletons' failure -- the host MPI launcher and "
+            f"the container's MPI/PMIx stack are incompatible. Aborting rather "
+            f"than running {intended} uncoordinated copies of the pipeline into "
+            f"the same output directory."
+        )
+
+
 def split_mpi_jobs(jobs, batch_size):
     """Split MPI Jobs.
 

src/shapepipe/run.py

@@ -20,7 +20,11 @@
 from shapepipe.pipeline.dependency_handler import DependencyHandler
 from shapepipe.pipeline.file_handler import FileHandler
 from shapepipe.pipeline.job_handler import JobHandler
-from shapepipe.pipeline.mpi_run import split_mpi_jobs, submit_mpi_jobs
+from shapepipe.pipeline.mpi_run import (
+    check_mpi_world,
+    split_mpi_jobs,
+    submit_mpi_jobs,
+)
 
 try:
     from mpi4py import MPI
@@ -372,6 +376,11 @@ def run_mpi(pipe, comm):
     # Assign master node
     master = comm.rank == 0
 
+    # Fail loudly if the MPI world did not form at the size the launcher
+    # requested (the "rank 0 of N singletons" launcher/container mismatch),
+    # rather than silently running redundant copies of the pipeline.
+    check_mpi_world(comm)
+
     # Get the module to be run
     modules = pipe.modules if master else None
     modules = comm.bcast(modules, root=0)

tests/unit/test_mpi_world.py

@@ -0,0 +1,62 @@
+"""Guard against the silent "rank 0 of N singletons" MPI failure.
+
+When the host MPI launcher and the container's MPI/PMIx stack are
+incompatible, every process initialises standalone (``COMM_WORLD`` size 1),
+and ShapePipe would otherwise run N uncoordinated copies of the pipeline into
+the same output directory -- silently, with exit code 0.
+``check_mpi_world`` turns that into an immediate error by comparing the size
+that actually wired up against the launcher's intended ``OMPI_COMM_WORLD_SIZE``.
+The intended/actual pairs below are the values measured on candide for a
+healthy run and for the OpenMPI-4-container / OpenMPI-5-host mismatch.
+"""
+
+import pytest
+
+from shapepipe.pipeline.mpi_run import check_mpi_world
+
+
+class _FakeComm:
+    """Minimal stand-in exposing only ``Get_size``."""
+
+    def __init__(self, size):
+
+        self._size = size
+
+    def Get_size(self):
+
+        return self._size
+
+
+@pytest.mark.parametrize(
+    "intended, actual",
+    [
+        ("4", 4),  # healthy multi-node run
+        (None, 1),  # no launcher env -> legitimate single-rank run
+        ("1", 1),  # explicit single rank
+    ],
+    ids=["healthy", "no-launcher-env", "single-rank"],
+)
+def test_check_mpi_world_passes(monkeypatch, intended, actual):
+
+    if intended is None:
+        monkeypatch.delenv("OMPI_COMM_WORLD_SIZE", raising=False)
+    else:
+        monkeypatch.setenv("OMPI_COMM_WORLD_SIZE", intended)
+
+    check_mpi_world(_FakeComm(actual))
+
+
+@pytest.mark.parametrize(
+    "intended, actual",
+    [
+        ("4", 1),  # the measured singleton failure
+        ("4", 3),  # partial wire-up
+    ],
+    ids=["singletons", "partial-wireup"],
+)
+def test_check_mpi_world_aborts_on_mismatch(monkeypatch, intended, actual):
+
+    monkeypatch.setenv("OMPI_COMM_WORLD_SIZE", intended)
+
+    with pytest.raises(RuntimeError, match="MPI world mismatch"):
+        check_mpi_world(_FakeComm(actual))