|
1 | 1 | # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
2 | 2 | # SPDX-License-Identifier: Apache-2.0 |
3 | 3 |
|
| 4 | +from pathlib import Path |
| 5 | + |
4 | 6 | import pytest |
5 | 7 |
|
6 | 8 | from cuda.pathfinder._dynamic_libs import load_nvidia_dynamic_lib as load_mod |
@@ -163,10 +165,69 @@ def _run_find_steps_disabled(ctx, steps): |
163 | 165 |
|
164 | 166 |
|
165 | 167 | # --------------------------------------------------------------------------- |
166 | | -# Search order tests (Conda-specific, since Conda is not covered by real CI) |
| 168 | +# Search order tests |
167 | 169 | # --------------------------------------------------------------------------- |
168 | 170 |
|
169 | 171 |
|
| 172 | +@pytest.mark.parametrize("preferred_source", ("site-packages", "conda", "system-search")) |
| 173 | +@pytest.mark.agent_authored(model="unknown") |
| 174 | +def test_cuda_path_is_a_fallback_to_environment_and_native_search(tmp_path, mocker, monkeypatch, preferred_source): |
| 175 | + desc = load_mod.LIB_DESCRIPTORS["cublasLt"] |
| 176 | + platform = steps_mod.PLATFORM |
| 177 | + filename = desc.windows_dlls[0] if IS_WINDOWS else desc.linux_sonames[0] |
| 178 | + |
| 179 | + def create_library(root, relative_dir): |
| 180 | + path = root / relative_dir / filename |
| 181 | + path.parent.mkdir(parents=True, exist_ok=True) |
| 182 | + path.touch() |
| 183 | + return path |
| 184 | + |
| 185 | + cuda_root = tmp_path / "toolkit" |
| 186 | + create_library(cuda_root, platform.anchor_rel_dirs(desc)[0]) |
| 187 | + mocker.patch(f"{_STEPS_MODULE}.get_cuda_path_or_home", return_value=str(cuda_root)) |
| 188 | + |
| 189 | + conda_root = tmp_path / "conda" |
| 190 | + conda_anchor = Path(platform.conda_anchor_point(str(conda_root))) |
| 191 | + conda_path = create_library(conda_anchor, platform.anchor_rel_dirs(desc)[0]) |
| 192 | + if preferred_source == "system-search": |
| 193 | + monkeypatch.delenv("CONDA_PREFIX", raising=False) |
| 194 | + else: |
| 195 | + monkeypatch.setenv("CONDA_PREFIX", str(conda_root)) |
| 196 | + |
| 197 | + wheel_root = tmp_path / "wheels" |
| 198 | + wheel_path = create_library(wheel_root, platform.site_packages_rel_dirs(desc)[0]) |
| 199 | + |
| 200 | + def find_wheel_dirs(sub_dirs): |
| 201 | + directory = wheel_root.joinpath(*sub_dirs) |
| 202 | + return [str(directory)] if preferred_source == "site-packages" and directory.is_dir() else [] |
| 203 | + |
| 204 | + mocker.patch( |
| 205 | + "cuda.pathfinder._dynamic_libs.search_platform.find_sub_dirs_all_sitepackages", |
| 206 | + side_effect=find_wheel_dirs, |
| 207 | + ) |
| 208 | + mocker.patch.object(load_mod.LOADER, "check_if_already_loaded_from_elsewhere", return_value=None) |
| 209 | + mocker.patch(f"{_MODULE}.load_dependencies") |
| 210 | + native_path = tmp_path / "native" / filename |
| 211 | + native_search = mocker.patch.object( |
| 212 | + load_mod.LOADER, |
| 213 | + "load_with_system_search", |
| 214 | + return_value=_make_loaded_dl(str(native_path), "system-search"), |
| 215 | + ) |
| 216 | + mocker.patch.object( |
| 217 | + load_mod.LOADER, |
| 218 | + "load_with_abs_path", |
| 219 | + side_effect=lambda _desc, path, via: _make_loaded_dl(path, via), |
| 220 | + ) |
| 221 | + |
| 222 | + loaded = _load_lib_no_cache("cublasLt") |
| 223 | + |
| 224 | + expected_path = {"site-packages": wheel_path, "conda": conda_path, "system-search": native_path}[preferred_source] |
| 225 | + assert loaded.found_via == preferred_source |
| 226 | + assert loaded.abs_path == str(expected_path) |
| 227 | + if preferred_source != "system-search": |
| 228 | + native_search.assert_not_called() |
| 229 | + |
| 230 | + |
170 | 231 | def test_cupti_search_order_conda_before_cuda_home(tmp_path, mocker, monkeypatch): |
171 | 232 | """Test that conda is searched before CUDA_HOME (CTK). |
172 | 233 |
|
|
0 commit comments