Coverage for cuda/bindings/_test_helpers/mempool.py: 25.00%

36 statements  

« prev     ^ index     » next       coverage.py v7.14.1, created at 2026-06-13 01:38 +0000

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

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

3 

4import sys 

5 

6import pytest 

7 

8from cuda.bindings import driver, runtime 

9 

10 

11# Keep in sync with the fallback in cuda_core/tests/conftest.py. The cuda_core 

12# copy is intentionally simpler because it only handles cuda_core CUDAError 

13# exceptions when this helper is absent from older published bindings. 

14def is_windows_mcdm_device(device=0): 

15 if sys.platform != "win32": 

16 return False 

17 import cuda.bindings.nvml as nvml 

18 

19 device_id = int(getattr(device, "device_id", device)) 

20 (err,) = driver.cuInit(0) 

21 if err != driver.CUresult.CUDA_SUCCESS: 

22 return False 

23 err, pci_bus_id = driver.cuDeviceGetPCIBusId(13, device_id) 

24 if err != driver.CUresult.CUDA_SUCCESS: 

25 return False 

26 pci_bus_id = pci_bus_id.split(b"\x00", 1)[0].decode("ascii") 

27 nvml.init_v2() 

28 try: 

29 handle = nvml.device_get_handle_by_pci_bus_id_v2(pci_bus_id) 

30 current, _ = nvml.device_get_driver_model_v2(handle) 

31 return current == nvml.DriverModel.DRIVER_MCDM 

32 finally: 

33 nvml.shutdown() 

34 

35 

36def xfail_if_mempool_oom(err_or_exc, api_name=None, device=0): 

37 if api_name is not None and not isinstance(api_name, str): 1bcd

38 device = api_name 

39 api_name = None 

40 

41 is_oom = err_or_exc in ( 1bcd

42 driver.CUresult.CUDA_ERROR_OUT_OF_MEMORY, 

43 runtime.cudaError_t.cudaErrorMemoryAllocation, 

44 ) or "CUDA_ERROR_OUT_OF_MEMORY" in str(err_or_exc) 

45 

46 if not is_oom: 1bcd

47 return 1bcd

48 try: 

49 is_windows_mcdm = is_windows_mcdm_device(device) 

50 except Exception: 

51 # If MCDM detection fails, leave the primary test failure visible. 

52 return 

53 if not is_windows_mcdm: 

54 return 

55 

56 api_context = f"{api_name} " if api_name else "" 

57 pytest.xfail(f"{api_context}could not reserve VA for mempool operations on Windows MCDM")