feat(pathfinder): add CTK-root canary fallback to find_nvidia_binary_utility

aryanputta · aryanputta · commit c3d39f026d3d · 2026-06-10T21:11:39.000-04:00
After the deterministic search over the explicit trusted directories (NVIDIA wheel bin/, CONDA_PREFIX, CUDA_HOME/CUDA_PATH) misses, fall back to a CTK-root canary probe: resolve cudart through the OS dynamic loader, which honors LD_LIBRARY_PATH on Linux and the native DLL search on Windows, derive the CUDA Toolkit root from its absolute path, and search that root's bin layout. This addresses the concern raised on #2196: users who follow the CUDA Linux installation guide set LD_LIBRARY_PATH for libraries and PATH for executables. The bounded finder alone would stop finding the utility for them because PATH is intentionally never consulted. The canary fallback recovers that case through LD_LIBRARY_PATH instead of PATH. LD_LIBRARY_PATH is still an attack vector, but a significantly harder one to exploit than PATH, and the ambient PATH and process CWD remain unused. The canary runs only after the explicit trusted dirs miss, so the common wheel/conda/CUDA_HOME cases never spawn the resolver subprocess. The canary -> CTK-root resolution is factored into a shared resolve_ctk_root_via_canary helper reused by the dynamic-library CTK-root canary flow. Adds tests for the fallback (found, ordering, Windows bin layout, not consulted when found earlier, cached) and for resolve_ctk_root_via_canary. Adds 1.6.0 release notes for the minor version bump.
diff --git a/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binary_utility.py b/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binary_utility.py
@@ -9,6 +9,11 @@
 from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages
 from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
 
+# CUDA Toolkit canary library used to derive the toolkit root when it is only
+# visible through the dynamic loader. ``cudart`` always ships with the CTK and
+# matches the anchor used by the dynamic-library CTK-root canary flow.
+_CTK_ROOT_CANARY_ANCHOR_LIBNAME = "cudart"
+
 
 class UnsupportedBinaryError(Exception):
     def __init__(self, utility: str) -> None:
@@ -41,6 +46,35 @@ def _is_executable_file(path: str) -> bool:
     return os.access(path, os.X_OK)
 
 
+def _ctk_bin_subdirs(root: str) -> list[str]:
+    """Return the bin directories to search under a CUDA Toolkit ``root``.
+
+    On Windows the CTK ships binaries under ``bin/x64`` (CTK 13), ``bin/x86_64``,
+    and ``bin`` (CTK 12); on Linux they live in ``bin``.
+    """
+    if IS_WINDOWS:
+        return [
+            os.path.join(root, "bin", "x64"),
+            os.path.join(root, "bin", "x86_64"),
+            os.path.join(root, "bin"),
+        ]
+    return [os.path.join(root, "bin")]
+
+
+def _resolve_ctk_root_via_canary() -> str | None:
+    """Derive the CUDA Toolkit root from the ``cudart`` canary library.
+
+    ``cudart`` is resolved by the OS dynamic loader, which honors
+    ``LD_LIBRARY_PATH`` on Linux and the native DLL search on Windows, and the
+    toolkit root is derived from its absolute path. The ambient ``PATH`` is
+    never consulted. The loader module is imported lazily to avoid pulling the
+    dynamic-library machinery in at import time.
+    """
+    from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import resolve_ctk_root_via_canary
+
+    return resolve_ctk_root_via_canary(_CTK_ROOT_CANARY_ANCHOR_LIBNAME)
+
+
 def _resolve_in_trusted_dirs(normalized_name: str, dirs: list[str]) -> str | None:
     """Resolve ``normalized_name`` against ``dirs`` only, in order.
 
@@ -99,6 +133,15 @@ def find_nvidia_binary_utility(utility_name: str) -> str | None:
              ``bin/x64``, ``bin/x86_64``, and ``bin`` subdirectories on Windows,
              or just ``bin`` on Linux.
 
+        4. **CTK-root canary fallback**
+
+           - Only when steps 1-3 miss: resolve the ``cudart`` library through the
+             OS dynamic loader (which honors ``LD_LIBRARY_PATH`` on Linux and the
+             native DLL search on Windows), derive the CUDA Toolkit root from it,
+             and search that root's bin layout. This finds the utility for users
+             who follow the CUDA install guide and set ``LD_LIBRARY_PATH`` for
+             libraries without also setting ``CUDA_HOME`` / ``CUDA_PATH``.
+
     Note:
         Results are cached using ``@functools.cache`` for performance. The cache
         persists for the lifetime of the process.
@@ -107,8 +150,9 @@ def find_nvidia_binary_utility(utility_name: str) -> str | None:
         (``.exe``, ``.bat``, ``.cmd``). On Unix-like systems, executables
         are identified by the ``X_OK`` (execute) permission bit.
 
-        Lookup is restricted to the trusted directories listed above; the
-        process working directory and the ambient ``PATH`` are never consulted.
+        Lookup is restricted to the trusted directories and the canary-derived
+        CTK root listed above; the process working directory and the ambient
+        ``PATH`` are never consulted.
 
     Example:
         >>> from cuda.pathfinder import find_nvidia_binary_utility
@@ -135,10 +179,17 @@ def find_nvidia_binary_utility(utility_name: str) -> str | None:
 
     # 3. Search in CUDA Toolkit (CUDA_HOME/CUDA_PATH)
     if (cuda_home := get_cuda_path_or_home()) is not None:
-        if IS_WINDOWS:
-            dirs.append(os.path.join(cuda_home, "bin", "x64"))
-            dirs.append(os.path.join(cuda_home, "bin", "x86_64"))
-        dirs.append(os.path.join(cuda_home, "bin"))
+        dirs.extend(_ctk_bin_subdirs(cuda_home))
 
     normalized_name = _normalize_utility_name(utility_name)
-    return _resolve_in_trusted_dirs(normalized_name, dirs)
+    found = _resolve_in_trusted_dirs(normalized_name, dirs)
+    if found is not None:
+        return found
+
+    # 4. CTK-root canary fallback: only when the explicit trusted dirs above
+    #    miss. Resolve cudart via the dynamic loader (honors LD_LIBRARY_PATH),
+    #    derive the toolkit root, and search its bin layout. PATH is never used.
+    ctk_root = _resolve_ctk_root_via_canary()
+    if ctk_root is not None:
+        return _resolve_in_trusted_dirs(normalized_name, _ctk_bin_subdirs(ctk_root))
+    return None
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
@@ -136,13 +136,25 @@ def _loadable_via_canary_subprocess(libname: str, *, timeout: float = _CANARY_PR
     return _resolve_system_loaded_abs_path_in_subprocess(libname, timeout=timeout) is not None
 
 
+def resolve_ctk_root_via_canary(canary_libname: str) -> str | None:
+    """Resolve the CUDA Toolkit root from a system-loadable canary library.
+
+    The canary library's absolute path is resolved by the OS dynamic loader in
+    an isolated subprocess, which honors ``LD_LIBRARY_PATH`` on Linux and the
+    native DLL search on Windows. The toolkit root is then derived from that
+    path. Returns ``None`` if the canary cannot be resolved or no root can be
+    derived. The ambient ``PATH`` is never consulted.
+    """
+    canary_abs_path = _resolve_system_loaded_abs_path_in_subprocess(canary_libname)
+    if canary_abs_path is None:
+        return None
+    return derive_ctk_root(canary_abs_path)
+
+
 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:
-        canary_abs_path = _resolve_system_loaded_abs_path_in_subprocess(canary_libname)
-        if canary_abs_path is None:
-            continue
-        ctk_root = derive_ctk_root(canary_abs_path)
+        ctk_root = resolve_ctk_root_via_canary(canary_libname)
         if ctk_root is None:
             continue
         find = find_via_ctk_root(ctx, ctk_root)
diff --git a/cuda_pathfinder/docs/source/release/1.6.0-notes.rst b/cuda_pathfinder/docs/source/release/1.6.0-notes.rst
@@ -0,0 +1,27 @@
+.. SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+.. SPDX-License-Identifier: Apache-2.0
+
+.. py:currentmodule:: cuda.pathfinder
+
+``cuda-pathfinder`` 1.6.0 Release notes
+=======================================
+
+Highlights
+----------
+
+* :func:`find_nvidia_binary_utility` now resolves binaries through a bounded,
+  deterministic search of trusted directories instead of ``shutil.which``. The
+  process working directory and the ambient ``PATH`` are never consulted, which
+  closes a lookup ambiguity on Windows where ``shutil.which`` prepends the CWD
+  even when an explicit search path is supplied.
+  (`PR #2196 <https://github.com/NVIDIA/cuda-python/pull/2196>`_)
+
+* :func:`find_nvidia_binary_utility` gains a CTK-root canary fallback. When the
+  NVIDIA wheel, ``CONDA_PREFIX``, and ``CUDA_HOME`` / ``CUDA_PATH`` directories
+  all miss, ``cudart`` is resolved through the OS dynamic loader, which honors
+  ``LD_LIBRARY_PATH`` on Linux and the native DLL search on Windows. The CUDA
+  Toolkit root is derived from that path and its ``bin`` layout is searched.
+  This locates the utility for users who follow the CUDA installation guide and
+  set ``LD_LIBRARY_PATH`` for libraries without also setting ``CUDA_HOME`` /
+  ``CUDA_PATH``, while still never falling back to ``PATH``.
+  (`PR #2196 <https://github.com/NVIDIA/cuda-python/pull/2196>`_)
diff --git a/cuda_pathfinder/tests/test_ctk_root_discovery.py b/cuda_pathfinder/tests/test_ctk_root_discovery.py
@@ -17,6 +17,7 @@
     _load_lib_no_cache,
     _resolve_system_loaded_abs_path_in_subprocess,
     _try_ctk_root_canary,
+    resolve_ctk_root_via_canary,
 )
 from cuda.pathfinder._dynamic_libs.search_steps import (
     SearchContext,
@@ -369,6 +370,35 @@ def test_canary_skips_when_abs_path_none(mocker):
     assert _try_ctk_root_canary(_ctx("nvvm")) is None
 
 
+# ---------------------------------------------------------------------------
+# resolve_ctk_root_via_canary (shared by lib and binary discovery)
+# ---------------------------------------------------------------------------
+
+
+def test_resolve_ctk_root_via_canary_returns_root(tmp_path, mocker):
+    ctk_root = tmp_path / "cuda-13"
+    _create_cudart_in_ctk(ctk_root)
+    probe = mocker.patch(
+        f"{_MODULE}._resolve_system_loaded_abs_path_in_subprocess",
+        return_value=_fake_canary_path(ctk_root),
+    )
+    assert resolve_ctk_root_via_canary("cudart") == str(ctk_root)
+    probe.assert_called_once_with("cudart")
+
+
+def test_resolve_ctk_root_via_canary_none_when_probe_fails(mocker):
+    mocker.patch(f"{_MODULE}._resolve_system_loaded_abs_path_in_subprocess", return_value=None)
+    assert resolve_ctk_root_via_canary("cudart") is None
+
+
+def test_resolve_ctk_root_via_canary_none_when_unrecognized(mocker):
+    mocker.patch(
+        f"{_MODULE}._resolve_system_loaded_abs_path_in_subprocess",
+        return_value=os.path.join(os.sep, "weird", "path", "libcudart.so.13"),
+    )
+    assert resolve_ctk_root_via_canary("cudart") is None
+
+
 # ---------------------------------------------------------------------------
 # _load_lib_no_cache search-order
 # ---------------------------------------------------------------------------
diff --git a/cuda_pathfinder/tests/test_find_nvidia_binaries.py b/cuda_pathfinder/tests/test_find_nvidia_binaries.py
@@ -76,6 +76,7 @@ def test_find_binary_search_path_includes_site_packages_conda_cuda(monkeypatch,
     )
     monkeypatch.setenv("CONDA_PREFIX", conda_prefix)
     mocker.patch.object(binary_finder_module, "get_cuda_path_or_home", return_value=cuda_home)
+    mocker.patch.object(binary_finder_module, "_resolve_ctk_root_via_canary", return_value=None)
     expected_dirs = [
         site_dir,
         os.path.join(conda_prefix, "bin"),
@@ -109,6 +110,7 @@ def test_find_binary_windows_extension_and_search_dirs(monkeypatch, mocker):
     )
     monkeypatch.setenv("CONDA_PREFIX", conda_prefix)
     mocker.patch.object(binary_finder_module, "get_cuda_path_or_home", return_value=cuda_home)
+    mocker.patch.object(binary_finder_module, "_resolve_ctk_root_via_canary", return_value=None)
     expected_dirs = [
         site_dir,
         os.path.join(conda_prefix, "Library", "bin"),
@@ -142,6 +144,7 @@ def test_find_binary_first_matching_dir_wins(monkeypatch, mocker):
     mocker.patch.object(binary_finder_module, "find_sub_dirs_all_sitepackages", return_value=[site_dir])
     monkeypatch.setenv("CONDA_PREFIX", conda_prefix)
     mocker.patch.object(binary_finder_module, "get_cuda_path_or_home", return_value=cuda_home)
+    mocker.patch.object(binary_finder_module, "_resolve_ctk_root_via_canary", return_value=None)
     conda_nvcc = os.path.join(conda_prefix, "bin", "nvcc")
     cuda_nvcc = os.path.join(cuda_home, "bin", "nvcc")
     checked = _patch_exec_probe(mocker, existing=[conda_nvcc, cuda_nvcc])
@@ -153,6 +156,72 @@ def test_find_binary_first_matching_dir_wins(monkeypatch, mocker):
     assert checked == [os.path.join(site_dir, "nvcc"), conda_nvcc]
 
 
+@pytest.mark.usefixtures("clear_find_binary_cache")
+def test_find_binary_ctk_root_canary_fallback(monkeypatch, mocker):
+    # When the explicit trusted dirs (wheels, conda, CUDA_HOME/PATH) all miss,
+    # the cudart-canary-derived CTK root is searched last.
+    ctk_root = os.path.join(os.sep, "opt", "cuda")
+
+    mocker.patch.object(binary_finder_module, "IS_WINDOWS", new=False)
+    mocker.patch.object(binary_finder_module.supported_nvidia_binaries, "SITE_PACKAGES_BINDIRS", {})
+    mocker.patch.object(binary_finder_module, "find_sub_dirs_all_sitepackages", return_value=[])
+    monkeypatch.delenv("CONDA_PREFIX", raising=False)
+    mocker.patch.object(binary_finder_module, "get_cuda_path_or_home", return_value=None)
+    canary_mock = mocker.patch.object(binary_finder_module, "_resolve_ctk_root_via_canary", return_value=ctk_root)
+    ctk_nvcc = os.path.join(ctk_root, "bin", "nvcc")
+    checked = _patch_exec_probe(mocker, existing=[ctk_nvcc])
+
+    result = find_nvidia_binary_utility("nvcc")
+
+    assert result == ctk_nvcc
+    canary_mock.assert_called_once_with()
+    # No earlier trusted dirs existed, so the only probe is the canary bin dir.
+    assert checked == [ctk_nvcc]
+
+
+@pytest.mark.usefixtures("clear_find_binary_cache")
+def test_find_binary_canary_windows_bin_layout(monkeypatch, mocker):
+    ctk_root = os.path.join("C:", os.sep, "cuda")
+
+    mocker.patch.object(binary_finder_module, "IS_WINDOWS", new=True)
+    mocker.patch.object(binary_finder_module.supported_nvidia_binaries, "SITE_PACKAGES_BINDIRS", {})
+    mocker.patch.object(binary_finder_module, "find_sub_dirs_all_sitepackages", return_value=[])
+    monkeypatch.delenv("CONDA_PREFIX", raising=False)
+    mocker.patch.object(binary_finder_module, "get_cuda_path_or_home", return_value=None)
+    mocker.patch.object(binary_finder_module, "_resolve_ctk_root_via_canary", return_value=ctk_root)
+    expected_dirs = [
+        os.path.join(ctk_root, "bin", "x64"),
+        os.path.join(ctk_root, "bin", "x86_64"),
+        os.path.join(ctk_root, "bin"),
+    ]
+    checked = _patch_exec_probe(mocker)
+
+    result = find_nvidia_binary_utility("nvcc")
+
+    assert result is None
+    assert checked == [os.path.join(d, "nvcc.exe") for d in expected_dirs]
+
+
+@pytest.mark.usefixtures("clear_find_binary_cache")
+def test_find_binary_canary_not_consulted_when_found_earlier(monkeypatch, mocker):
+    # An earlier trusted dir hit must short-circuit before the canary subprocess.
+    conda_prefix = os.path.join(os.sep, "conda")
+
+    mocker.patch.object(binary_finder_module, "IS_WINDOWS", new=False)
+    mocker.patch.object(binary_finder_module.supported_nvidia_binaries, "SITE_PACKAGES_BINDIRS", {})
+    mocker.patch.object(binary_finder_module, "find_sub_dirs_all_sitepackages", return_value=[])
+    monkeypatch.setenv("CONDA_PREFIX", conda_prefix)
+    mocker.patch.object(binary_finder_module, "get_cuda_path_or_home", return_value=None)
+    canary_mock = mocker.patch.object(binary_finder_module, "_resolve_ctk_root_via_canary", return_value=None)
+    conda_nvcc = os.path.join(conda_prefix, "bin", "nvcc")
+    _patch_exec_probe(mocker, existing=[conda_nvcc])
+
+    result = find_nvidia_binary_utility("nvcc")
+
+    assert result == conda_nvcc
+    canary_mock.assert_not_called()
+
+
 @pytest.mark.usefixtures("clear_find_binary_cache")
 def test_find_binary_returns_none_with_no_candidates(monkeypatch, mocker):
     site_key = os.path.join("nvidia", "cuda_nvcc", "bin")
@@ -166,6 +235,7 @@ def test_find_binary_returns_none_with_no_candidates(monkeypatch, mocker):
     find_sub_dirs_mock = mocker.patch.object(binary_finder_module, "find_sub_dirs_all_sitepackages", return_value=[])
     monkeypatch.delenv("CONDA_PREFIX", raising=False)
     mocker.patch.object(binary_finder_module, "get_cuda_path_or_home", return_value=None)
+    mocker.patch.object(binary_finder_module, "_resolve_ctk_root_via_canary", return_value=None)
     checked = _patch_exec_probe(mocker)
 
     result = find_nvidia_binary_utility("nvcc")
@@ -186,6 +256,7 @@ def test_find_binary_without_site_packages_entry(monkeypatch, mocker):
     find_sub_dirs_mock = mocker.patch.object(binary_finder_module, "find_sub_dirs_all_sitepackages", return_value=[])
     monkeypatch.setenv("CONDA_PREFIX", conda_prefix)
     mocker.patch.object(binary_finder_module, "get_cuda_path_or_home", return_value=cuda_home)
+    mocker.patch.object(binary_finder_module, "_resolve_ctk_root_via_canary", return_value=None)
     expected_dirs = [
         os.path.join(conda_prefix, "bin"),
         os.path.join(cuda_home, "bin"),
@@ -206,15 +277,18 @@ def test_find_binary_cache_negative_result(monkeypatch, mocker):
     mocker.patch.object(binary_finder_module, "find_sub_dirs_all_sitepackages", return_value=[])
     monkeypatch.delenv("CONDA_PREFIX", raising=False)
     cuda_home_mock = mocker.patch.object(binary_finder_module, "get_cuda_path_or_home", return_value=None)
+    canary_mock = mocker.patch.object(binary_finder_module, "_resolve_ctk_root_via_canary", return_value=None)
     _patch_exec_probe(mocker)
 
     first = find_nvidia_binary_utility("nvcc")
     second = find_nvidia_binary_utility("nvcc")
 
     assert first is None
     assert second is None
-    # The second call is served from @functools.cache, so the body runs only once.
+    # The second call is served from @functools.cache, so the body runs only
+    # once, including the canary fallback.
     cuda_home_mock.assert_called_once_with()
+    canary_mock.assert_called_once_with()
 
 
 class TestResolveInTrustedDirs:
