A case for `pytest` in PyTorch

This is a short post about why I think it would be beneficial for PyTorch to not only use pytest as test runner, but also rely on the other features it provides.

Disclaimer

My experience with the PyTorch test suite is limited as of now. Thus, it might very well be that my view on things is too naive. In that case I'm happy to hear about examples where and adoption of pytest would make a use case significantly harder or outright impossible.

Setup

`test_pytorch.py`:

from torch.testing._internal.common_utils import TestCase
from torch.testing._internal.common_device_type import (
    instantiate_device_type_tests,
    onlyCPU,
    onlyCUDA,
    onlyOnCPUAndCUDA,
)


class TestFoo(TestCase):
    @onlyCPU
    def test_bar(self, device):
        pass


instantiate_device_type_tests(TestFoo, globals())


class TestSpam(TestCase):
    @onlyCUDA
    def test_ham(self, device):
        pass

    @onlyOnCPUAndCUDA
    def test_eggs(self, device):
        if str(device) == "cpu":
            raise AssertionError


instantiate_device_type_tests(TestSpam, globals())

`test_pytest.py`:

import torch
import pytest


def devices(*devices_):
    return pytest.mark.parametrize(
        "device", [torch.device(device) for device in devices_], ids=devices_
    )


cpu = devices("cpu")
cuda = devices("cuda")
cpu_and_cuda = devices("cpu", "cuda")


class TestFoo:
    @cpu
    def test_bar(self, device):
        pass


class TestSpam:
    @cuda
    def test_ham(self, device):
        pass

    @cpu_and_cuda
    def test_eggs(self, device):
        if str(device) == "cpu":
            raise AssertionError

Output

pytest -rA test_pytorch.py

PASSED test_pytorch.py::TestFooCPU::test_bar_cpu
PASSED test_pytorch.py::TestSpamCUDA::test_eggs_cuda
PASSED test_pytorch.py::TestSpamCUDA::test_ham_cuda
SKIPPED [2] test_pytorch.py:11: Only runs on cpu
SKIPPED [2] test_pytorch.py:20: Only runs on cuda
SKIPPED [1] test_pytorch.py:24: onlyOnCPUAndCUDA: doesn't run on meta
FAILED test_pytorch.py::TestSpamCPU::test_eggs_cpu - AssertionError

pytest -rA test_pytest.py

PASSED test_pytest.py::TestFoo::test_bar[cpu]
PASSED test_pytest.py::TestSpam::test_ham[cuda]
PASSED test_pytest.py::TestSpam::test_eggs[cuda]
FAILED test_pytest.py::TestSpam::test_eggs[cpu] - AssertionError

In PyTorch style the tests are double namespaced with the device.
In PyTorch style all tests are instatiated with every device and skipped if the test should be run with a specific device. This clutters the output.
In PyTorch style the message of the skipped tests is not expressive as it misses the name of the test that is skipped.

Selection

Situation	PyTorch style	`pytest` style
Select one test case and run it against all devices	`test_pytorch.py -k "TestSpam"`	`test_pytest.py::TestSpam`
Select one test case and run it against a specific device	`test_pytorch.py::TestSpamCUDA`	`test_pytest.py::TestSpam -k cuda`
Select one test and run it against all devices	`test_pytorch.py -k "TestSpam and test_eggs"`	`test_pytest.py::TestSpam::test_eggs`
Select one test and run it against a specific device	`test_pytorch.py::TestSpamCPU::test_eggs_cpu`	`test_pytest.py::TestSpam::test_eggs[cpu]`

Popular IDEs (PyCharm, VSCode) provide the option to run / debug tests without dropping into a terminal. They rely on the name within the file to select the test (case) to be run. Since in PyTorch style the tests are renamed, one cannot use this feature.

pmeier/pytorch_pytest.md

A case for pytest in PyTorch