mpi-pytest is now on pip (#4048)

Author: connorjward

pyproject.toml

@@ -67,11 +67,11 @@ spydump = "pyop2.scripts.spydump:main"
 
 [project.optional-dependencies]
 test = [
-  "pylit",
+  "mpi-pytest",
   "nbval",
+  "pylit",
   "pytest",
   "pytest-xdist",
-  "pytest-mpi @ git+https://github.com/firedrakeproject/pytest-mpi.git@main",
 ]
 dev = [
   "flake8",

requirements-ext.txt

@@ -3,6 +3,7 @@ Cython>=3.0
 decorator<=4.4.2
 flake8
 mpi4py
+mpi-pytest
 nbval
 numpy
 packaging

requirements-git.txt

@@ -2,6 +2,5 @@ git+https://github.com/firedrakeproject/ufl.git#egg=fenics-ufl
 git+https://github.com/firedrakeproject/fiat.git#egg=fenics-fiat
 git+https://github.com/dolfin-adjoint/pyadjoint.git#egg=pyadjoint-ad
 git+https://github.com/firedrakeproject/loopy.git@main#egg=loopy
-git+https://github.com/firedrakeproject/pytest-mpi.git@main#egg=pytest-mpi
 git+https://github.com/firedrakeproject/petsc.git@firedrake#egg=petsc
 git+https://github.com/firedrakeproject/libsupermesh.git#egg=libsupermesh

scripts/firedrake-run-split-tests

@@ -24,7 +24,7 @@
 #
 #   * pytest
 #   * pytest-split
-#   * pytest-mpi (https://github.com/firedrakeproject/pytest-mpi)
+#   * mpi-pytest
 #   * GNU parallel
 
 num_procs=$1

tests/firedrake/regression/test_ensembleparallelism.py

@@ -2,6 +2,7 @@
 from firedrake.petsc import DEFAULT_DIRECT_SOLVER_PARAMETERS
 from pyop2.mpi import MPI
 import pytest
+from pytest_mpi import parallel_assert
 
 from operator import mul
 from functools import reduce
@@ -25,24 +26,6 @@
             pytest.param(False, id="nonblocking")]
 
 
-def parallel_assert(assertion, subset=None, msg=""):
-    """ Move this functionality to pytest-mpi
-    """
-    if subset:
-        if MPI.COMM_WORLD.rank in subset:
-            evaluation = assertion()
-        else:
-            evaluation = True
-    else:
-        evaluation = assertion()
-    all_results = MPI.COMM_WORLD.allgather(evaluation)
-    if not all(all_results):
-        raise AssertionError(
-            "Parallel assertion failed on ranks: "
-            f"{[ii for ii, b in enumerate(all_results) if not b]}\n" + msg
-        )
-
-
 # unique profile on each mixed function component on each ensemble rank
 def function_profile(x, y, rank, cpt):
     return sin(cpt + (rank+1)*pi*x)*cos(cpt + (rank+1)*pi*y)
@@ -204,13 +187,13 @@ def test_ensemble_reduce(ensemble, mesh, W, urank, urank_sum, root, blocking):
     root_ranks = {ii + root*ensemble.comm.size for ii in range(ensemble.comm.size)}
     parallel_assert(
         lambda: error < 1e-12,
-        subset=root_ranks,
+        participating=COMM_WORLD.rank in root_ranks,
         msg=f"{error = :.5f}"  # noqa: E203, E251
     )
     error = errornorm(Function(W).assign(10), u_reduce)
     parallel_assert(
         lambda: error < 1e-12,
-        subset={range(COMM_WORLD.size)} - root_ranks,
+        participating=COMM_WORLD.rank not in root_ranks,
         msg=f"{error = :.5f}"  # noqa: E203, E251
     )
 
@@ -221,9 +204,7 @@ def test_ensemble_reduce(ensemble, mesh, W, urank, urank_sum, root, blocking):
     with u_reduce.dat.vec as v:
         states[spatial_rank] = v.stateGet()
     ensemble.comm.Allgather(MPI.IN_PLACE, states)
-    parallel_assert(
-        lambda: len(set(states)) == 1,
-    )
+    parallel_assert(lambda: len(set(states)) == 1)
 
 
 @pytest.mark.parallel(nprocs=2)
@@ -346,7 +327,7 @@ def test_send_and_recv(ensemble, mesh, W, blocking):
     root_ranks |= {ii + rank1*ensemble.comm.size for ii in range(ensemble.comm.size)}
     parallel_assert(
         lambda: error < 1e-12,
-        subset=root_ranks,
+        participating=COMM_WORLD.rank in root_ranks,
         msg=f"{error = :.5f}"  # noqa: E203, E251
     )