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) 56 ctx1xyzcdeABfZghiCDjk0123ElmnopqFGHIJKrsLbMtNOPQRS4TUVuvWXwY

35 

36 if isinstance(expected_device_arch, nvml.DeviceArch): 56 ctx1xyzcdeABfZghiCDjk0123ElmnopqFGHIJKrsLbMtNOPQRS4TUVuvWXwY

37 expected_device_arch_int = int(expected_device_arch) 7 ctx1Z0123b4

38 elif expected_device_arch == "FERMI": 50 ctx1xyzcdeABfghiCDjkElmnopqFGHIJKrsLbMtNOPQRSTUVuvWXwY

39 expected_device_arch_int = 1 2 ctx1bM

40 else: 

41 expected_device_arch_int = 0 49 ctx1xyzcdeABfghiCDjkElmnopqFGHIJKrsLbtNOPQRSTUVuvWXwY

42 

43 if expected_device_arch is None or expected_device_arch == "HAS_INFOROM" or device_arch == nvml.DeviceArch.UNKNOWN: 56 ctx1xyzcdeABfZghiCDjk0123ElmnopqFGHIJKrsLbMtNOPQRS4TUVuvWXwY

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: 47 ctx1xyzcdeABfghiCDjkElmnopqFGHIJrsLbtNPQRSTUVuvWXwY

52 yield 47 ctx1xyzcdeABfghiCDjkElmnopqFGHIJrsLbtNPQRSTUVuvWXwY

53 except (nvml.NotSupportedError, nvml.FunctionNotFoundError, NvmlSymbolNotFoundError): 22 ctx1cdefghijklmnopqrsbtuvw

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( 22 ctx1cdefghijklmnopqrsbtuvw

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: 10 ctx1Z0123KbMO4

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 10 ctx1Z0123KbMO4