fix: add option to disable MPI (#68)

CalcJob plugins can specify default values for whether codes should use
MPI or not: using the `metadata.options.withmpi` variable.

aiida-testing can run calculations with a single MPI process if an MPI
library is installed. However, the MPI library is then also required in
order to fetch the results from the test database, because the mock
executable will be run with `mpirun -np 1 ...`.

Here we add an option `--mock-disable-mpi` that allows users to disable
the use of MPI. If supported by the AiiDA version (>=2.1.0), this option is
enabled automatically by `--mock-fail-on-missing`.
The documentation now recommends using `--mock-fail-on-missing` on CI
by default.

aiida_testing/mock_code/_fixtures.py

@@ -10,11 +10,13 @@
 import warnings
 import collections
 import os
+from pkg_resources import parse_version
 
 import click
 import pytest
 
 from aiida.orm import Code
+from aiida import __version__ as aiida_version
 
 from ._env_keys import MockVariables
 from ._hasher import InputHasher
@@ -25,6 +27,7 @@
     "testing_config_action",
     "mock_regenerate_test_data",
     "mock_fail_on_missing",
+    "mock_disable_mpi",
     "testing_config",
     "mock_code_factory",
 )
@@ -51,6 +54,12 @@ def pytest_addoption(parser):
         default=False,
         help="Fail if cached data is not found, rather than regenerating it.",
     )
+    parser.addoption(
+        "--mock-disable-mpi",
+        action="store_true",
+        default=False,
+        help="Run all calculations with `metadata.options.usempi=False`.",
+    )
 
 
 @pytest.fixture(scope='session')
@@ -71,6 +80,15 @@ def mock_fail_on_missing(request):
     return request.config.getoption("--mock-fail-on-missing")
 
 
+@pytest.fixture(scope='function')
+def mock_disable_mpi(request):
+    """Enforce `withmpi=False` based on `--mock-disable-mpi` cli option.
+
+    This is achieved by monkey-patching the `CalcJob.get_option()` method.
+    """
+    return request.config.getoption("--mock-disable-mpi")
+
+
 @pytest.fixture(scope='session')
 def testing_config(testing_config_action):  # pylint: disable=redefined-outer-name
     """Get content of .aiida-testing-config.yml
@@ -89,18 +107,27 @@ def testing_config(testing_config_action):  # pylint: disable=redefined-outer-na
         config.to_file()
 
 
+def _forget_mpi_decorator(func):
+    """Modify :py:meth:`aiida.orm.Code.get_prepend_cmdline_params` to discard MPI parameters."""
+
+    def _get_prepend_cmdline_params(self, mpi_args=None, extra_mpirun_params=None):  # pylint: disable=unused-argument
+        return func(self)
+
+    return _get_prepend_cmdline_params
+
+
 @pytest.fixture(scope='function')
 def mock_code_factory(
     aiida_localhost, testing_config, testing_config_action, mock_regenerate_test_data,
-    mock_fail_on_missing, request: pytest.FixtureRequest, tmp_path: pathlib.Path
-):  # pylint: disable=too-many-arguments,redefined-outer-name
+    mock_fail_on_missing, mock_disable_mpi, monkeypatch, request: pytest.FixtureRequest,
+    tmp_path: pathlib.Path
+):  # pylint: disable=too-many-arguments,redefined-outer-name,unused-argument,too-many-statements
     """
     Fixture to create a mock AiiDA Code.
 
     testing_config_action :
         Read config file if present ('read'), require config file ('require') or generate new config file ('generate').
 
-
     """
     log_file = tmp_path.joinpath("_aiida_mock_code.log")
     log_file.touch()
@@ -117,6 +144,7 @@ def _get_mock_code(
         _config_action: str = testing_config_action,
         _regenerate_test_data: bool = mock_regenerate_test_data,
         _fail_on_missing: bool = mock_fail_on_missing,
+        _disable_mpi: bool = mock_disable_mpi,
     ):  # pylint: disable=too-many-arguments,too-many-branches,too-many-locals
         """
         Creates a mock AiiDA code. If the same inputs have been run previously,
@@ -237,6 +265,24 @@ def _get_mock_code(
         code.set_prepend_text(variables.to_env())
 
         code.store()
+
+        # Monkeypatch MPI behavior of code class, if requested either directly via `--mock-disable-mpi` or
+        # indirectly via `--mock-fail-on-missing` (no need to use MPI in this case)
+        if _disable_mpi or _fail_on_missing:
+            is_mpi_disable_supported = parse_version(aiida_version) >= parse_version('2.1.0')
+
+            if not is_mpi_disable_supported:
+                if _disable_mpi:
+                    raise ValueError(
+                        "Upgrade to AiiDA >= 2.1.0 in order to use `--mock-disable-mpi`"
+                    )
+                # if only _fail_on_missing, we silently do not disable MPI
+            else:
+                monkeypatch.setattr(
+                    code.__class__, 'get_prepend_cmdline_params',
+                    _forget_mpi_decorator(code.__class__.get_prepend_cmdline_params)
+                )
+
         return code
 
     yield _get_mock_code

docs/source/user_guide/mock_code.rst

@@ -45,6 +45,7 @@ Second, we need to tell the mock executable where to find the *actual* ``diff``
       # code-label: absolute path
       diff: /usr/bin/diff
 
+
 .. note::
    Why yet another configuration file?
 
@@ -77,10 +78,13 @@ The next time, it will recognise the inputs and directly use the outputs cached
     It does *not* rely on the hashing mechanism of AiiDA.
 
 
-Don't forget to add your data directory to your test data in order to make them available in CI and to other users of your plugin!
+Running continuous integration (CI) tests on your repository:
+
+ - Don't forget to commit changes to your data directory to make the cache available on CI
+ - Run tests on CI with ``pytest --mock-fail-on-missing`` to force a test failure when it fails when the committed cache is incomplete
 
-Since the ``.aiida-testing-config.yml`` is usually specific to your machine, it usually better not to commit it.
-Tests will run fine without it, and if other developers need to change test inputs, they can easily regenerate a template for it using ``pytest --testing-config-action=generate``.
+Since the ``.aiida-testing-config.yml`` file is usually specific to your machine, there is no need to commit it.
+As long as the test cache is complete, tests will run fine without it, and if other developers need to change test inputs, they can easily regenerate a template for it using ``pytest --testing-config-action=generate``.
 
 For further documentation on the pytest commandline options added by mock code, see:
 
@@ -96,6 +100,9 @@ For further documentation on the pytest commandline options added by mock code,
       --mock-regenerate-test-data
                             Regenerate test data.
 
+      --mock-fail-on-missing
+                            Fail if cached data is not found, rather than regenerating it.
+      --mock-disable-mpi    Run all calculations with `metadata.options.usempi=False`.
 
 Limitations
 -----------

pyproject.toml

@@ -51,11 +51,13 @@ testing = [
     "pgtest~=1.3.1",
     "aiida-diff",
     "pytest-datadir",
+    "pytest-mock",
 ]
 pre_commit = [
     "pre-commit",
     "pylint~=2.12.2",
     "mypy==0.930",
+    "types-setuptools==65.7.0.3",
     "types-PyYAML",
 ]
 dev = [

tests/mock_code/test_diff.py

@@ -5,11 +5,14 @@
 
 import shutil
 import os
+import json
 import tempfile
 from pathlib import Path
+from pkg_resources import parse_version
 
 import pytest
 
+from aiida import __version__ as aiida_version
 from aiida.engine import run_get_node
 from aiida.plugins import CalculationFactory
 
@@ -208,3 +211,34 @@ def test_regenerate_test_data_executable(mock_code_factory, generate_diff_inputs
     # check that ignore_paths works
     assert not (datadir / '_aiidasubmit.sh').is_file()
     assert (datadir / 'file1.txt').is_file()
+
+
+@pytest.mark.skipif(
+    parse_version(aiida_version) < parse_version('2.1.0'), reason='requires AiiDA v2.1.0+'
+)
+def test_disable_mpi(mock_code_factory, generate_diff_inputs):  # pylint: disable=unused-argument
+    """
+    Check that disabling MPI is respected.
+
+    Let a CalcJob explicitly request `withmpi=True`, and check it is still run without MPI.
+    """
+    mock_code = mock_code_factory(
+        label='diff',
+        data_dir_abspath=TEST_DATA_DIR,
+        entry_point=CALC_ENTRY_POINT,
+        ignore_paths=('_aiidasubmit.sh', 'file*txt'),
+        _disable_mpi=True,
+    )
+
+    inputs = generate_diff_inputs()
+    inputs['metadata']['options']['withmpi'] = True
+
+    res, node = run_get_node(CalculationFactory(CALC_ENTRY_POINT), code=mock_code, **inputs)
+    assert node.exit_status == 0, f"diff calculation failed with exit status {node.exit_status}"
+    assert node.is_finished_ok
+    check_diff_output(res)
+
+    # check that the submit script does not contain mpirun
+    job_tmpl = json.loads(node.base.repository.get_object_content('.aiida/job_tmpl.json'))
+    assert not job_tmpl['codes_info'][0]['prepend_cmdline_params']
+    assert 'mpirun' not in node.base.repository.get_object_content('_aiidasubmit.sh')