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

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)

11from cuda.bindings._internal.utils import FunctionNotFoundError as NvmlSymbolNotFoundError (empty)

12 

13 

14@cache (empty)

15def hardware_supports_nvml(): (empty)

16 """ 

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

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

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

20 """ 

21 nvml.init_v2() (empty)

22 try: (empty)

23 nvml.system_get_driver_branch() (empty)

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

25 return False 

26 else: 

27 return True (empty)

28 finally: 

29 nvml.shutdown() (empty)

30 

31 

32@contextmanager (empty)

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

34 device_arch = nvml.device_get_architecture(device) 54 ctx1wxycdzeVfghABijWXYZCklmnopDEFGH0qrIbJsK1LMNO2PQRtuSTvU

35 

36 if isinstance(expected_device_arch, nvml.DeviceArch): 54 ctx1wxycdzeVfghABijWXYZCklmnopDEFGH0qrIbJsK1LMNO2PQRtuSTvU

37 expected_device_arch_int = int(expected_device_arch) 9 ctx1VWXYZ0b12

38 elif expected_device_arch == "FERMI": 46 ctx1wxycdzefghABijCklmnopDEFGHqrIbJsKLMNOPQRtuSTvU

39 expected_device_arch_int = 1 3 ctx1GbJ

40 else: 

41 expected_device_arch_int = 0 44 ctx1wxycdzefghABijCklmnopDEFHqrIbsKLMNOPQRtuSTvU

42 

43 if expected_device_arch is None or expected_device_arch == "HAS_INFOROM" or device_arch == nvml.DeviceArch.UNKNOWN: 54 ctx1wxycdzeVfghABijWXYZCklmnopDEFGH0qrIbJsK1LMNO2PQRtuSTvU

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

45 # does or does not. 

46 

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

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

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

50 

51 try: 44 ctx1wxycdzefghABijCklmnopDEFHqrIbsKLMNOPQRtuSTvU

52 yield 44 ctx1wxycdzefghABijCklmnopDEFHqrIbsKLMNOPQRtuSTvU

53 except (nvml.NotSupportedError, nvml.FunctionNotFoundError, NvmlSymbolNotFoundError): 21 ctx1cdefghijklmnopqrbstuv

54 # The API call raised NotSupportedError, NVML status FunctionNotFoundError, 

55 # or NvmlSymbolNotFoundError (symbol absent from the loaded NVML DLL), so we 

56 # skip the test but don't fail it 

57 pytest.skip( 21 ctx1cdefghijklmnopqrbstuv

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

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

60 ) 

61 # If the API call worked, just continue 

62 elif int(device_arch) < expected_device_arch_int: 11 ctx1VWXYZG0bJ12

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

64 with pytest.raises(nvml.NotSupportedError): 

65 yield 

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

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

68 else: 

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

70 yield 11 ctx1VWXYZG0bJ12