Mark CV-CUDA tests with needs_cvcuda (#9305)

Author: zy1git

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
 

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.

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"])