Using `unittest` and multiple values for `nprocs`

[ddundo]

Hi! I am trying to use this plugin with unittest-based tests. It works well if I only set a single value for nprocs, for example:

import unittest

import pytest
from firedrake.utility_meshes import UnitCubeMesh
from pyop2.mpi import COMM_WORLD

def count_elements(mesh):
    return COMM_WORLD.allreduce(mesh.coordinates.cell_set.size)

def count_vertices(mesh):
    return COMM_WORLD.allreduce(mesh.coordinates.node_set.size)


class TestParallel(unittest.TestCase):

    @pytest.mark.parallel(nprocs=3)
    def test_counting_3d(self):
        self.assertEqual(COMM_WORLD.size, 3)
        mesh = UnitCubeMesh(3, 3, 3)
        self.assertEqual(count_elements(mesh), 162)
        self.assertEqual(count_vertices(mesh), 64)

or @pytest.mark.parallel(nprocs=2) and so on. But if I give multiple values for nprocs, such as@pytest.mark.parallel(nprocs=[2, 3]), I get this error after calling the test with mpiexec -n 3 python3 -m pytest -v -m parallel[3] .:

Traceback

collected 1 item                                                                                                                                                                                                                                                                   
INTERNALERROR> Traceback (most recent call last):
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/_pytest/main.py", line 283, in wrap_session
INTERNALERROR>     session.exitstatus = doit(config, session) or 0
INTERNALERROR>                          ^^^^^^^^^^^^^^^^^^^^^
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/_pytest/main.py", line 336, in _main
INTERNALERROR>     config.hook.pytest_collection(session=session)
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/pluggy/_hooks.py", line 513, in __call__
INTERNALERROR>     return self._hookexec(self.name, self._hookimpls.copy(), kwargs, firstresult)
INTERNALERROR>            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/pluggy/_manager.py", line 120, in _hookexec
INTERNALERROR>     return self._inner_hookexec(hook_name, methods, kwargs, firstresult)
INTERNALERROR>            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/pluggy/_callers.py", line 139, in _multicall
INTERNALERROR>     raise exception.with_traceback(exception.__traceback__)
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/pluggy/_callers.py", line 122, in _multicall
INTERNALERROR>     teardown.throw(exception)  # type: ignore[union-attr]
INTERNALERROR>     ^^^^^^^^^^^^^^^^^^^^^^^^^
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/_pytest/logging.py", line 790, in pytest_collection
INTERNALERROR>     return (yield)
INTERNALERROR>             ^^^^^
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/pluggy/_callers.py", line 122, in _multicall
INTERNALERROR>     teardown.throw(exception)  # type: ignore[union-attr]
INTERNALERROR>     ^^^^^^^^^^^^^^^^^^^^^^^^^
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/_pytest/warnings.py", line 121, in pytest_collection
INTERNALERROR>     return (yield)
INTERNALERROR>             ^^^^^
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/pluggy/_callers.py", line 122, in _multicall
INTERNALERROR>     teardown.throw(exception)  # type: ignore[union-attr]
INTERNALERROR>     ^^^^^^^^^^^^^^^^^^^^^^^^^
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/_pytest/config/__init__.py", line 1417, in pytest_collection
INTERNALERROR>     return (yield)
INTERNALERROR>             ^^^^^
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/pluggy/_callers.py", line 103, in _multicall
INTERNALERROR>     res = hook_impl.function(*args)
INTERNALERROR>           ^^^^^^^^^^^^^^^^^^^^^^^^^
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/_pytest/main.py", line 347, in pytest_collection
INTERNALERROR>     session.perform_collect()
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/_pytest/main.py", line 812, in perform_collect
INTERNALERROR>     hook.pytest_collection_modifyitems(
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/pluggy/_hooks.py", line 513, in __call__
INTERNALERROR>     return self._hookexec(self.name, self._hookimpls.copy(), kwargs, firstresult)
INTERNALERROR>            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/pluggy/_manager.py", line 120, in _hookexec
INTERNALERROR>     return self._inner_hookexec(hook_name, methods, kwargs, firstresult)
INTERNALERROR>            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/pluggy/_callers.py", line 139, in _multicall
INTERNALERROR>     raise exception.with_traceback(exception.__traceback__)
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/pluggy/_callers.py", line 122, in _multicall
INTERNALERROR>     teardown.throw(exception)  # type: ignore[union-attr]
INTERNALERROR>     ^^^^^^^^^^^^^^^^^^^^^^^^^
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/_pytest/cacheprovider.py", line 443, in pytest_collection_modifyitems
INTERNALERROR>     res = yield
INTERNALERROR>           ^^^^^
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/pluggy/_callers.py", line 122, in _multicall
INTERNALERROR>     teardown.throw(exception)  # type: ignore[union-attr]
INTERNALERROR>     ^^^^^^^^^^^^^^^^^^^^^^^^^
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/_pytest/cacheprovider.py", line 373, in pytest_collection_modifyitems
INTERNALERROR>     res = yield
INTERNALERROR>           ^^^^^
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/pluggy/_callers.py", line 103, in _multicall
INTERNALERROR>     res = hook_impl.function(*args)
INTERNALERROR>           ^^^^^^^^^^^^^^^^^^^^^^^^^
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/pytest_mpi/plugin.py", line 100, in pytest_collection_modifyitems
INTERNALERROR>     nprocs = _extract_nprocs_for_single_test(item)
INTERNALERROR>              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
INTERNALERROR>   File "/home/firedrake/.local/lib/python3.12/site-packages/pytest_mpi/plugin.py", line 261, in _extract_nprocs_for_single_test
INTERNALERROR>     nprocs, = _parse_marker_nprocs(marker)
INTERNALERROR>     ^^^^^^^
INTERNALERROR> ValueError: too many values to unpack (expected 1)

So this seems to skip mpi-pytest's pytest_generate_tests, which then causes an issue since the test isn't split per number of processes. Do you please have an idea how this could be fixed?

For now I have the following workaround. I mark tests with @pytest.mark.parallel_dynamic which I then replace with the regular @pytest.mark.parallel marker in conftest.py:

import pytest
from pyop2.mpi import COMM_WORLD


def pytest_configure(config):
    config.addinivalue_line(
        "markers", 
        "parallel_dynamic: mark test to run with the current MPI size"
    )

def pytest_collection_modifyitems(config, items):
    """
    Replace parallel_dynamic markers with pytest-mpi's parallel marker such that
    the test is run with the current MPI size.
    """
    rank_size = COMM_WORLD.Get_size()
    
    for item in items:
        # Check if a parallel_dynamic marker is set
        markers = [marker for marker in item.own_markers 
                  if marker.name == 'parallel_dynamic']
        
        if markers:
            # Remove parallel_dynamic markers
            for marker in markers:
                item.own_markers.remove(marker)
            
            # Add regular parallel marker with current process count
            item.add_marker(pytest.mark.parallel(nprocs=rank_size))

So the test gets the @pytest.mark.parallel(nprocs=2) marker if I call mpiexec -n 2 python3 -m pytest -v -m parallel[2] . and similarly with nprocs=3 and so on. This works well I think. Do you foresee some issues with this approach?

[JHopeCollins]

I'm not sure about why it isn't working properly with unittest, but I've been thinking about some updates to mpi-pytest literally this morning that would allow an equivalent of your parallel_dynamic. I'm going to open an issue later and I'll link it here.

[connorjward]

I'm not sure how I feel about parallel_dynamic. It feels better to marks tests with a specific level of parallelism since it is explicit and won't lead to any unexpected effects.

Until now I had not considered marking tests inside of classes so I am not terribly surprised that it fails to do the right thing.

[ddundo]

Thanks both. The issue doesn't seem to be caused by the test being inside a class. I.e., this works fine

class TestParallel:
    @pytest.mark.parallel(nprocs=[2, 3])
    def test_counting_3d(self):
        mesh = UnitCubeMesh(3, 3, 3)
        assert count_elements(mesh) == 162
        assert count_vertices(mesh) == 64

I'll keep an eye on how @JHopeCollins's updates shape up.

[JHopeCollins]

See #16