Use subtests where appropriate everywhere

cuda_bindings/pixi.toml

@@ -14,7 +14,7 @@ cuda-version = ["12.*", "13.3.*"]
 
 [feature.test.dependencies]
 cuda-bindings = { path = "." }
-pytest = ">=6.2.4"
+pytest = ">=9.0"
 pytest-benchmark = ">=3.4.1"
 pytest-randomly = "*"
 pytest-repeat = "*"

cuda_bindings/tests/nvml/test_compute_mode.py

@@ -18,15 +18,18 @@
 
 
 @pytest.mark.skipif(sys.platform == "win32", reason="Test not supported on Windows")
-def test_compute_mode_supported_nonroot(all_devices):
+def test_compute_mode_supported_nonroot(all_devices, subtests):
     for device in all_devices:
-        with unsupported_before(device, None):
-            original_compute_mode = nvml.device_get_compute_mode(device)
-
-        for cm in COMPUTE_MODES:
-            try:
-                nvml.device_set_compute_mode(device, cm)
-            except nvml.NoPermissionError:
-                pytest.skip("Insufficient permissions to set compute mode")
-            nvml.device_set_compute_mode(device, original_compute_mode)
-            assert original_compute_mode == nvml.device_get_compute_mode(device), "Compute mode shouldn't have changed"
+        with subtests.test(device=device):
+            with unsupported_before(device, None):
+                original_compute_mode = nvml.device_get_compute_mode(device)
+
+            for cm in COMPUTE_MODES:
+                try:
+                    nvml.device_set_compute_mode(device, cm)
+                except nvml.NoPermissionError:
+                    pytest.skip("Insufficient permissions to set compute mode")
+                nvml.device_set_compute_mode(device, original_compute_mode)
+                assert original_compute_mode == nvml.device_get_compute_mode(device), (
+                    "Compute mode shouldn't have changed"
+                )

cuda_bindings/tests/nvml/test_device.py

@@ -38,11 +38,12 @@ def test_clk_mon_status_t():
     assert not hasattr(obj, "clk_mon_list_size")
 
 
-def test_current_clock_freqs(all_devices):
+def test_current_clock_freqs(all_devices, subtests):
     for device in all_devices:
-        with unsupported_before(device, None):
-            clk_freqs = nvml.device_get_current_clock_freqs(device)
-        assert isinstance(clk_freqs, str)
+        with subtests.test(device=device):
+            with unsupported_before(device, None):
+                clk_freqs = nvml.device_get_current_clock_freqs(device)
+            assert isinstance(clk_freqs, str)
 
 
 def test_grid_licensable_features(all_devices):
@@ -69,17 +70,18 @@ def test_get_handle_by_uuidv(all_devices):
         assert new_handle == device
 
 
-def test_get_nv_link_supported_bw_modes(all_devices):
+def test_get_nv_link_supported_bw_modes(all_devices, subtests):
     for device in all_devices:
-        with unsupported_before(device, None):
-            modes = nvml.device_get_nvlink_supported_bw_modes(device)
-        assert isinstance(modes, nvml.NvlinkSupportedBwModes_v1)
-        # #define NVML_NVLINK_TOTAL_SUPPORTED_BW_MODES 23
-        assert len(modes.bw_modes) <= 23
-        assert not hasattr(modes, "total_bw_modes")
+        with subtests.test(device=device):
+            with unsupported_before(device, None):
+                modes = nvml.device_get_nvlink_supported_bw_modes(device)
+            assert isinstance(modes, nvml.NvlinkSupportedBwModes_v1)
+            # #define NVML_NVLINK_TOTAL_SUPPORTED_BW_MODES 23
+            assert len(modes.bw_modes) <= 23
+            assert not hasattr(modes, "total_bw_modes")
 
-        for mode in modes.bw_modes:
-            assert isinstance(mode, np.uint8)
+            for mode in modes.bw_modes:
+                assert isinstance(mode, np.uint8)
 
 
 def test_device_get_pdi(all_devices):
@@ -88,62 +90,67 @@ def test_device_get_pdi(all_devices):
         assert isinstance(pdi, int)
 
 
-def test_device_get_performance_modes(all_devices):
+def test_device_get_performance_modes(all_devices, subtests):
     for device in all_devices:
-        with unsupported_before(device, None):
-            modes = nvml.device_get_performance_modes(device)
-        assert isinstance(modes, str)
+        with subtests.test(device=device):
+            with unsupported_before(device, None):
+                modes = nvml.device_get_performance_modes(device)
+            assert isinstance(modes, str)
 
 
 @pytest.mark.skipif(cuda_version_less_than(13010), reason="Introduced in 13.1")
-def test_device_get_unrepairable_memory_flag(all_devices):
+def test_device_get_unrepairable_memory_flag(all_devices, subtests):
     for device in all_devices:
-        with unsupported_before(device, None):
-            status = nvml.device_get_unrepairable_memory_flag_v1(device)
-        assert isinstance(status, int)
+        with subtests.test(device=device):
+            with unsupported_before(device, None):
+                status = nvml.device_get_unrepairable_memory_flag_v1(device)
+            assert isinstance(status, int)
 
 
-def test_device_vgpu_get_heterogeneous_mode(all_devices):
+def test_device_vgpu_get_heterogeneous_mode(all_devices, subtests):
     for device in all_devices:
-        with unsupported_before(device, None):
-            mode = nvml.device_get_vgpu_heterogeneous_mode(device)
-        assert isinstance(mode, int)
+        with subtests.test(device=device):
+            with unsupported_before(device, None):
+                mode = nvml.device_get_vgpu_heterogeneous_mode(device)
+            assert isinstance(mode, int)
 
 
 @pytest.mark.skipif(cuda_version_less_than(13010), reason="Introduced in 13.1")
-def test_read_prm_counters(all_devices):
+def test_read_prm_counters(all_devices, subtests):
     for device in all_devices:
-        counters = nvml.PRMCounter_v1(5)
-        with unsupported_before(device, None):
-            read_counters = nvml.device_read_prm_counters_v1(device, counters)
-        assert counters is read_counters
-        assert len(read_counters) == 5
+        with subtests.test(device=device):
+            counters = nvml.PRMCounter_v1(5)
+            with unsupported_before(device, None):
+                read_counters = nvml.device_read_prm_counters_v1(device, counters)
+            assert counters is read_counters
+            assert len(read_counters) == 5
 
 
 @pytest.mark.thread_unsafe(reason="API appears to be thread-unsafe (2026-06)")
-def test_read_write_prm(all_devices):
+def test_read_write_prm(all_devices, subtests):
     for device in all_devices:
-        # Docs say supported in BLACKWELL or later
-        with unsupported_before(device, None):
-            try:
-                result = nvml.device_read_write_prm_v1(device, b"012345678")
-            except nvml.NoPermissionError:
-                pytest.skip("No permission to read/write PRM")
-        assert isinstance(result, tuple)
-        assert isinstance(result[0], int)
-        assert isinstance(result[1], bytes)
-
-
-def test_get_power_management_limit(all_devices):
+        with subtests.test(device=device):
+            # Docs say supported in BLACKWELL or later
+            with unsupported_before(device, None):
+                try:
+                    result = nvml.device_read_write_prm_v1(device, b"012345678")
+                except nvml.NoPermissionError:
+                    pytest.skip("No permission to read/write PRM")
+            assert isinstance(result, tuple)
+            assert isinstance(result[0], int)
+            assert isinstance(result[1], bytes)
+
+
+def test_get_power_management_limit(all_devices, subtests):
     for device in all_devices:
         # Docs say supported on KEPLER or later
-        with unsupported_before(device, None):
+        with subtests.test(device=device), unsupported_before(device, nvml.DeviceArch.KEPLER):
             nvml.device_get_power_management_limit(device)
 
 
-def test_set_power_management_limit(all_devices):
+def test_set_power_management_limit(all_devices, subtests):
     for device in all_devices:
-        with unsupported_before(device, nvml.DeviceArch.KEPLER):
+        with subtests.test(device=device), unsupported_before(device, nvml.DeviceArch.KEPLER):
             try:
                 nvml.device_set_power_management_limit_v2(device, nvml.PowerScope.GPU, 10000)
             except nvml.NoPermissionError:
@@ -152,18 +159,19 @@ def test_set_power_management_limit(all_devices):
                 pytest.skip("Invalid argument when setting power management limit -- probably unsupported")
 
 
-def test_set_temperature_threshold(all_devices):
+def test_set_temperature_threshold(all_devices, subtests):
     for device in all_devices:
-        # Docs say supported on MAXWELL or newer
-        with unsupported_before(device, None):
-            temp = nvml.device_get_temperature_threshold(
-                device, nvml.TemperatureThresholds.TEMPERATURE_THRESHOLD_ACOUSTIC_CURR
-            )
-        try:
-            nvml.device_set_temperature_threshold(
-                device, nvml.TemperatureThresholds.TEMPERATURE_THRESHOLD_ACOUSTIC_CURR, temp
-            )
-        except nvml.NoPermissionError:
-            pytest.skip("No permission to set temperature threshold")
-        except nvml.InvalidArgumentError:
-            pytest.skip("Invalid argument when setting temperature threshold -- this is probably the temp type")
+        with subtests.test(device=device):
+            # Docs say supported on MAXWELL or newer
+            with unsupported_before(device, None):
+                temp = nvml.device_get_temperature_threshold(
+                    device, nvml.TemperatureThresholds.TEMPERATURE_THRESHOLD_ACOUSTIC_CURR
+                )
+            try:
+                nvml.device_set_temperature_threshold(
+                    device, nvml.TemperatureThresholds.TEMPERATURE_THRESHOLD_ACOUSTIC_CURR, temp
+                )
+            except nvml.NoPermissionError:
+                pytest.skip("No permission to set temperature threshold")
+            except nvml.InvalidArgumentError:
+                pytest.skip("Invalid argument when setting temperature threshold -- this is probably the temp type")

cuda_bindings/tests/nvml/test_gpu.py

@@ -10,7 +10,7 @@
 from .conftest import unsupported_before
 
 
-def test_gpu_get_module_id(nvml_init):
+def test_gpu_get_module_id(nvml_init, subtests):
     # Unique module IDs cannot exceed the number of GPUs on the system
     device_count = nvml.device_get_count_v2()
 
@@ -21,23 +21,25 @@ def test_gpu_get_module_id(nvml_init):
         if util.is_vgpu(device):
             continue
 
-        with unsupported_before(device, None):
-            module_id = nvml.device_get_module_id(device)
-        assert isinstance(module_id, int)
+        with subtests.test(i=i):
+            with unsupported_before(device, None):
+                module_id = nvml.device_get_module_id(device)
+            assert isinstance(module_id, int)
 
 
-def test_gpu_get_platform_info(all_devices):
+def test_gpu_get_platform_info(all_devices, subtests):
     for device in all_devices:
-        if util.is_vgpu(device):
-            pytest.skip(f"Not supported on vGPU device {device}")
+        with subtests.test(device=device):
+            if util.is_vgpu(device):
+                pytest.skip(f"Not supported on vGPU device {device}")
 
-        # Documentation says Blackwell or newer only, but this does seem to pass
-        # on some newer GPUs.
+            # Documentation says Blackwell or newer only, but this does seem to pass
+            # on some newer GPUs.
 
-        with unsupported_before(device, None):
-            platform_info = nvml.device_get_platform_info(device)
+            with unsupported_before(device, None):
+                platform_info = nvml.device_get_platform_info(device)
 
-        assert isinstance(platform_info, (nvml.PlatformInfo_v1, nvml.PlatformInfo_v2))
+            assert isinstance(platform_info, (nvml.PlatformInfo_v1, nvml.PlatformInfo_v2))
 
 
 # TODO: Test APIs related to GPU instances, which require specific hardware and root
@@ -58,10 +60,10 @@ def test_conf_compute_attestation_report_t(all_devices):
     assert report.nonce.dtype == np.uint8
 
 
-def test_gpu_conf_compute_attestation_report(all_devices):
+def test_gpu_conf_compute_attestation_report(all_devices, subtests):
     for device in all_devices:
         # Documentation says AMPERE or newer
-        with unsupported_before(device, None), pytest.raises(nvml.UnknownError):
+        with subtests.test(device=device), unsupported_before(device, None), pytest.raises(nvml.UnknownError):
             # The nonce string is nonsensical, so if this "works", we expect an UnknownError
             nvml.device_get_conf_compute_gpu_attestation_report(device, nonce=b"12345678")
 
@@ -74,9 +76,9 @@ def test_conf_compute_gpu_certificate_t():
     assert len(cert.attestation_cert_chain) == 0
 
 
-def test_conf_compute_gpu_certificate(all_devices):
+def test_conf_compute_gpu_certificate(all_devices, subtests):
     for device in all_devices:
         # Documentation says AMPERE or newer
-        with unsupported_before(device, None), pytest.raises(nvml.UnknownError):
+        with subtests.test(device=device), unsupported_before(device, None), pytest.raises(nvml.UnknownError):
             # This is expected to fail if the device doesn't have a proper certificate
             nvml.device_get_conf_compute_gpu_certificate(device)

cuda_bindings/tests/nvml/test_pci.py

@@ -9,12 +9,13 @@
 from .conftest import unsupported_before
 
 
-def test_discover_gpus(all_devices):
+def test_discover_gpus(all_devices, subtests):
     for device in all_devices:
-        pci_info = nvml.device_get_pci_info_v3(device)
-        # Docs say this should be supported on PASCAL and later
-        with unsupported_before(device, None), contextlib.suppress(nvml.OperatingSystemError):
-            nvml.device_discover_gpus(pci_info.ptr)
+        with subtests.test(device=device):
+            pci_info = nvml.device_get_pci_info_v3(device)
+            # Docs say this should be supported on PASCAL and later
+            with unsupported_before(device, None), contextlib.suppress(nvml.OperatingSystemError):
+                nvml.device_discover_gpus(pci_info.ptr)
 
 
 def test_bridge_chip_hierarchy_t():
@@ -24,12 +25,13 @@ def test_bridge_chip_hierarchy_t():
     assert isinstance(hierarchy.bridge_chip_info, nvml.BridgeChipInfo)
 
 
-def test_bridge_chip_info(all_devices):
+def test_bridge_chip_info(all_devices, subtests):
     for device in all_devices:
-        with unsupported_before(device, None):
-            info = nvml.device_get_bridge_chip_info(device)
-        assert isinstance(info, nvml.BridgeChipHierarchy)
-        for entry in info.bridge_chip_info:
-            assert isinstance(entry, nvml.BridgeChipInfo)
-            assert isinstance(entry.type, int)
-            assert isinstance(entry.fw_version, int)
+        with subtests.test(device=device):
+            with unsupported_before(device, None):
+                info = nvml.device_get_bridge_chip_info(device)
+            assert isinstance(info, nvml.BridgeChipHierarchy)
+            for entry in info.bridge_chip_info:
+                assert isinstance(entry, nvml.BridgeChipInfo)
+                assert isinstance(entry.type, int)
+                assert isinstance(entry.fw_version, int)