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
« 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
4import sys
6import pytest
8from cuda.bindings import driver, runtime
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
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()
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
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)
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
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")