Add a pytest plugin for testflo-style MPI test execution

[robfalck]

Summary

Adds testflo.pytest_plugin, a pytest plugin (registered via the pytest11 entry point, no conftest.py needed) that brings testflo's MPI execution model — spawn mpirun -n N, run the test on every rank, gather results — into pytest suites, including full pytest-xdist support. Ships alongside testflo rather than replacing it; testflo itself takes no new dependency since pytest/mpi4py/psutil are all optional extras (testflo[pytest]).

Why

testflo's own runner handles MPI tests well pytest has become the standard. Teams that have already standardized on pytest (fixtures, markers, -k/-m selection, xdist, JUnit XML, pytest-cov, IDE integration) had no way to get testflo's "just spawn mpirun and run the test on every rank" model without leaving pytest. This plugin gives them that, while staying optional and backward compatible with the existing testflo runner.

Now that we're also looking into multiprocessing, we needed a tool that could schedule tests and dedicate a given number of processors to them, either under an MPI environment or just a python environment with multiprocessing used within.

A few existing mpi-related pytest plugins either relied on the pytest invocation itself to be done under MPI, or spawned tests in a monolithic MPI environment.

Marking tests for parallelization

Two independent markers, so they can be selected separately:

import pytest

@pytest.mark.mpi(4)
def test_foo(comm):
    assert comm.size == 4

@pytest.mark.mpi([2, 4])          # parametrizes into test_bar[nprocs=2], test_bar[nprocs=4]
def test_bar(comm):
    assert comm.allreduce(1) == comm.size

@pytest.mark.multiprocessing(4)   # no mpirun spawned; runs in-process on a size-1 FakeComm
def test_pool():
    with multiprocessing.Pool(4) as pool:
        ...

@pytest.mark.mpi(N) is canonical; nprocs=N is an accepted alias. A bare @pytest.mark.mpi defaults to 2 ranks. multiprocessing has no default — a bare @pytest.mark.multiprocessing is a UsageError, since there's no sensible pool size to assume. The two compose: mpi(4) + multiprocessing(2) costs 4 * 2 = 8 cores (each rank spawns its own pool).

Existing testflo-style N_PROCS unittest.TestCase classes work unchanged — N_PROCS is treated as an alias for mpi(nprocs=N_PROCS), including for -m mpi selection.

pytest -m 'not mpi' runs everything that doesn't need mpirun (multiprocessing-only tests included) — useful for a fast pass that still exercises pool scaling without MPI's overhead or hang risk.

Execution modes

The plugin detects which of three contexts it's running in purely from environment variables (never by importing mpi4py in the host process — that would initialize MPI and make the host unsafe to fork mpirun from):

flowchart TD
    start([pytest process starts]) --> childcheck{TESTFLO_PYTEST_CHILD=1?}
    childcheck -- yes --> child["<b>Child mode</b><br/>test runs normally on this rank;<br/>results gathered on COMM_WORLD<br/>at sessionfinish, rank 0 writes JSON"]
    childcheck -- no --> outercheck{"OMPI_COMM_WORLD_SIZE / PMI_SIZE /<br/>MV2_COMM_WORLD_SIZE / MSMPI_RANK_SIZE<br/>&gt; 1?"}
    outercheck -- yes --> outer["<b>Outer-mpirun mode</b><br/>user ran `mpirun -n N pytest ...` themselves;<br/>tests whose mpi size == N run in place,<br/>everything else is skipped;<br/>xdist is rejected (UsageError)"]
    outercheck -- no --> launcher["<b>Launcher mode</b><br/>(the normal case)"]
    launcher --> spec{"item's mpi size &gt; 1?"}
    spec -- yes --> mpi["reserve cores, then spawn<br/>`mpirun -n N python -m pytest &lt;nodeid&gt;`,<br/>synthesize TestReports from the result"]
    spec -- no --> serial["reserve cores (if under xdist),<br/>run normally in this process<br/>(FakeComm size 1, or real COMM_WORLD<br/>if MPI happens to already be initialized)"]

Loading

Node IDs are identical between launcher and child invocations, so -k/-m selection, parametrize IDs, and JUnit XML all look like ordinary pytest output — a 4-rank MPI test just shows up as one test with per-rank tracebacks and captured output attached if it fails.

Scheduling under pytest-xdist

Every test — serial or parallel — is distributed freely by xdist; there's no special grouping. Instead, each test reserves the cores it needs from a budget shared across all workers just before it runs, so a big MPI test simply waits until enough cores are free while serial tests keep flowing around it. The budget is hosted by a multiprocessing.managers server the xdist controller starts (pytest_configure/pytest_configure_node); workers connect to it over loopback, so a waiting acquire() blocks on a condition variable in the server rather than polling.

Requests are served FIFO with one relaxation (a request may pass ones queued ahead of it as long as it still leaves room for the oldest), so a big test can't be starved indefinitely but small tests still flow. Because xdist can't hand a blocked test back once a worker has committed to it, the plugin defaults --dist to worksteal whenever xdist is active and the user hasn't set --dist explicitly — idle workers steal the tests queued behind a blocked one instead of everyone stalling.

A worker that dies mid-test (holding or waiting for cores) is detected and reaped once a second so it can't starve the rest of the run. A test whose cost alone exceeds the budget fails immediately instead of hanging forever.

Running it

# plain run — no setup needed, the plugin registers itself
pytest

# only tests that don't need mpirun (still runs multiprocessing-marked tests)
pytest -m 'not mpi'

# parallel run; defaults to --dist worksteal automatically
pytest -n auto

# cap or remove the core budget
pytest -n 4 --max-concurrent-cores=8
pytest --oversubscribe                      # or TESTFLO_PYTEST_OVERSUBSCRIBE=1

# run mpi-marked tests in-process on a size-1 FakeComm (no MPI needed at all)
pytest --nompi

# guard against deadlocks from a rank-local failure before a collective call
pytest --mpi-timeout=60

# point at an mpirun/mpiexec not on PATH
pytest --mpirun-exe=/opt/openmpi/bin/mpirun

Running with coverage

MPI-marked tests execute inside a spawned mpirun, never in the pytest process that reports them, so coverage has to be collected in the ranks and merged back in. This rides entirely on coverage.py's own subprocess support — nothing here is pytest-cov-specific, so pytest --cov=... and coverage run -m pytest behave the same:

# .coveragerc, or [tool.coverage.run] in pyproject.toml
[run]
patch = subprocess        # measure spawned mpirun ranks; implies parallel = true
source = mypkg

patch = subprocess makes coverage export a serialized config that the .pth file it installs picks up to start tracing in every freshly-spawned rank. parallel = true is forced on by it and is required — without it every rank writes the same data file and they clobber each other.

If a test also spawns its own worker pool (@pytest.mark.multiprocessing(M), including inside an mpi rank), add _exit too — a pool worker exits via os._exit, which skips the atexit hook coverage normally saves from:

[run]
patch = _exit, subprocess

pytest --cov=mypkg --cov-report=term-missing
# or
coverage run -m pytest
coverage combine     # same as you'd already do with xdist
coverage report

If none of this is configured, the plugin falls back to synthesizing the same COVERAGE_PROCESS_CONFIG environment for the ranks from whatever coverage is already running in the launcher, so a bare pytest --cov=mypkg still measures MPI-only code instead of silently reporting 0%. The one case this can't cover: a rank killed by --mpi-timeout or MPI_Abort never reaches its own save, so its coverage is missing (not corrupted) for that run.

Running with and without pytest-sugar

The plugin reports MPI results by synthesizing normal TestReport objects through the standard pytest_runtest_logreport hook, so any terminal reporter that just listens to that hook — including pytest-sugar — sees ordinary test outcomes and needs no awareness of this plugin:

pip install pytest-sugar
pytest -n auto                 # sugar's progress bar, same synthesized MPI reports
pytest -n auto -p no:sugar     # fall back to plain dot/verbose output for one run (useful for CI)

Known limitations

• MPICH on macOS: a networking issue in MPICH's OFI provider can surface as OFI poll failed (... Input/output error) during MPI finalization when ranks are forcefully killed (e.g. a --mpi-timeout SIGKILL). Not seen in normal use; switch to Open MPI if it shows up in stress/timeout-heavy runs.
• A test whose core cost alone exceeds --max-concurrent-cores is reported as a failure rather than run anyway.

Other changes bundled in this branch

• CI (test_workflow.yml): matrix expanded to Ubuntu × macOS, each against both Open MPI and MPICH, via pixi environments; added a dedicated "Test pytest plugin" step (pytest ... -n auto, run once normally and once with --nompi).
• testflo's own self-test suite: previously, testflo's self-tests validated its pass/fail/skip counting by mixing genuinely-passing checks with tests designed to fail/skip/xfail in the same files, so a plain testflo . always "failed" even when everything worked. Moved the intentionally-broken fixtures to self_report_fixtures.py (outside testflo's default discovery pattern) and added test_self_test.py, which runs testflo against those fixtures as a subprocess and asserts the counts — so a standard testflo run now legitimately exits 0.
• Fixed TESTFLO_RUNNING not being set under the pytest plugin, so tests written against that env var behave the same under testflo and pytest.

Test plan

• pytest testflo/tests/pytest_tests/test_pytest_plugin.py -n auto — 45 tests covering marker parsing/validation, core-budget fairness and starvation-avoidance, xdist worksteal defaulting, child-mode and outer-mpirun-mode result handling, coverage propagation (patch=subprocess and the fallback synthesis, plus multiprocessing + _exit), and --nompi/--mpi-timeout/--oversubscribe.
• Same suite run again with pytest-sugar installed — no change in outcome.
• CI matrix: Ubuntu + macOS, Open MPI + MPICH.

Related Issues

• Resolves #

Backwards incompatibilities

None

New Dependencies

None