diff --git a/cuda_core/tests/helpers/__init__.py b/cuda_core/tests/helpers/__init__.py index 02cbe6d8e9..ad9d281c16 100644 --- a/cuda_core/tests/helpers/__init__.py +++ b/cuda_core/tests/helpers/__init__.py @@ -1,9 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import functools import os import pathlib import sys +from typing import Union + +from cuda.core._utils.cuda_utils import handle_return CUDA_PATH = os.environ.get("CUDA_PATH") CUDA_INCLUDE_PATH = None @@ -24,3 +28,38 @@ # Import shared platform helpers for tests across repos sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[3] / "cuda_python_test_helpers")) from cuda_python_test_helpers import * # noqa: F403 + + +@functools.cache +def supports_ipc_mempool(device_id: Union[int, object]) -> bool: + """Return True if mempool IPC via POSIX file descriptor is supported. + + Uses cuDeviceGetAttribute(CU_DEVICE_ATTRIBUTE_MEMPOOL_SUPPORTED_HANDLE_TYPES) + to check for CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR support. Does not + require an active CUDA context. + """ + if IS_WSL: # noqa: F405 + return False + + try: + # Lazy import to avoid hard dependency when not running GPU tests + try: + from cuda.bindings import driver # type: ignore + except Exception: + from cuda import cuda as driver # type: ignore + + # Initialize CUDA + handle_return(driver.cuInit(0)) + + # Resolve device id from int or Device-like object + dev_id = int(getattr(device_id, "device_id", device_id)) + + # Query supported mempool handle types bitmask + attr = driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MEMPOOL_SUPPORTED_HANDLE_TYPES + mask = handle_return(driver.cuDeviceGetAttribute(attr, dev_id)) + + # Check POSIX FD handle type support via bitmask + posix_fd = driver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR + return (int(mask) & int(posix_fd)) != 0 + except Exception: + return False diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 0f6288bcf7..71adb4ffc7 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -38,7 +38,7 @@ from cuda.core._memory import IPCBufferDescriptor from cuda.core._utils.cuda_utils import CUDAError, handle_return from cuda.core.utils import StridedMemoryView -from helpers import IS_WINDOWS +from helpers import IS_WINDOWS, supports_ipc_mempool from helpers.buffers import DummyUnifiedMemoryResource from conftest import ( @@ -46,7 +46,6 @@ skip_if_managed_memory_unsupported, skip_if_pinned_memory_unsupported, ) -from cuda_python_test_helpers import supports_ipc_mempool POOL_SIZE = 2097152 # 2MB size diff --git a/cuda_python_test_helpers/cuda_python_test_helpers/__init__.py b/cuda_python_test_helpers/cuda_python_test_helpers/__init__.py index fca190c103..70ecf42ea0 100644 --- a/cuda_python_test_helpers/cuda_python_test_helpers/__init__.py +++ b/cuda_python_test_helpers/cuda_python_test_helpers/__init__.py @@ -2,20 +2,15 @@ # SPDX-License-Identifier: Apache-2.0 import ctypes -import functools import os import platform import sys from contextlib import suppress -from typing import Union - -from cuda.core._utils.cuda_utils import handle_return __all__ = [ "IS_WINDOWS", "IS_WSL", "libc", - "supports_ipc_mempool", "under_compute_sanitizer", ] @@ -68,38 +63,3 @@ def under_compute_sanitizer() -> bool: # Another common indicator: sanitizer injectors are configured via env vars. inj = os.environ.get("CUDA_INJECTION64_PATH", "") return "compute-sanitizer" in inj or "cuda-memcheck" in inj - - -@functools.cache -def supports_ipc_mempool(device_id: Union[int, object]) -> bool: - """Return True if mempool IPC via POSIX file descriptor is supported. - - Uses cuDeviceGetAttribute(CU_DEVICE_ATTRIBUTE_MEMPOOL_SUPPORTED_HANDLE_TYPES) - to check for CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR support. Does not - require an active CUDA context. - """ - if _detect_wsl(): - return False - - try: - # Lazy import to avoid hard dependency when not running GPU tests - try: - from cuda.bindings import driver # type: ignore - except Exception: - from cuda import cuda as driver # type: ignore - - # Initialize CUDA - handle_return(driver.cuInit(0)) - - # Resolve device id from int or Device-like object - dev_id = int(getattr(device_id, "device_id", device_id)) - - # Query supported mempool handle types bitmask - attr = driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MEMPOOL_SUPPORTED_HANDLE_TYPES - mask = handle_return(driver.cuDeviceGetAttribute(attr, dev_id)) - - # Check POSIX FD handle type support via bitmask - posix_fd = driver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR - return (int(mask) & int(posix_fd)) != 0 - except Exception: - return False