PR #417: from firedrakeproject/JDBetteridge/parallel_tests

Set up parallel test suite

.github/workflows/build.yml

@@ -26,23 +26,37 @@ jobs:
         run: |
           cd ..
           rm -rf build
-      - name: Setup python
-        uses: actions/setup-python@v4
-        with:
-          python-version: 3.8
       - name: Install Gusto
         run: |
           . /home/firedrake/firedrake/bin/activate
           python -m pip install -r requirements.txt
           python -m pip install -e .
           python -m pip install  \
-            pytest-cov pytest-timeout pytest-xdist
+            pytest-timeout pytest-xdist
       - name: Gusto tests
         run: |
           . /home/firedrake/firedrake/bin/activate
-          which firedrake-clean
+          firedrake-clean
+          export GUSTO_PARALLEL_LOG=FILE
           python -m pytest \
             -n 12 --dist worksteal \
             --durations=100 \
-            --cov gusto \
+            --timeout=3600 \
+            --timeout-method=thread \
+            -o faulthandler_timeout=3660 \
             -v unit-tests integration-tests examples
+        timeout-minutes: 120
+      - name: Prepare logs
+        if: always()
+        run: |
+          mkdir logs
+          cd /tmp/pytest-of-firedrake/pytest-0/
+          find . -name "*.log" -exec cp --parents {} /__w/gusto/gusto/logs/ \;
+      - name: Upload artifact
+        if: always()
+        uses: actions/upload-pages-artifact@v1
+        with:
+          name: log-files
+          path: /__w/gusto/gusto/logs
+          retention-days: 5
+

examples/shallow_water/shallow_water_1d.py

@@ -1,4 +1,5 @@
 import numpy as np
+import sys
 
 from firedrake import *
 from gusto import *
@@ -9,6 +10,10 @@
 delta = L/n
 mesh = PeriodicIntervalMesh(128, L)
 dt = 0.0001
+if '--running-tests' in sys.argv:
+    T = 0.0005
+else:
+    T = 1
 
 domain = Domain(mesh, dt, 'CG', 1)
 
@@ -61,4 +66,4 @@
 
 D += parameters.H
 
-stepper.run(0, 1)
+stepper.run(0, T)

examples/test_examples_run.py

@@ -3,6 +3,7 @@
 import subprocess
 import glob
 import sys
+import os
 
 
 examples_dir = abspath(dirname(__file__))
@@ -19,4 +20,19 @@ def test_example_runs(example_file, tmpdir, monkeypatch):
     # This ensures that the test writes output in a temporary
     # directory, rather than where pytest was run from.
     monkeypatch.chdir(tmpdir)
-    subprocess.check_call([sys.executable, example_file, "--running-tests"])
+    subprocess.run(
+        [sys.executable, example_file, "--running-tests"],
+        check=True,
+        env=os.environ | {"PYOP2_CFLAGS": "-O0"}
+    )
+
+
+def test_example_runs_parallel(example_file, tmpdir, monkeypatch):
+    # This ensures that the test writes output in a temporary
+    # directory, rather than where pytest was run from.
+    monkeypatch.chdir(tmpdir)
+    subprocess.run(
+        ["mpiexec", "-n", "4", sys.executable, example_file, "--running-tests"],
+        check=True,
+        env=os.environ | {"PYOP2_CFLAGS": "-O0"}
+    )

gusto/solvers/linear_solvers.py

@@ -17,7 +17,10 @@
 from pyop2.profiling import timed_function, timed_region
 
 from gusto.equations.active_tracers import TracerVariableType
-from gusto.core.logging import logger, DEBUG, logging_ksp_monitor_true_residual
+from gusto.core.logging import (
+    logger, DEBUG, logging_ksp_monitor_true_residual,
+    attach_custom_monitor
+)
 from gusto.core.labels import linearisation, time_derivative, hydrostatic
 from gusto.equations import thermodynamics
 from gusto.recovery.recovery_kernels import AverageWeightings, AverageKernel
@@ -56,8 +59,8 @@ def __init__(self, equations, alpha=0.5, solver_parameters=None,
                 solver_parameters = p
             self.solver_parameters = solver_parameters
 
-        # ~ if logger.isEnabledFor(DEBUG):
-            # ~ self.solver_parameters["ksp_monitor_true_residual"] = None
+        if logger.isEnabledFor(DEBUG):
+            self.solver_parameters["ksp_monitor_true_residual"] = None
 
         # setup the solver
         self._setup_solver()
@@ -353,7 +356,6 @@ def L_tr(f):
         # Log residuals on hybridized solver
         self.log_ksp_residuals(self.hybridized_solver.snes.ksp)
         # Log residuals on the trace system too
-        from gusto.core.logging import attach_custom_monitor
         python_context = self.hybridized_solver.snes.ksp.pc.getPythonContext()
         attach_custom_monitor(python_context, logging_ksp_monitor_true_residual)
 

gusto/timestepping/semi_implicit_quasi_newton.py

@@ -262,12 +262,12 @@ def timestep(self):
         x_after_slow(self.field_name).assign(xn(self.field_name))
         if len(self.slow_physics_schemes) > 0:
             with timed_stage("Slow physics"):
-                logger.info('SIQN: Slow physics')
+                logger.info('Semi-implicit Quasi Newton: Slow physics')
                 for _, scheme in self.slow_physics_schemes:
                     scheme.apply(x_after_slow(scheme.field_name), x_after_slow(scheme.field_name))
 
         with timed_stage("Apply forcing terms"):
-            logger.info('SIQN: Explicit forcing')
+            logger.info('Semi-implicit Quasi Newton: Explicit forcing')
             # Put explicit forcing into xstar
             self.forcing.apply(x_after_slow, xn, xstar(self.field_name), "explicit")
 
@@ -281,14 +281,14 @@ def timestep(self):
                 self.io.log_courant(self.fields, 'transporting_velocity',
                                     message=f'transporting velocity, outer iteration {outer}')
                 for name, scheme in self.active_transport:
-                    logger.info(f'SIQN: Transport {outer}: {name}')
+                    logger.info(f'Semi-implicit Quasi Newton: Transport {outer}: {name}')
                     # transports a field from xstar and puts result in xp
                     scheme.apply(xp(name), xstar(name))
 
             x_after_fast(self.field_name).assign(xp(self.field_name))
             if len(self.fast_physics_schemes) > 0:
                 with timed_stage("Fast physics"):
-                    logger.info(f'SIQN: Fast physics {outer}')
+                    logger.info(f'Semi-implicit Quasi Newton: Fast physics {outer}')
                     for _, scheme in self.fast_physics_schemes:
                         scheme.apply(x_after_fast(scheme.field_name), x_after_fast(scheme.field_name))
 
@@ -300,14 +300,14 @@ def timestep(self):
                 # TODO: this is where to update the reference state
 
                 with timed_stage("Apply forcing terms"):
-                    logger.info(f'SIQN: Implicit forcing {(outer, inner)}')
+                    logger.info(f'Semi-implicit Quasi Newton: Implicit forcing {(outer, inner)}')
                     self.forcing.apply(xp, xnp1, xrhs, "implicit")
 
                 xrhs -= xnp1(self.field_name)
                 xrhs += xrhs_phys
 
                 with timed_stage("Implicit solve"):
-                    logger.info(f'SIQN: Mixed solve {(outer, inner)}')
+                    logger.info(f'Semi-implicit Quasi Newton: Mixed solve {(outer, inner)}')
                     self.linear_solver.solve(xrhs, dy)  # solves linear system and places result in dy
 
                 xnp1X = xnp1(self.field_name)
@@ -320,17 +320,21 @@ def timestep(self):
 
         for name, scheme in self.auxiliary_schemes:
             # transports a field from xn and puts result in xnp1
+            logger.debug(f"Semi-implicit Quasi-Newton auxiliary scheme for {name}")
             scheme.apply(xnp1(name), xn(name))
 
         with timed_stage("Diffusion"):
             for name, scheme in self.diffusion_schemes:
+                logger.debug(f"Semi-implicit Quasi-Newton diffusing {name}")
                 scheme.apply(xnp1(name), xnp1(name))
 
         if len(self.final_physics_schemes) > 0:
             with timed_stage("Final Physics"):
                 for _, scheme in self.final_physics_schemes:
                     scheme.apply(xnp1(scheme.field_name), xnp1(scheme.field_name))
 
+        logger.debug("Leaving Semi-implicit Quasi-Newton timestep method")
+
     def run(self, t, tmax, pick_up=False):
         """
         Runs the model for the specified time, from t to tmax.

gusto/timestepping/timestepper.py

@@ -171,6 +171,7 @@ def run(self, t, tmax, pick_up=False):
 
         # Set up dump, which may also include an initial dump
         with timed_stage("Dump output"):
+            logger.debug('Dumping output to disk')
             self.io.setup_dump(self.fields, t, pick_up)
 
         self.t.assign(t)
@@ -197,7 +198,7 @@ def run(self, t, tmax, pick_up=False):
         if self.io.output.checkpoint and self.io.output.checkpoint_method == 'dumbcheckpoint':
             self.io.chkpt.close()
 
-        logger.info(f'TIMELOOP complete. t={float(self.t)}, tmax={tmax}')
+        logger.info(f'TIMELOOP complete. t={float(self.t):.5f}, {tmax=:.5f}')
 
     def set_reference_profiles(self, reference_profiles):
         """