diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py index 83f78ef881..f30afe5667 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py @@ -290,6 +290,12 @@ class DescriptorSpec: anchor_rel_dirs_windows=("extras/CUPTI/lib64", "bin"), ctk_root_canary_anchor_libnames=("cudart",), ), + DescriptorSpec( + name="cudla", + packaged_with="ctk", + linux_sonames=("libcudla.so.1",), + site_packages_linux=("nvidia/cu13/lib",), + ), # ----------------------------------------------------------------------- # Third-party / separately packaged libraries # ----------------------------------------------------------------------- @@ -386,6 +392,11 @@ class DescriptorSpec: linux_sonames=("libcuda.so.1",), windows_dlls=("nvcuda.dll",), ), + DescriptorSpec( + name="nvcudla", + packaged_with="driver", + linux_sonames=("libnvcudla.so",), + ), DescriptorSpec( name="nvml", packaged_with="driver", diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py index f8df1f75e4..a7a8965d2e 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py @@ -99,14 +99,18 @@ def _raise_canary_probe_child_process_error( @functools.cache -def _resolve_system_loaded_abs_path_in_subprocess(libname: str) -> str | None: +def _resolve_system_loaded_abs_path_in_subprocess( + libname: str, + *, + timeout: float = _CANARY_PROBE_TIMEOUT_SECONDS, +) -> str | None: """Resolve a canary library's absolute path in a fresh Python subprocess.""" try: result = subprocess.run( # noqa: S603 - trusted argv: current interpreter + internal probe module build_dynamic_lib_subprocess_command(MODE_CANARY, libname), capture_output=True, text=True, - timeout=_CANARY_PROBE_TIMEOUT_SECONDS, + timeout=timeout, check=False, cwd=DYNAMIC_LIB_SUBPROCESS_CWD, ) @@ -127,6 +131,11 @@ def _resolve_system_loaded_abs_path_in_subprocess(libname: str) -> str | None: return None +def _loadable_via_canary_subprocess(libname: str, *, timeout: float = _CANARY_PROBE_TIMEOUT_SECONDS) -> bool: + """Return True if the canary subprocess can resolve ``libname`` via system search.""" + return _resolve_system_loaded_abs_path_in_subprocess(libname, timeout=timeout) is not None + + def _try_ctk_root_canary(ctx: SearchContext) -> str | None: """Try CTK-root canary fallback for descriptor-configured libraries.""" for canary_libname in ctx.desc.ctk_root_canary_anchor_libnames: diff --git a/cuda_pathfinder/cuda/pathfinder/_headers/header_descriptor_catalog.py b/cuda_pathfinder/cuda/pathfinder/_headers/header_descriptor_catalog.py index 81f93638c5..493445c8ad 100644 --- a/cuda_pathfinder/cuda/pathfinder/_headers/header_descriptor_catalog.py +++ b/cuda_pathfinder/cuda/pathfinder/_headers/header_descriptor_catalog.py @@ -134,6 +134,13 @@ class HeaderDescriptorSpec: site_packages_dirs=("nvidia/cu13/include", "nvidia/cuda_nvcc/nvvm/include"), anchor_include_rel_dirs=("nvvm/include",), ), + HeaderDescriptorSpec( + name="cudla", + packaged_with="ctk", + header_basename="cudla.h", + site_packages_dirs=("nvidia/cu13/include",), + available_on_windows=False, + ), # ----------------------------------------------------------------------- # Third-party / separately packaged headers # ----------------------------------------------------------------------- diff --git a/cuda_pathfinder/pyproject.toml b/cuda_pathfinder/pyproject.toml index 2e08bdffae..5f86c0afb9 100644 --- a/cuda_pathfinder/pyproject.toml +++ b/cuda_pathfinder/pyproject.toml @@ -36,6 +36,7 @@ cu13 = [ "cuda-toolkit[cufile]==13.*; sys_platform != 'win32'", "cutensor-cu13", "nvidia-cublasmp-cu13; sys_platform != 'win32'", + "nvidia-cudla; platform_system == 'Linux' and platform_machine == 'aarch64'", "nvidia-cudss-cu13", "nvidia-cufftmp-cu13; sys_platform != 'win32'", "nvidia-cusolvermp-cu13; sys_platform != 'win32'", diff --git a/cuda_pathfinder/tests/conftest.py b/cuda_pathfinder/tests/conftest.py index 47f8ff1612..e8a5e11b39 100644 --- a/cuda_pathfinder/tests/conftest.py +++ b/cuda_pathfinder/tests/conftest.py @@ -29,3 +29,15 @@ def _append(message): request.config.custom_info.append(f"{request.node.name}: {message}") return _append + + +def skip_if_missing_libnvcudla_so(libname: str, *, timeout: float) -> None: + if libname not in ("cudla", "nvcudla"): + return + # Keep the import inside the helper so unrelated import issues do not fail + # pytest collection for the whole test suite. + from cuda.pathfinder._dynamic_libs import load_nvidia_dynamic_lib as load_nvidia_dynamic_lib_module + + if load_nvidia_dynamic_lib_module._loadable_via_canary_subprocess("nvcudla", timeout=timeout): + return + pytest.skip("libnvcudla.so is not loadable via canary subprocess on this host.") diff --git a/cuda_pathfinder/tests/test_driver_lib_loading.py b/cuda_pathfinder/tests/test_driver_lib_loading.py index bfc8e87f49..bf62a17d70 100644 --- a/cuda_pathfinder/tests/test_driver_lib_loading.py +++ b/cuda_pathfinder/tests/test_driver_lib_loading.py @@ -16,6 +16,7 @@ run_load_nvidia_dynamic_lib_in_subprocess, ) +from conftest import skip_if_missing_libnvcudla_so from cuda.pathfinder._dynamic_libs.lib_descriptor import LIB_DESCRIPTORS from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError, LoadedDL from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import ( @@ -147,6 +148,7 @@ def raise_child_process_failed(): error_label="Load subprocess child process", ) if payload.status == STATUS_NOT_FOUND: + skip_if_missing_libnvcudla_so(libname, timeout=timeout) if STRICTNESS == "all_must_work": raise_child_process_failed() info_summary_append(f"Not found: {libname=!r}") diff --git a/cuda_pathfinder/tests/test_find_nvidia_headers.py b/cuda_pathfinder/tests/test_find_nvidia_headers.py index a4ca8df602..e28f64d352 100644 --- a/cuda_pathfinder/tests/test_find_nvidia_headers.py +++ b/cuda_pathfinder/tests/test_find_nvidia_headers.py @@ -21,6 +21,7 @@ import pytest import cuda.pathfinder._headers.find_nvidia_headers as find_nvidia_headers_module +from conftest import skip_if_missing_libnvcudla_so from cuda.pathfinder import LocatedHeaderDir, find_nvidia_header_directory, locate_nvidia_header_directory from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import ( _resolve_system_loaded_abs_path_in_subprocess, @@ -158,6 +159,8 @@ def test_locate_ctk_headers(info_summary_append, libname): h_filename = SUPPORTED_HEADERS_CTK[libname] assert os.path.isfile(os.path.join(hdr_dir, h_filename)) if STRICTNESS == "all_must_work": + if libname == "cudla": + skip_if_missing_libnvcudla_so(libname, timeout=30) assert hdr_dir is not None diff --git a/cuda_pathfinder/tests/test_load_nvidia_dynamic_lib.py b/cuda_pathfinder/tests/test_load_nvidia_dynamic_lib.py index 36487bd58e..401e7dc13f 100644 --- a/cuda_pathfinder/tests/test_load_nvidia_dynamic_lib.py +++ b/cuda_pathfinder/tests/test_load_nvidia_dynamic_lib.py @@ -11,10 +11,14 @@ ) from local_helpers import have_distribution +from conftest import skip_if_missing_libnvcudla_so from cuda.pathfinder import DynamicLibNotAvailableError, DynamicLibUnknownError, load_nvidia_dynamic_lib from cuda.pathfinder._dynamic_libs import load_nvidia_dynamic_lib as load_nvidia_dynamic_lib_module from cuda.pathfinder._dynamic_libs import supported_nvidia_libs -from cuda.pathfinder._dynamic_libs.subprocess_protocol import STATUS_NOT_FOUND, parse_dynamic_lib_subprocess_payload +from cuda.pathfinder._dynamic_libs.subprocess_protocol import ( + STATUS_NOT_FOUND, + parse_dynamic_lib_subprocess_payload, +) from cuda.pathfinder._utils.platform_aware import IS_WINDOWS, quote_for_shell STRICTNESS = os.environ.get("CUDA_PATHFINDER_TEST_LOAD_NVIDIA_DYNAMIC_LIB_STRICTNESS", "see_what_works") @@ -117,6 +121,7 @@ def raise_child_process_failed(): raise RuntimeError(build_child_process_failed_for_libname_message(libname, result)) if result.returncode != 0: + skip_if_missing_libnvcudla_so(libname, timeout=timeout) raise_child_process_failed() assert not result.stderr payload = parse_dynamic_lib_subprocess_payload( @@ -125,6 +130,7 @@ def raise_child_process_failed(): error_label="Load subprocess child process", ) if payload.status == STATUS_NOT_FOUND: + skip_if_missing_libnvcudla_so(libname, timeout=timeout) if STRICTNESS == "all_must_work" and not _is_expected_load_nvidia_dynamic_lib_failure(libname): raise_child_process_failed() info_summary_append(f"Not found: {libname=!r}")