Coverage for cuda / bindings / _test_helpers / arch_check.py: 83.87%

1# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 

2# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE 

3 

4 

5from contextlib import contextmanager (empty)

6from functools import cache (empty)

7 

8import pytest (empty)

9 

10from cuda.bindings import nvml (empty)

11 

12 

13@cache (empty)

14def hardware_supports_nvml(): (empty)

15 """ 

16 Tries to call the simplest NVML API possible to see if just the basics 

17 works. If not we are probably on one of the platforms where NVML is not 

18 supported at all (e.g. Jetson Orin). 

19 """ 

20 nvml.init_v2() (empty)

21 try: (empty)

22 nvml.system_get_driver_branch() (empty)

23 except (nvml.NotSupportedError, nvml.UnknownError): 

24 return False 

25 else: 

26 return True (empty)

27 finally: 

28 nvml.shutdown() (empty)

29 

30 

31@contextmanager (empty)

32def unsupported_before(device: int, expected_device_arch: nvml.DeviceArch | str | None): (empty)

33 device_arch = nvml.device_get_architecture(device) 50 ctx1wxycdzeRfghABijSTUVCklmnopDEFGqWrsHbtIXJKLMYNOuPQv

34 

35 if isinstance(expected_device_arch, nvml.DeviceArch): 50 ctx1wxycdzeRfghABijSTUVCklmnopDEFGqWrsHbtIXJKLMYNOuPQv

36 expected_device_arch_int = int(expected_device_arch) 9 ctx1RSTUVWbXY

37 elif expected_device_arch == "FERMI": 42 ctx1wxycdzefghABijCklmnopDEFGqrsHbtIJKLMNOuPQv

38 expected_device_arch_int = 1 2 ctx1Gb

39 else: 

40 expected_device_arch_int = 0 41 ctx1wxycdzefghABijCklmnopDEFqrsHbtIJKLMNOuPQv

41 

42 if expected_device_arch is None or expected_device_arch == "HAS_INFOROM" or device_arch == nvml.DeviceArch.UNKNOWN: 50 ctx1wxycdzeRfghABijSTUVCklmnopDEFGqWrsHbtIXJKLMYNOuPQv

43 # In this case, we don't /know/ if it will fail, but we are ok if it 

44 # does or does not. 

45 

46 # TODO: There are APIs that are documented as supported only if the 

47 # device has an InfoROM, but I couldn't find a way to detect that. For 

48 # now, they are just handled as "possibly failing". 

49 

50 try: 41 ctx1wxycdzefghABijCklmnopDEFqrsHbtIJKLMNOuPQv

51 yield 41 ctx1wxycdzefghABijCklmnopDEFqrsHbtIJKLMNOuPQv

52 except nvml.NotSupportedError: 21 ctx1cdefghijklmnopqrsbtuv

53 # The API call raised NotSupportedError, so we skip the test, but 

54 # don't fail it 

55 pytest.skip( 21 ctx1cdefghijklmnopqrsbtuv

56 f"Unsupported call for device architecture {nvml.DeviceArch(device_arch).name} " 

57 f"on device '{nvml.device_get_name(device)}'" 

58 ) 

59 # If the API call worked, just continue 

60 elif int(device_arch) < expected_device_arch_int: 10 ctx1RSTUVGWbXY

61 # In this case, we /know/ if will fail, and we want to assert that it does. 

62 with pytest.raises(nvml.NotSupportedError): 

63 yield 

64 # The above call was unsupported, so the rest of the test is skipped 

65 pytest.skip(f"Unsupported before {expected_device_arch.name}, got {nvml.device_get_name(device)}") 

66 else: 

67 # In this case, we /know/ it should work, and if it fails, the test should fail. 

68 yield 10 ctx1RSTUVGWbXY