diff --git a/test/common_utils.py b/test/common_utils.py index e3fa464b5ea..24ebb1376c3 100644 --- a/test/common_utils.py +++ b/test/common_utils.py @@ -29,6 +29,7 @@ IN_RE_WORKER = os.environ.get("INSIDE_RE_WORKER") is not None IN_FBCODE = os.environ.get("IN_FBCODE_TORCHVISION") == "1" CUDA_NOT_AVAILABLE_MSG = "CUDA device not available" +CVCUDA_NOT_AVAILABLE_MSG = "CV-CUDA not available" MPS_NOT_AVAILABLE_MSG = "MPS device not available" OSS_CI_GPU_NO_CUDA_MSG = "We're in an OSS GPU machine, and this test doesn't need cuda." @@ -134,6 +135,12 @@ def needs_cuda(test_func): return pytest.mark.needs_cuda(test_func) +def needs_cvcuda(test_func): + import pytest # noqa + + return pytest.mark.needs_cvcuda(test_func) + + def needs_mps(test_func): import pytest # noqa diff --git a/test/conftest.py b/test/conftest.py index a9768598ded..dadfc9088c7 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -5,7 +5,9 @@ import torch from common_utils import ( + _is_cvcuda_available, CUDA_NOT_AVAILABLE_MSG, + CVCUDA_NOT_AVAILABLE_MSG, IN_FBCODE, IN_OSS_CI, IN_RE_WORKER, @@ -17,6 +19,7 @@ def pytest_configure(config): # register an additional marker (see pytest_collection_modifyitems) config.addinivalue_line("markers", "needs_cuda: mark for tests that rely on a CUDA device") + config.addinivalue_line("markers", "needs_cvcuda: mark for tests that rely on CV-CUDA") config.addinivalue_line("markers", "needs_mps: mark for tests that rely on a MPS device") config.addinivalue_line("markers", "dont_collect: mark for tests that should not be collected") config.addinivalue_line("markers", "opcheck_only_one: only opcheck one parametrization") @@ -41,6 +44,9 @@ def pytest_collection_modifyitems(items): # @pytest.mark.parametrize('device', cpu_and_cuda()) # the "instances" of the tests where device == 'cuda' will have the 'needs_cuda' mark, # and the ones with device == 'cpu' won't have the mark. + needs_cvcuda = item.get_closest_marker("needs_cvcuda") is not None + if needs_cvcuda: + item.add_marker(pytest.mark.needs_cuda) needs_cuda = item.get_closest_marker("needs_cuda") is not None needs_mps = item.get_closest_marker("needs_mps") is not None @@ -52,6 +58,9 @@ def pytest_collection_modifyitems(items): if needs_mps and not torch.backends.mps.is_available(): item.add_marker(pytest.mark.skip(reason=MPS_NOT_AVAILABLE_MSG)) + if needs_cvcuda and not _is_cvcuda_available(): + item.add_marker(pytest.mark.skip(reason=CVCUDA_NOT_AVAILABLE_MSG)) + if IN_FBCODE: # fbcode doesn't like skipping tests, so instead we just don't collect the test # so that they don't even "exist", hence the continue statements. diff --git a/test/test_transforms_v2.py b/test/test_transforms_v2.py index 3ce603c3ed2..b9f440bd545 100644 --- a/test/test_transforms_v2.py +++ b/test/test_transforms_v2.py @@ -37,6 +37,7 @@ make_video, make_video_tensor, needs_cuda, + needs_cvcuda, set_rng_seed, ) @@ -52,17 +53,8 @@ from torchvision.transforms.v2 import functional as F from torchvision.transforms.v2._utils import check_type, is_pure_tensor from torchvision.transforms.v2.functional._geometry import _get_perspective_coeffs, _parallelogram_to_bounding_boxes -from torchvision.transforms.v2.functional._utils import ( - _get_kernel, - _import_cvcuda, - _is_cvcuda_available, - _register_kernel_internal, -) - +from torchvision.transforms.v2.functional._utils import _get_kernel, _import_cvcuda, _register_kernel_internal -CVCUDA_AVAILABLE = _is_cvcuda_available() -if CVCUDA_AVAILABLE: - cvcuda = _import_cvcuda() # turns all warnings into errors for this module pytestmark = [pytest.mark.filterwarnings("error")] @@ -1242,7 +1234,7 @@ def test_kernel_video(self): make_image, pytest.param( make_image_cvcuda, - marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="CVCUDA is not available"), + marks=pytest.mark.needs_cvcuda, ), make_bounding_boxes, make_segmentation_mask, @@ -1262,7 +1254,7 @@ def test_functional(self, make_input): pytest.param( F._geometry._horizontal_flip_image_cvcuda, None, - marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="CVCUDA is not available"), + marks=pytest.mark.needs_cvcuda, ), (F.horizontal_flip_bounding_boxes, tv_tensors.BoundingBoxes), (F.horizontal_flip_mask, tv_tensors.Mask), @@ -1283,7 +1275,7 @@ def test_functional_signature(self, kernel, input_type): make_image, pytest.param( make_image_cvcuda, - marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="CVCUDA is not available"), + marks=pytest.mark.needs_cvcuda, ), make_bounding_boxes, make_segmentation_mask, @@ -1304,7 +1296,7 @@ def test_transform(self, make_input, device): make_image, pytest.param( make_image_cvcuda, - marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="CVCUDA is not available"), + marks=pytest.mark.needs_cvcuda, ), ], ) @@ -1372,7 +1364,7 @@ def test_keypoints_correctness(self, fn): make_image, pytest.param( make_image_cvcuda, - marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="CVCUDA is not available"), + marks=pytest.mark.needs_cvcuda, ), make_bounding_boxes, make_segmentation_mask, @@ -1884,7 +1876,7 @@ def test_kernel_video(self): make_image, pytest.param( make_image_cvcuda, - marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="CVCUDA is not available"), + marks=pytest.mark.needs_cvcuda, ), make_bounding_boxes, make_segmentation_mask, @@ -1904,7 +1896,7 @@ def test_functional(self, make_input): pytest.param( F._geometry._vertical_flip_image_cvcuda, None, - marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="CVCUDA is not available"), + marks=pytest.mark.needs_cvcuda, ), (F.vertical_flip_bounding_boxes, tv_tensors.BoundingBoxes), (F.vertical_flip_mask, tv_tensors.Mask), @@ -1925,7 +1917,7 @@ def test_functional_signature(self, kernel, input_type): make_image, pytest.param( make_image_cvcuda, - marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="CVCUDA is not available"), + marks=pytest.mark.needs_cvcuda, ), make_bounding_boxes, make_segmentation_mask, @@ -1944,7 +1936,7 @@ def test_transform(self, make_input, device): make_image, pytest.param( make_image_cvcuda, - marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="CVCUDA is not available"), + marks=pytest.mark.needs_cvcuda, ), ], ) @@ -2008,7 +2000,7 @@ def test_keypoints_correctness(self, fn): make_image, pytest.param( make_image_cvcuda, - marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="CVCUDA is not available"), + marks=pytest.mark.needs_cvcuda, ), make_bounding_boxes, make_segmentation_mask, @@ -6794,8 +6786,7 @@ def test_functional_error(self): F.pil_to_tensor(object()) -@pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA") -@needs_cuda +@needs_cvcuda class TestToCVCUDATensor: @pytest.mark.parametrize("image_type", (torch.Tensor, tv_tensors.Image)) @pytest.mark.parametrize("dtype", [torch.uint8, torch.uint16, torch.float32, torch.float64]) @@ -6813,7 +6804,7 @@ def test_functional_and_transform(self, image_type, dtype, device, color_space, assert is_pure_tensor(image) output = fn(image) - assert isinstance(output, cvcuda.Tensor) + assert isinstance(output, _import_cvcuda().Tensor) assert F.get_size(output) == F.get_size(image) assert output is not None @@ -6856,9 +6847,8 @@ def test_round_trip(self, dtype, device, color_space, batch_size): assert result_tensor.shape[0] == batch_size -@pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA") -@needs_cuda -class TestCVDUDAToTensor: +@needs_cvcuda +class TestCVCUDAToTensor: @pytest.mark.parametrize("dtype", [torch.uint8, torch.uint16, torch.float32, torch.float64]) @pytest.mark.parametrize("device", cpu_and_cuda()) @pytest.mark.parametrize("color_space", ["RGB", "GRAY"])