Parallel test improvements (#3865)

* Refactor test_demos, avoid mpiexec invocation
* Fix multigrid test, __main__ is not valid when importing

demos/multigrid/geometric_multigrid.py.rst

@@ -191,7 +191,7 @@ bilinear form to the solver ourselves: ::
       "fieldsplit_0_pc_type": "mg",
       "fieldsplit_1_ksp_type": "preonly",
       "fieldsplit_1_pc_type": "python",
-      "fieldsplit_1_pc_python_type": "__main__.Mass",
+      "fieldsplit_1_pc_python_type": "geometric_multigrid.Mass",
       "fieldsplit_1_aux_pc_type": "bjacobi",
       "fieldsplit_1_aux_sub_pc_type": "icc",
   }
@@ -227,7 +227,7 @@ approximations.
         "mg_coarse_fieldsplit_0_pc_type": "lu",
         "mg_coarse_fieldsplit_1_ksp_type": "preonly",
         "mg_coarse_fieldsplit_1_pc_type": "python",
-        "mg_coarse_fieldsplit_1_pc_python_type": "__main__.Mass",
+        "mg_coarse_fieldsplit_1_pc_python_type": "geometric_multigrid.Mass",
         "mg_coarse_fieldsplit_1_aux_pc_type": "cholesky",
         "mg_levels_ksp_type": "richardson",
         "mg_levels_ksp_max_it": 1,
@@ -245,7 +245,7 @@ approximations.
         "mg_levels_fieldsplit_1_ksp_richardson_self_scale": None,
         "mg_levels_fieldsplit_1_ksp_max_it": 3,
         "mg_levels_fieldsplit_1_pc_type": "python",
-        "mg_levels_fieldsplit_1_pc_python_type": "__main__.Mass",
+        "mg_levels_fieldsplit_1_pc_python_type": "geometric_multigrid.Mass",
         "mg_levels_fieldsplit_1_aux_pc_type": "bjacobi",
         "mg_levels_fieldsplit_1_aux_sub_pc_type": "icc",
   }

tests/demos/test_demos_run.py

@@ -1,60 +1,120 @@
-import pytest
-from os.path import abspath, basename, dirname, join, splitext
+import glob
+import importlib
 import os
 import subprocess
-import glob
 import sys
+from collections import namedtuple
+from pathlib import Path
+from os.path import abspath, basename, dirname, join, splitext
+
+import pyadjoint
+import pytest
+
 from firedrake.petsc import get_external_packages
 
 
-cwd = abspath(dirname(__file__))
-demo_dir = join(cwd, "..", "..", "demos")
-VTK_DEMOS = [
-    "benney_luke.py",
-    "burgers.py",
-    "camassaholm.py",
-    "geometric_multigrid.py",
-    "helmholtz.py",
-    "higher_order_mass_lumping.py",
-    "linear_fluid_structure_interaction.py",
-    "linear_wave_equation.py",
-    "ma-demo.py",
-    "navier_stokes.py",
-    "netgen_mesh.py",
-    "poisson_mixed.py",
-    "qg_1layer_wave.py",
-    "qgbasinmodes.py",
-    "qg_winddrivengyre.py",
-    "rayleigh-benard.py",
-    "stokes.py",
-    "test_extrusion_lsw.py",
+Demo = namedtuple("Demo", ["loc", "requirements"])
+
+
+CWD = abspath(dirname(__file__))
+DEMO_DIR = join(CWD, "..", "..", "demos")
+
+SERIAL_DEMOS = [
+    Demo(("benney_luke", "benney_luke"), ["vtk"]),
+    Demo(("burgers", "burgers"), ["vtk"]),
+    Demo(("camassa-holm", "camassaholm"), ["vtk"]),
+    Demo(("DG_advection", "DG_advection"), ["matplotlib"]),
+    Demo(("eigenvalues_QG_basinmodes", "qgbasinmodes"), ["matplotlib", "slepc", "vtk"]),
+    Demo(("extruded_continuity", "extruded_continuity"), []),
+    Demo(("helmholtz", "helmholtz"), ["vtk"]),
+    Demo(("higher_order_mass_lumping", "higher_order_mass_lumping"), ["vtk"]),
+    Demo(("immersed_fem", "immersed_fem"), []),
+    Demo(("linear_fluid_structure_interaction", "linear_fluid_structure_interaction"), ["vtk"]),
+    Demo(("linear-wave-equation", "linear_wave_equation"), ["vtk"]),
+    Demo(("ma-demo", "ma-demo"), ["vtk"]),
+    Demo(("matrix_free", "navier_stokes"), ["mumps", "vtk"]),
+    Demo(("matrix_free", "poisson"), []),
+    Demo(("matrix_free", "rayleigh-benard"), ["hypre", "mumps", "vtk"]),
+    Demo(("matrix_free", "stokes"), ["hypre", "mumps", "vtk"]),
+    Demo(("multigrid", "geometric_multigrid"), ["vtk"]),
+    Demo(("netgen", "netgen_mesh"), ["mumps", "ngsPETSc", "netgen", "slepc", "vtk"]),
+    Demo(("nonlinear_QG_winddrivengyre", "qg_winddrivengyre"), ["vtk"]),
+    Demo(("parallel-printing", "parprint"), []),
+    Demo(("poisson", "poisson_mixed"), ["vtk"]),
+    Demo(("quasigeostrophy_1layer", "qg_1layer_wave"), ["hypre", "vtk"]),
+    Demo(("saddle_point_pc", "saddle_point_systems"), ["hypre", "mumps"]),
 ]
-
-parallel_demos = [
-    "full_waveform_inversion.py",
+PARALLEL_DEMOS = [
+    Demo(("full_waveform_inversion", "full_waveform_inversion"), ["adjoint"]),
 ]
 
 
-# Discover the demo files by globbing the demo directory
-@pytest.fixture(params=glob.glob("%s/*/*.py.rst" % demo_dir),
-                ids=lambda x: basename(x))
-def rst_file(request):
-    return abspath(request.param)
-
-
 @pytest.fixture
 def env():
     env = os.environ.copy()
     env["MPLBACKEND"] = "pdf"
     return env
 
 
-@pytest.fixture
-def py_file(rst_file, tmpdir, monkeypatch):
+def test_no_missing_demos():
+    all_demo_locs = {
+        demo.loc
+        for demos in [SERIAL_DEMOS, PARALLEL_DEMOS]
+        for demo in demos
+    }
+    for rst_file in glob.glob(f"{DEMO_DIR}/*/*.py.rst"):
+        rst_path = Path(rst_file)
+        demo_dir = rst_path.parent.name
+        demo_name, _, _ = rst_path.name.split(".")
+        demo_loc = (demo_dir, demo_name)
+        assert demo_loc in all_demo_locs
+        all_demo_locs.remove(demo_loc)
+    assert not all_demo_locs, "Unrecognised demos listed"
+
+
+def _maybe_skip_demo(demo):
+    # Add pytest skips for missing imports or packages
+    if "mumps" in demo.requirements and "mumps" not in get_external_packages():
+        pytest.skip("MUMPS not installed with PETSc")
+
+    if "hypre" in demo.requirements and "hypre" not in get_external_packages():
+        pytest.skip("hypre not installed with PETSc")
+
+    if "slepc" in demo.requirements:
+        try:
+            # Do not use `pytest.importorskip` to check for slepc4py:
+            # It isn't sufficient to actually detect whether slepc4py
+            # is installed. Both petsc4py and slepc4py require
+            # `from xy4py import Xy`
+            # to actually load the library.
+            from slepc4py import SLEPc  # noqa: F401
+        except ImportError:
+            pytest.skip("SLEPc unavailable")
+
+    if "matplotlib" in demo.requirements:
+        pytest.importorskip("matplotlib", reason="Matplotlib unavailable")
+
+    if "netgen" in demo.requirements:
+        pytest.importorskip("netgen", reason="Netgen unavailable")
+
+    if "ngsPETSc" in demo.requirements:
+        pytest.importorskip("ngsPETSc", reason="ngsPETSc unavailable")
+
+    if "vtk" in demo.requirements:
+        try:
+            import vtkmodules.vtkCommonDataModel  # noqa: F401
+        except ImportError:
+            pytest.skip("VTK unavailable")
+
+
+def _prepare_demo(demo, monkeypatch, tmpdir):
     # Change to the temporary directory (monkeypatch ensures that this
     # is undone when the fixture usage disappears)
     monkeypatch.chdir(tmpdir)
 
+    demo_dir, demo_name = demo.loc
+    rst_file = f"{DEMO_DIR}/{demo_dir}/{demo_name}.py.rst"
+
     # Check if we need to generate any meshes
     geos = glob.glob("%s/*.geo" % dirname(rst_file))
     for geo in geos:
@@ -70,67 +130,38 @@ def py_file(rst_file, tmpdir, monkeypatch):
 
     # Get the name of the python file that pylit will make
     name = splitext(basename(rst_file))[0]
-    output = str(tmpdir.join(name))
+    py_file = str(tmpdir.join(name))
     # Convert rst demo to runnable python file
-    subprocess.check_call(["pylit", rst_file, output])
-    return output
+    subprocess.check_call(["pylit", rst_file, py_file])
+    return Path(py_file)
 
 
-@pytest.mark.skipcomplex  # Will need to add a seperate case for a complex demo.
-def test_demo_runs(py_file, env):
-    # Add pytest skips for missing imports or packages
-    if basename(py_file) in ("stokes.py", "rayleigh-benard.py", "saddle_point_systems.py", "navier_stokes.py", "netgen_mesh.py"):
-        if "mumps" not in get_external_packages():
-            pytest.skip("MUMPS not installed with PETSc")
+def _exec_file(py_file):
+    # To execute a file we import it. We therefore need to modify sys.path so the
+    # tempdir can be found.
+    sys.path.insert(0, str(py_file.parent))
+    importlib.import_module(py_file.with_suffix("").name)
+    sys.path.pop(0)  # cleanup
 
-    if basename(py_file) in ("stokes.py", "rayleigh-benard.py", "saddle_point_systems.py", "qg_1layer_wave.py"):
-        if "hypre" not in get_external_packages():
-            pytest.skip("hypre not installed with PETSc")
 
-    if basename(py_file) == "qgbasinmodes.py":
-        try:
-            # Do not use `pytest.importorskip` to check for slepc4py:
-            # It isn't sufficient to actually detect whether slepc4py
-            # is installed. Both petsc4py and slepc4py require
-            # `from xy4py import Xy`
-            # to actually load the library.
-            from slepc4py import SLEPc  # noqa: F401
-        except ImportError:
-            pytest.skip(reason="SLEPc unavailable, skipping qgbasinmodes.py")
-
-    if basename(py_file) in ("DG_advection.py", "qgbasinmodes.py"):
-        pytest.importorskip(
-            "matplotlib",
-            reason=f"Matplotlib unavailable, skipping {basename(py_file)}"
-        )
-
-    if basename(py_file) == "netgen_mesh.py":
-        pytest.importorskip(
-            "netgen",
-            reason="Netgen unavailable, skipping Netgen test."
-        )
-        pytest.importorskip(
-            "ngsPETSc",
-            reason="ngsPETSc unavailable, skipping Netgen test."
-        )
-        try:
-            from slepc4py import SLEPc  # noqa: F401, F811
-        except ImportError:
-            pytest.skip(reason="SLEPc unavailable, skipping netgen_mesh.py")
+@pytest.mark.skipcomplex
+@pytest.mark.parametrize("demo", SERIAL_DEMOS, ids=["/".join(d.loc) for d in SERIAL_DEMOS])
+def test_serial_demo(demo, env, monkeypatch, tmpdir):
+    _maybe_skip_demo(demo)
+    py_file = _prepare_demo(demo, monkeypatch, tmpdir)
+    _exec_file(py_file)
 
-    if basename(py_file) in VTK_DEMOS:
-        try:
-            import vtkmodules.vtkCommonDataModel  # noqa: F401
-        except ImportError:
-            pytest.skip(reason=f"VTK unavailable, skipping {basename(py_file)}")
-    if basename(py_file) in parallel_demos:
-        if basename(py_file) == "full_waveform_inversion.py":
-            processes = 2
-        else:
-            raise NotImplementedError("You need to specify the number of processes for this test")
-
-        executable = ["mpiexec", "-n", str(processes), sys.executable, py_file]
-    else:
-        executable = [sys.executable, py_file]
-
-    subprocess.check_call(executable, env=env)
+    if "adjoint" in demo.requirements:
+        pyadjoint.get_working_tape().clear_tape()
+
+
+@pytest.mark.parallel(2)
+@pytest.mark.skipcomplex
+@pytest.mark.parametrize("demo", PARALLEL_DEMOS, ids=["/".join(d.loc) for d in PARALLEL_DEMOS])
+def test_parallel_demo(demo, env, monkeypatch, tmpdir):
+    _maybe_skip_demo(demo)
+    py_file = _prepare_demo(demo, monkeypatch, tmpdir)
+    _exec_file(py_file)
+
+    if "adjoint" in demo.requirements:
+        pyadjoint.get_working_tape().clear_tape()

tests/regression/test_ensembleparallelism.py

@@ -205,13 +205,13 @@ def test_ensemble_reduce(ensemble, mesh, W, urank, urank_sum, root, blocking):
     parallel_assert(
         lambda: error < 1e-12,
         subset=root_ranks,
-        msg=f"{error = :.5f}"
+        msg=f"{error=:.5f}"
     )
     error = errornorm(Function(W).assign(10), u_reduce)
     parallel_assert(
         lambda: error < 1e-12,
         subset={range(COMM_WORLD.size)} - root_ranks,
-        msg=f"{error = :.5f}"
+        msg=f"{error=:.5f}"
     )
 
     # check that u_reduce dat vector is still synchronised
@@ -347,7 +347,7 @@ def test_send_and_recv(ensemble, mesh, W, blocking):
     parallel_assert(
         lambda: error < 1e-12,
         subset=root_ranks,
-        msg=f"{error = :.5f}"
+        msg=f"{error=:.5f}"
     )
 
 

tests/regression/test_vertex_based_limiter.py

@@ -1,8 +1,6 @@
 import pytest
 from firedrake import *
 import numpy as np
-import subprocess
-import sys
 
 
 @pytest.fixture(params=["periodic-interval",
@@ -122,44 +120,23 @@ def test_step_function_loop(mesh, iterations=100):
     assert np.min(u.dat.data_ro) >= 0.0, "Failed by exceeding min values"
 
 
+@pytest.mark.parallel
 @pytest.mark.skipcomplex
-def test_parallel_limiting(tmpdir):
-    import pickle
-    mesh = RectangleMesh(10, 4, 5000., 1000.)
+def test_parallel_limiting():
+    serial_result = _apply_limiter_with_comm(COMM_SELF)
+    parallel_result = _apply_limiter_with_comm(COMM_WORLD)
+    assert np.allclose(serial_result, parallel_result)
+
+
+def _apply_limiter_with_comm(comm):
+    mesh = RectangleMesh(10, 4, 5000., 1000., comm=comm)
     V = space(mesh)
     f = Function(V)
     x, *_ = SpatialCoordinate(mesh)
     f.project(sin(2*pi*x/3000.))
     limiter = VertexBasedLimiter(V)
     limiter.apply(f)
 
-    expect = np.asarray([norm(f),
-                         norm(limiter.centroids),
-                         norm(limiter.min_field),
-                         norm(limiter.max_field)])
-
-    tmpfile = tmpdir.join("a")
-    code = """
-import pickle
-from firedrake import *
-mesh = RectangleMesh(10, 4, 5000., 1000.)
-element = BrokenElement(mesh.coordinates.function_space().ufl_element().sub_elements[0])
-V = FunctionSpace(mesh, element)
-f = Function(V)
-x, *_ = SpatialCoordinate(mesh)
-f.project(sin(2*pi*x/3000.))
-limiter = VertexBasedLimiter(V)
-limiter.apply(f)
-
-fnorm = norm(f)
-centroid_norm = norm(limiter.centroids)
-min_norm = norm(limiter.min_field)
-max_norm = norm(limiter.max_field)
-if mesh.comm.rank == 0:
-    with open("{file}", "wb") as f:
-        pickle.dump([fnorm, centroid_norm, min_norm, max_norm], f)
-""".format(file=tmpfile)
-    subprocess.check_call(["mpiexec", "-n", "3", sys.executable, "-c", code])
-    with tmpfile.open("rb") as f:
-        actual = np.asarray(pickle.load(f))
-    assert np.allclose(expect, actual)
+    return np.asarray([
+        norm(f), norm(limiter.centroids), norm(limiter.min_field), norm(limiter.max_field)
+    ])