add compute_capabilities_after

Signed-off-by: YunLiu <[email protected]>

Author: KumoLiu

monai/utils/__init__.py

@@ -123,6 +123,7 @@
     run_eval,
     version_geq,
     version_leq,
+    compute_capabilities_after,
 )
 from .nvtx import Range
 from .ordering import Ordering

monai/utils/module.py

@@ -26,6 +26,7 @@
 from re import match
 from types import FunctionType, ModuleType
 from typing import Any, cast
+import pynvml
 
 import torch
 
@@ -634,3 +635,36 @@ def pytorch_after(major: int, minor: int, patch: int = 0, current_ver_string: st
     if is_prerelease:
         return False
     return True
+
+@functools.lru_cache(None)
+def compute_capabilities_after(major: int, minor: int = 0, current_ver_string: str | None = None) -> bool:
+    """
+    Compute whether the current system GPU CUDA compute capability is after or equal to the specified version.
+    The current system GPU CUDA compute capability is determined by the first GPU in the system.
+    The compared version is a string in the form of "major.minor".
+    
+    Args:
+        major: major version number to be compared with.
+        minor: minor version number to be compared with. Defaults to 0.
+        current_ver_string: if None, the current system GPU CUDA compute capability will be used.
+    
+    Returns:
+        True if the current system GPU CUDA compute capability is greater than or equal to the specified version.
+    """
+    if current_ver_string is None:
+        pynvml.nvmlInit()
+        handle = pynvml.nvmlDeviceGetHandleByIndex(0)  # get the first GPU
+        major_c, minor_c = pynvml.nvmlDeviceGetCudaComputeCapability(handle)
+        pynvml.nvmlShutdown()
+        current_ver_string = f"{major_c}.{minor_c}"
+    
+    ver, has_ver = optional_import("packaging.version", name="parse")
+    if has_ver:
+        return ver(".".join((f"{major}", f"{minor}"))) <= ver(f"{current_ver_string}")  # type: ignore
+    parts = f"{current_ver_string}".split("+", 1)[0].split(".", 2)
+    while len(parts) < 2:
+        parts += ["0"]
+    c_major, c_minor = parts[:2]
+    c_mn = int(c_major), int(c_minor)
+    mn = int(major), int(minor)
+    return c_mn > mn

tests/test_pytorch_version_after.py

@@ -21,7 +21,7 @@
 from monai.networks import trt_compile
 from monai.networks.nets import UNet, cell_sam_wrapper, vista3d132
 from monai.utils import min_version, optional_import
-from tests.utils import SkipIfAtLeastPyTorchVersion, skip_if_no_cuda, skip_if_quick, skip_if_windows
+from tests.utils import SkipIfAtLeastPyTorchVersion, skip_if_no_cuda, skip_if_quick, skip_if_windows, SkipIfBeforeComputeCapabilityVersion
 
 trt, trt_imported = optional_import("tensorrt", "10.1.0", min_version)
 polygraphy, polygraphy_imported = optional_import("polygraphy")
@@ -36,6 +36,7 @@
 @skip_if_quick
 @unittest.skipUnless(trt_imported, "tensorrt is required")
 @unittest.skipUnless(polygraphy_imported, "polygraphy is required")
+@SkipIfBeforeComputeCapabilityVersion((7, 0))
 class TestTRTCompile(unittest.TestCase):
 
     def setUp(self):

tests/test_trt_compile.py

@@ -47,7 +47,7 @@
 from monai.networks import convert_to_onnx, convert_to_torchscript
 from monai.utils import optional_import
 from monai.utils.misc import MONAIEnvVars
-from monai.utils.module import pytorch_after
+from monai.utils.module import pytorch_after, compute_capabilities_after
 from monai.utils.tf32 import detect_default_tf32
 from monai.utils.type_conversion import convert_data_type
 
@@ -286,6 +286,20 @@ def __call__(self, obj):
         )(obj)
 
 
+class SkipIfBeforeComputeCapabilityVersion:
+    """Decorator to be used if test should be skipped
+    with Compute Capability older than that given."""
+
+    def __init__(self, compute_capability_tuple):
+        self.min_version = compute_capability_tuple
+        self.version_too_old = not compute_capabilities_after(*compute_capability_tuple)
+
+    def __call__(self, obj):
+        return unittest.skipIf(
+            self.version_too_old, f"Skipping tests that fail on Compute Capability versions before: {self.min_version}"
+        )(obj)
+
+
 def is_main_test_process():
     ps = torch.multiprocessing.current_process()
     if not ps or not hasattr(ps, "name"):