OpenMPI forking mode support (#21)

* pytest-plugin: support openmpi

* Add parallel marker tests

---------

Co-authored-by: qbisi <[email protected]>

Author: connorjward

.github/workflows/ci.yml

@@ -0,0 +1,55 @@
+name: CI
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+  schedule:
+    - cron: '1 5 * * 1'
+
+jobs:
+  tests:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python: ['3.8', '3.9', '3.10', '3.11', '3.12', '3.13']
+        mpi: ['mpich', 'openmpi']
+    defaults:
+      run:
+        shell: bash -l {0}
+    env:
+      PYTEST_MPI_MAX_NPROCS: 3
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install Conda environment with Micromamba
+        uses: mamba-org/setup-micromamba@v1
+        with:
+          environment-file: ".github/etc/test_environment_${{ matrix.mpi }}.yml"
+          create-args: >-
+              python=${{ matrix.python }}
+
+      - name: Install mpi-pytest
+        run: pip install --no-deps -e .
+
+      - name: Run tests (MPICH)
+        if: matrix.mpi == 'mpich'
+        run: |
+          : # 'forking' mode
+          pytest -v tests
+          : # 'non-forking' mode
+          mpiexec -n 1 pytest -v -m "not parallel or parallel[1]" tests
+          mpiexec -n 2 pytest -v -m parallel[2] tests
+          mpiexec -n 3 pytest -v -m parallel[3] tests
+
+      - name: Run tests (OpenMPI)
+        if: matrix.mpi == 'openmpi'
+        run: |
+          : # 'forking' mode
+          pytest -v tests
+          : # 'non-forking' mode
+          mpiexec --oversubscribe -n 1 pytest -v -m "not parallel or parallel[1]" tests
+          mpiexec --oversubscribe -n 2 pytest -v -m parallel[2] tests
+          mpiexec --oversubscribe -n 3 pytest -v -m parallel[3] tests

.github/workflows/ci_pipeline.yml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "mpi-pytest"
 # <year.month.patch>
-version = "2025.4.0"
+version = "2025.5.0.dev0"
 dependencies = ["mpi4py", "pytest"]
 authors = [
   { name="Connor Ward", email="[email protected]" },
@@ -14,7 +14,7 @@ authors = [
 description = "A pytest plugin for executing tests in parallel with MPI"
 readme = "README.md"
 license = { file = "LICENSE" }
-requires-python = ">=3.7"
+requires-python = ">=3.8"
 classifiers = [
     "Programming Language :: Python :: 3",
     "Framework :: Pytest",

pyproject.toml

@@ -1,4 +1,5 @@
 import collections
+import enum
 import numbers
 import os
 import subprocess
@@ -230,9 +231,14 @@ def _set_parallel_callback(item):
         "--disable-warnings", "--show-capture=no"
     ]
 
-    cmd = [
-        "mpiexec", "-n", "1", "-genv", CHILD_PROCESS_FLAG, "1", *executable
-    ] + pytest_args + [
+    impl = detect_mpi_implementation()
+    if impl == MPIImplementation.OPENMPI:
+        cmd = ["mpiexec", "-n", "1", "-x", f"{CHILD_PROCESS_FLAG}=1", *executable]
+    else:
+        assert impl == MPIImplementation.MPICH
+        cmd = ["mpiexec", "-n", "1", "-genv", CHILD_PROCESS_FLAG, "1", *executable]
+
+    cmd += pytest_args + [
         ":", "-n", f"{nprocs-1}", *executable
     ] + quieter_pytest_args
 
@@ -286,3 +292,35 @@ def _parse_marker_nprocs(marker):
 
 def _as_tuple(arg):
     return tuple(arg) if isinstance(arg, collections.abc.Iterable) else (arg,)
+
+
+class MPIImplementation(enum.Enum):
+    OPENMPI = enum.auto()
+    MPICH = enum.auto()
+
+
+def detect_mpi_implementation() -> MPIImplementation:
+    try:
+        result = subprocess.run(
+            ["mpiexec", "--version"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+            check=True
+        )
+    except FileNotFoundError:
+        raise FileNotFoundError(
+            "'mpiexec' not found on your PATH, please run in non-forking mode "
+            "where you can specify a different MPI executable"
+        )
+
+    output = result.stdout.lower()
+    if "open mpi" in output or "open-rte" in output:
+        return MPIImplementation.OPENMPI
+    elif "mpich" in output:
+        return MPIImplementation.MPICH
+    else:
+        raise RuntimeError(
+            "MPI distribution is not recognised, please run in non-forking "
+            "mode where you can specify your MPI executable"
+        )

pytest_mpi/plugin.py

@@ -0,0 +1,17 @@
+import pytest
+from mpi4py import MPI
+
+
+@pytest.mark.parallel
+def test_parallel_marker_no_args():
+    assert MPI.COMM_WORLD.size == 3
+
+
+@pytest.mark.parallel(2)
+def test_parallel_marker_with_int():
+    assert MPI.COMM_WORLD.size == 2
+
+
+@pytest.mark.parallel([2, 3])
+def test_parallel_marker_with_list():
+    assert MPI.COMM_WORLD.size in {2, 3}