Skip to content

Commit 3a7222b

Browse files
authored
Add support for pathfinder.load_nvidia_dynamic_lib("cupti") (#1693)
* Add Linux support for loading libcupti.so.12 and libcupti.so.13 This commit adds support for finding and loading CUPTI libraries on Linux through cuda.pathfinder. It implements support for all enumerated installation methods: - Site-packages: nvidia/cuda_cupti/lib (CUDA 12) and nvidia/cu13/lib (CUDA 13) - Conda: $CONDA_PREFIX/lib (colocated with other CUDA libraries) - CTK via CUDA_HOME: $CUDA_HOME/extras/CUPTI/lib64 - CTK via canary probe: system CTK root discovery (similar to nvvm) Changes: - Add 'cupti' to supported library names and SONAMEs - Add site-packages paths for CUDA 12 and 13 - Add cupti to CTK root canary discoverable libraries - Update find_nvidia_dynamic_lib to handle extras/CUPTI/lib64 path - Add logic to distinguish CTK (extras/CUPTI/lib64) vs conda (lib) paths - Update _find_so_using_lib_dir to support versioned libraries via glob - Add comprehensive mock tests covering all installation methods Fixes #1572 (Linux support) Made-with: Cursor * Update cupti tests to use new SearchContext-based API Migrated test_load_nvidia_dynamic_lib_using_mocker.py from the old _FindNvidiaDynamicLib API to the new descriptor-based SearchContext API. Changes: - Replace _FindNvidiaDynamicLib imports with search_steps and load_nvidia_dynamic_lib modules - Update mocks to use run_find_steps, LOADER, and SearchContext - Use LIB_DESCRIPTORS to get cupti descriptor - Update all test functions to work with the new search step architecture Made-with: Cursor * Remove unused CTK canary variables from supported_nvidia_libs.py These variables (_CTK_ROOT_CANARY_ANCHOR_LIBNAMES and _CTK_ROOT_CANARY_DISCOVERABLE_LIBNAMES) were added in the cupti PR but are not used in the new descriptor-based architecture. The new code uses desc.ctk_root_canary_anchor_libnames directly from descriptors. Made-with: Cursor * Improve comment for change in LinuxSearchPlatform.find_in_lib_dir() * Add cputi to cu12, cu13 groups in cuda_pathfinder/pyproject.toml * Add cuda_cupti to cuda-components in .github/actions/fetch_ctk/action.yml * Add windows_dlls, site_packages_windows, anchor_rel_dirs_windows for cupti in /descriptor_catalog.py * test: Refactor cupti mock tests to focus on Conda and error paths Remove tests covered by real CI: - Site-packages tests (CUDA 12 and 13) - covered by real CI - CTK tests (CUDA_HOME and canary probe) - covered by real CI - Search order tests involving site-packages/CTK - covered by real CI Keep tests not covered by real CI: - Conda discovery test - Conda not covered by real CI - Error path test (not found) - error path not covered - Conda vs CTK search order test - Conda not covered by real CI Also remove unused imports and helper functions. Made-with: Cursor * Add reverse=True to glob sorting
1 parent ef9253b commit 3a7222b

File tree

5 files changed

+210
-4
lines changed

5 files changed

+210
-4
lines changed

.github/actions/fetch_ctk/action.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ inputs:
1414
cuda-components:
1515
description: "A list of the CTK components to install as a comma-separated list. e.g. 'cuda_nvcc,cuda_nvrtc,cuda_cudart'"
1616
required: false
17-
default: "cuda_nvcc,cuda_cudart,cuda_crt,libnvvm,cuda_nvrtc,cuda_profiler_api,cuda_cccl,libnvjitlink,libcufile,libnvfatbin"
17+
default: "cuda_nvcc,cuda_cudart,cuda_crt,libnvvm,cuda_nvrtc,cuda_profiler_api,cuda_cccl,cuda_cupti,libnvjitlink,libcufile,libnvfatbin"
1818
cuda-path:
1919
description: "where the CTK components will be installed to, relative to $PWD"
2020
required: false

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -266,6 +266,29 @@ class DescriptorSpec:
266266
linux_sonames=("libcufile.so.0",),
267267
site_packages_linux=("nvidia/cu13/lib", "nvidia/cufile/lib"),
268268
),
269+
DescriptorSpec(
270+
name="cupti",
271+
packaged_with="ctk",
272+
linux_sonames=("libcupti.so.12", "libcupti.so.13"),
273+
windows_dlls=(
274+
"cupti64_2025.4.1.dll",
275+
"cupti64_2025.3.1.dll",
276+
"cupti64_2025.2.1.dll",
277+
"cupti64_2025.1.1.dll",
278+
"cupti64_2024.3.2.dll",
279+
"cupti64_2024.2.1.dll",
280+
"cupti64_2024.1.1.dll",
281+
"cupti64_2023.3.1.dll",
282+
"cupti64_2023.2.2.dll",
283+
"cupti64_2023.1.1.dll",
284+
"cupti64_2022.4.1.dll",
285+
),
286+
site_packages_linux=("nvidia/cu13/lib", "nvidia/cuda_cupti/lib"),
287+
site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cuda_cupti/bin"),
288+
anchor_rel_dirs_linux=("extras/CUPTI/lib64", "lib"),
289+
anchor_rel_dirs_windows=("extras/CUPTI/lib64", "bin"),
290+
ctk_root_canary_anchor_libnames=("cudart",),
291+
),
269292
# -----------------------------------------------------------------------
270293
# Third-party / separately packaged libraries
271294
# -----------------------------------------------------------------------

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -141,10 +141,20 @@ def find_in_lib_dir(
141141
error_messages: list[str],
142142
attachments: list[str],
143143
) -> str | None:
144+
# Most libraries have both unversioned and versioned files/symlinks (exact match first)
144145
so_name = os.path.join(lib_dir, lib_searched_for)
145146
if os.path.isfile(so_name):
146147
return so_name
147-
error_messages.append(f"No such file: {so_name}")
148+
# Some libraries only exist as versioned files (e.g., libcupti.so.13 in conda),
149+
# so the glob fallback is needed
150+
file_wild = lib_searched_for + "*"
151+
# Only one match is expected, but to ensure deterministic behavior in unexpected
152+
# situations, and to be internally consistent, we sort in reverse order with the
153+
# intent to return the newest version first.
154+
for so_name in sorted(glob.glob(os.path.join(lib_dir, file_wild)), reverse=True):
155+
if os.path.isfile(so_name):
156+
return so_name
157+
error_messages.append(f"No such file: {file_wild}")
148158
attachments.append(f' listdir("{lib_dir}"):')
149159
if not os.path.isdir(lib_dir):
150160
attachments.append(" DIRECTORY DOES NOT EXIST")

cuda_pathfinder/pyproject.toml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ test = [
1919
]
2020
# Internal organization of test dependencies.
2121
cu12 = [
22-
"cuda-toolkit[nvcc,cublas,nvrtc,cudart,cufft,curand,cusolver,cusparse,npp,nvfatbin,nvjitlink,nvjpeg,cccl]==12.*",
22+
"cuda-toolkit[nvcc,cublas,nvrtc,cudart,cufft,curand,cusolver,cusparse,npp,nvfatbin,nvjitlink,nvjpeg,cccl,cupti]==12.*",
2323
"cuda-toolkit[cufile]==12.*; sys_platform != 'win32'",
2424
"cutensor-cu12",
2525
"nvidia-cublasmp-cu12; sys_platform != 'win32'",
@@ -31,7 +31,7 @@ cu12 = [
3131
"nvidia-nvshmem-cu12; sys_platform != 'win32'",
3232
]
3333
cu13 = [
34-
"cuda-toolkit[nvcc,cublas,nvrtc,cudart,cufft,curand,cusolver,cusparse,npp,nvfatbin,nvjitlink,nvjpeg,cccl,nvvm]==13.*",
34+
"cuda-toolkit[nvcc,cublas,nvrtc,cudart,cufft,curand,cusolver,cusparse,npp,nvfatbin,nvjitlink,nvjpeg,cccl,cupti,nvvm]==13.*",
3535
"cuda-toolkit[cufile]==13.*; sys_platform != 'win32'",
3636
"cutensor-cu13",
3737
"nvidia-cublasmp-cu13; sys_platform != 'win32'",
Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,173 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
import pytest
5+
6+
from cuda.pathfinder._dynamic_libs import load_nvidia_dynamic_lib as load_mod
7+
from cuda.pathfinder._dynamic_libs import search_steps as steps_mod
8+
from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError, LoadedDL
9+
from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import (
10+
_load_lib_no_cache,
11+
_resolve_system_loaded_abs_path_in_subprocess,
12+
)
13+
from cuda.pathfinder._dynamic_libs.search_steps import EARLY_FIND_STEPS
14+
from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
15+
16+
_MODULE = "cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib"
17+
_STEPS_MODULE = "cuda.pathfinder._dynamic_libs.search_steps"
18+
19+
20+
@pytest.fixture(autouse=True)
21+
def _clear_canary_subprocess_probe_cache():
22+
_resolve_system_loaded_abs_path_in_subprocess.cache_clear()
23+
yield
24+
_resolve_system_loaded_abs_path_in_subprocess.cache_clear()
25+
26+
27+
def _make_loaded_dl(path, found_via):
28+
return LoadedDL(path, False, 0xDEAD, found_via)
29+
30+
31+
def _create_cupti_in_ctk(ctk_root):
32+
"""Create a fake cupti lib in extras/CUPTI/lib64."""
33+
if IS_WINDOWS:
34+
cupti_dir = ctk_root / "extras" / "CUPTI" / "lib64"
35+
cupti_dir.mkdir(parents=True, exist_ok=True)
36+
cupti_lib = cupti_dir / "cupti64_2025.4.1.dll"
37+
else:
38+
cupti_dir = ctk_root / "extras" / "CUPTI" / "lib64"
39+
cupti_dir.mkdir(parents=True, exist_ok=True)
40+
cupti_lib = cupti_dir / "libcupti.so.13"
41+
# Create symlink like real CTK installations
42+
cupti_symlink = cupti_dir / "libcupti.so"
43+
cupti_symlink.symlink_to("libcupti.so.13")
44+
cupti_lib.write_bytes(b"fake")
45+
return cupti_lib
46+
47+
48+
# ---------------------------------------------------------------------------
49+
# Conda tests
50+
# Note: Site-packages and CTK are covered by real CI tests.
51+
# Mock tests focus on Conda (not covered by real CI) and error paths.
52+
# ---------------------------------------------------------------------------
53+
54+
55+
def test_cupti_found_in_conda(tmp_path, mocker, monkeypatch):
56+
"""Test finding cupti in conda environment."""
57+
if IS_WINDOWS:
58+
pytest.skip("Windows support for cupti not yet implemented")
59+
60+
# Create conda structure
61+
conda_prefix = tmp_path / "conda_env"
62+
conda_lib_dir = conda_prefix / "lib"
63+
conda_lib_dir.mkdir(parents=True)
64+
cupti_lib = conda_lib_dir / "libcupti.so.13"
65+
cupti_lib.write_bytes(b"fake")
66+
67+
# Mock conda discovery
68+
monkeypatch.setenv("CONDA_PREFIX", str(conda_prefix))
69+
70+
# Disable site-packages search
71+
def _run_find_steps_without_site_packages(ctx, steps):
72+
if steps is EARLY_FIND_STEPS:
73+
# Skip site-packages, only run conda
74+
from cuda.pathfinder._dynamic_libs.search_steps import find_in_conda
75+
76+
result = find_in_conda(ctx)
77+
return result
78+
return steps_mod.run_find_steps(ctx, steps)
79+
80+
mocker.patch(f"{_MODULE}.run_find_steps", side_effect=_run_find_steps_without_site_packages)
81+
mocker.patch.object(load_mod.LOADER, "check_if_already_loaded_from_elsewhere", return_value=None)
82+
mocker.patch(f"{_MODULE}.load_dependencies")
83+
mocker.patch.object(load_mod.LOADER, "load_with_system_search", return_value=None)
84+
mocker.patch(f"{_STEPS_MODULE}.get_cuda_home_or_path", return_value=None)
85+
mocker.patch(f"{_MODULE}._resolve_system_loaded_abs_path_in_subprocess", return_value=None)
86+
mocker.patch.object(
87+
load_mod.LOADER,
88+
"load_with_abs_path",
89+
side_effect=lambda _desc, path, via: _make_loaded_dl(path, via),
90+
)
91+
92+
result = _load_lib_no_cache("cupti")
93+
assert result.found_via == "conda"
94+
assert result.abs_path == str(cupti_lib)
95+
96+
97+
# ---------------------------------------------------------------------------
98+
# Error path tests
99+
# ---------------------------------------------------------------------------
100+
101+
102+
def test_cupti_not_found_raises_error(mocker):
103+
"""Test that DynamicLibNotFoundError is raised when cupti is not found."""
104+
if IS_WINDOWS:
105+
pytest.skip("Windows support for cupti not yet implemented")
106+
107+
# Mock all search paths to return None
108+
def _run_find_steps_disabled(ctx, steps):
109+
return None
110+
111+
mocker.patch(f"{_MODULE}.run_find_steps", side_effect=_run_find_steps_disabled)
112+
mocker.patch.object(load_mod.LOADER, "check_if_already_loaded_from_elsewhere", return_value=None)
113+
mocker.patch(f"{_MODULE}.load_dependencies")
114+
mocker.patch.object(load_mod.LOADER, "load_with_system_search", return_value=None)
115+
mocker.patch(f"{_STEPS_MODULE}.get_cuda_home_or_path", return_value=None)
116+
mocker.patch(
117+
f"{_MODULE}._resolve_system_loaded_abs_path_in_subprocess",
118+
return_value=None,
119+
)
120+
121+
with pytest.raises(DynamicLibNotFoundError):
122+
_load_lib_no_cache("cupti")
123+
124+
125+
# ---------------------------------------------------------------------------
126+
# Search order tests (Conda-specific, since Conda is not covered by real CI)
127+
# ---------------------------------------------------------------------------
128+
129+
130+
def test_cupti_search_order_conda_before_cuda_home(tmp_path, mocker, monkeypatch):
131+
"""Test that conda is searched before CUDA_HOME (CTK).
132+
133+
This test is important because Conda is not covered by real CI tests,
134+
so we need to verify the search order between Conda and CTK.
135+
"""
136+
if IS_WINDOWS:
137+
pytest.skip("Windows support for cupti not yet implemented")
138+
139+
# Create both conda and CUDA_HOME structures
140+
conda_prefix = tmp_path / "conda_env"
141+
conda_lib_dir = conda_prefix / "lib"
142+
conda_lib_dir.mkdir(parents=True)
143+
conda_cupti_lib = conda_lib_dir / "libcupti.so.13"
144+
conda_cupti_lib.write_bytes(b"fake")
145+
146+
ctk_root = tmp_path / "cuda-13.1"
147+
_create_cupti_in_ctk(ctk_root)
148+
149+
# Mock discovery - disable site-packages, enable conda
150+
def _run_find_steps_without_site_packages(ctx, steps):
151+
if steps is EARLY_FIND_STEPS:
152+
# Skip site-packages, only run conda
153+
from cuda.pathfinder._dynamic_libs.search_steps import find_in_conda
154+
155+
result = find_in_conda(ctx)
156+
return result
157+
return steps_mod.run_find_steps(ctx, steps)
158+
159+
mocker.patch(f"{_MODULE}.run_find_steps", side_effect=_run_find_steps_without_site_packages)
160+
monkeypatch.setenv("CONDA_PREFIX", str(conda_prefix))
161+
mocker.patch.object(load_mod.LOADER, "check_if_already_loaded_from_elsewhere", return_value=None)
162+
mocker.patch(f"{_MODULE}.load_dependencies")
163+
mocker.patch.object(load_mod.LOADER, "load_with_system_search", return_value=None)
164+
mocker.patch(f"{_STEPS_MODULE}.get_cuda_home_or_path", return_value=str(ctk_root))
165+
mocker.patch.object(
166+
load_mod.LOADER,
167+
"load_with_abs_path",
168+
side_effect=lambda _desc, path, via: _make_loaded_dl(path, via),
169+
)
170+
171+
result = _load_lib_no_cache("cupti")
172+
assert result.found_via == "conda"
173+
assert result.abs_path == str(conda_cupti_lib)

0 commit comments

Comments
 (0)