diff --git a/doc/conf.py b/doc/conf.py index 70a55f0c8b0af0ad56a6b73303ef6393b33fc99c..133844183bdd4eb38fffa5ae56cc79fbccd510b0 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -29,8 +29,11 @@ intersphinx_mapping = { "https://numpy.org/doc/stable": None, "https://documen.tician.de/pymbolic/": None, "https://documen.tician.de/loopy/": None, + "https://docs.pytest.org/en/stable/": None, } nitpick_ignore_regex = [ ["py:class", r"typing_extensions\.(.+)"], ] + +nitpicky = True diff --git a/doc/index.rst b/doc/index.rst index 25f26c6244c9c8c679e05669e463079aa0b0c373..1a1e9c444a55b9f7aed9aeb27b0bf8bfc725daea 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -11,6 +11,7 @@ Welcome to pytools's documentation! graph tag codegen + mpi misc 🚀 Github <https://github.com/inducer/pytools> 💾 Download Releases <https://pypi.python.org/pypi/pytools> diff --git a/doc/mpi.rst b/doc/mpi.rst new file mode 100644 index 0000000000000000000000000000000000000000..e7bcae56843a03fa83c124578322ab89da923610 --- /dev/null +++ b/doc/mpi.rst @@ -0,0 +1 @@ +.. automodule:: pytools.mpi diff --git a/pytools/mpi.py b/pytools/mpi.py index e1650197bf6f6e8fed8e8b75b7b6d5666af1e3e5..9961d76614e1b623b7e96618b03221bc5abc7bce 100644 --- a/pytools/mpi.py +++ b/pytools/mpi.py @@ -1,3 +1,41 @@ +__copyright__ = """ +Copyright (C) 2009-2019 Andreas Kloeckner +Copyright (C) 2022 University of Illinois Board of Trustees +""" + +__license__ = """ +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + +__doc__ = """ +MPI helper functionality +======================== + +.. autofunction:: check_for_mpi_relaunch +.. autofunction:: run_with_mpi_ranks +.. autofunction:: pytest_raises_on_rank +""" + +from contextlib import contextmanager, AbstractContextManager +from typing import Generator, Tuple, Union, Type + + def check_for_mpi_relaunch(argv): if argv[1] != "--mpi-relaunch": return @@ -26,3 +64,23 @@ def run_with_mpi_ranks(py_script, ranks, callable_, args=(), kwargs=None): check_call(["mpirun", "-np", str(ranks), sys.executable, py_script, "--mpi-relaunch", callable_and_args], env=newenv) + + +@contextmanager +def pytest_raises_on_rank(my_rank: int, fail_rank: int, + expected_exception: Union[Type[BaseException], + Tuple[Type[BaseException], ...]]) \ + -> Generator[AbstractContextManager, None, None]: + """ + Like :func:`pytest.raises`, but only expect an exception on rank *fail_rank*. + """ + import pytest + from contextlib import nullcontext + + if my_rank == fail_rank: + cm: AbstractContextManager = pytest.raises(expected_exception) + else: + cm = nullcontext() + + with cm as exc: + yield exc diff --git a/test/test_mpi.py b/test/test_mpi.py new file mode 100644 index 0000000000000000000000000000000000000000..4bbd34771d80dbcdcb5f0c53c06375bcdea7624a --- /dev/null +++ b/test/test_mpi.py @@ -0,0 +1,46 @@ +__copyright__ = "Copyright (C) 2022 University of Illinois Board of Trustees" + +__license__ = """ +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + +import pytest + + +def test_pytest_raises_on_rank(): + from pytools.mpi import pytest_raises_on_rank + + def fail(my_rank: int, fail_rank: int) -> None: + if my_rank == fail_rank: + raise ValueError("test failure") + + with pytest.raises(ValueError): + fail(0, 0) + + fail(0, 1) + + with pytest_raises_on_rank(0, 0, ValueError): + # Generates an exception, and pytest_raises_on_rank + # expects one. + fail(0, 0) + + with pytest_raises_on_rank(0, 1, ValueError): + # Generates no exception, and pytest_raises_on_rank + # does not expect one. + fail(0, 1)