pathfinder: cache static and bitcode locate results

Move static and bitcode caching to the shared locate layer so strict-mode public APIs reuse the same discovery boundary after process-wide guard-rails indirection. Add symmetric wrapper cache clears and a regression test that exercises the strict-mode path.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: rwgk

cuda_pathfinder/cuda/pathfinder/_process_wide_compatibility_guard_rails.py

@@ -261,6 +261,12 @@ def find_nvidia_binary_utility(utility_name: str) -> str | None:
 locate_nvidia_header_directory.cache_clear = _cache_clear_with_process_state_reset(  # type: ignore[attr-defined]
     _locate_nvidia_header_directory.cache_clear
 )
+_locate_static_lib_cache_clear = _cache_clear_with_process_state_reset(_locate_static_lib.cache_clear)
+locate_static_lib.cache_clear = _locate_static_lib_cache_clear  # type: ignore[attr-defined]
+find_static_lib.cache_clear = _locate_static_lib_cache_clear  # type: ignore[attr-defined]
+_locate_bitcode_lib_cache_clear = _cache_clear_with_process_state_reset(_locate_bitcode_lib.cache_clear)
+locate_bitcode_lib.cache_clear = _locate_bitcode_lib_cache_clear  # type: ignore[attr-defined]
+find_bitcode_lib.cache_clear = _locate_bitcode_lib_cache_clear  # type: ignore[attr-defined]
 find_nvidia_binary_utility.cache_clear = _cache_clear_with_process_state_reset(  # type: ignore[attr-defined]
     _find_nvidia_binary_utility.cache_clear
 )

cuda_pathfinder/cuda/pathfinder/_static_libs/find_bitcode_lib.py

@@ -137,6 +137,7 @@ def raise_not_found_error(self) -> NoReturn:
         raise BitcodeLibNotFoundError(f'Failure finding "{self.filename}": {err}\n{att}')
 
 
+@functools.cache
 def locate_bitcode_lib(name: str) -> LocatedBitcodeLib:
     """Locate a bitcode library by name.
 
@@ -176,7 +177,6 @@ def locate_bitcode_lib(name: str) -> LocatedBitcodeLib:
     finder.raise_not_found_error()
 
 
-@functools.cache
 def find_bitcode_lib(name: str) -> str:
     """Find the absolute path to a bitcode library.
 
@@ -185,3 +185,6 @@ def find_bitcode_lib(name: str) -> str:
         BitcodeLibNotFoundError: If the bitcode library cannot be found.
     """
     return locate_bitcode_lib(name).abs_path
+
+
+find_bitcode_lib.cache_clear = locate_bitcode_lib.cache_clear  # type: ignore[attr-defined]

cuda_pathfinder/cuda/pathfinder/_static_libs/find_static_lib.py

@@ -122,6 +122,7 @@ def raise_not_found_error(self) -> NoReturn:
         raise StaticLibNotFoundError(f'Failure finding "{self.filename}": {err}\n{att}')
 
 
+@functools.cache
 def locate_static_lib(name: str) -> LocatedStaticLib:
     """Locate a static library by name.
 
@@ -161,7 +162,6 @@ def locate_static_lib(name: str) -> LocatedStaticLib:
     finder.raise_not_found_error()
 
 
-@functools.cache
 def find_static_lib(name: str) -> str:
     """Find the absolute path to a static library.
 
@@ -170,3 +170,6 @@ def find_static_lib(name: str) -> str:
         StaticLibNotFoundError: If the static library cannot be found.
     """
     return locate_static_lib(name).abs_path
+
+
+find_static_lib.cache_clear = locate_static_lib.cache_clear  # type: ignore[attr-defined]

cuda_pathfinder/tests/test_compatibility_guard_rails_public.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import sys
+from pathlib import Path
 
 import pytest
 from compatibility_guard_rails_test_utils import (
@@ -21,6 +22,8 @@
     process_wide_module,
 )
 
+import cuda.pathfinder._static_libs.find_bitcode_lib as find_bitcode_lib_module
+import cuda.pathfinder._static_libs.find_static_lib as find_static_lib_module
 from cuda import pathfinder
 from cuda.pathfinder import (
     CompatibilityCheckError,
@@ -30,6 +33,7 @@
     LocatedHeaderDir,
     process_wide_compatibility_guard_rails,
 )
+from cuda.pathfinder._utils.env_vars import get_cuda_path_or_home
 
 
 def test_process_wide_compatibility_guard_rails_is_public_singleton():
@@ -313,3 +317,76 @@ def test_public_apis_share_process_wide_guard_rails_state(monkeypatch, tmp_path)
     assert loaded.abs_path == lib_path
     with pytest.raises(CompatibilityCheckError, match=r"companion tag 'api_nvrtc'"):
         pathfinder.find_nvidia_header_directory("nvrtc")
+
+
+@pytest.mark.parametrize(
+    ("public_find_name", "public_locate_name", "raw_module", "finder_class_name", "name", "relative_path"),
+    [
+        (
+            "find_static_lib",
+            "locate_static_lib",
+            find_static_lib_module,
+            "_FindStaticLib",
+            "cudadevrt",
+            Path(find_static_lib_module._SUPPORTED_STATIC_LIBS_INFO["cudadevrt"]["ctk_rel_paths"][0])
+            / find_static_lib_module._SUPPORTED_STATIC_LIBS_INFO["cudadevrt"]["filename"],
+        ),
+        (
+            "find_bitcode_lib",
+            "locate_bitcode_lib",
+            find_bitcode_lib_module,
+            "_FindBitcodeLib",
+            "device",
+            Path(find_bitcode_lib_module._SUPPORTED_BITCODE_LIBS_INFO["device"]["rel_path"])
+            / find_bitcode_lib_module._SUPPORTED_BITCODE_LIBS_INFO["device"]["filename"],
+        ),
+    ],
+)
+def test_public_strict_mode_static_and_bitcode_reuse_cached_locate_path(
+    monkeypatch,
+    tmp_path,
+    public_find_name,
+    public_locate_name,
+    raw_module,
+    finder_class_name,
+    name,
+    relative_path,
+):
+    ctk_root = tmp_path / "cuda-12.9"
+    abs_path = _touch_ctk_file(ctk_root, "12.9.20250531", relative_path)
+    finder_class = getattr(raw_module, finder_class_name)
+    original_try_with_cuda_home = finder_class.try_with_cuda_home
+    try_with_cuda_home_calls: list[str] = []
+
+    def counting_try_with_cuda_home(self):
+        try_with_cuda_home_calls.append(self.name)
+        return original_try_with_cuda_home(self)
+
+    monkeypatch.setattr(raw_module, "find_sub_dirs_all_sitepackages", lambda _sub_dir: [])
+    monkeypatch.setattr(finder_class, "try_with_cuda_home", counting_try_with_cuda_home)
+    monkeypatch.delenv("CONDA_PREFIX", raising=False)
+    monkeypatch.delenv("CUDA_PATH", raising=False)
+    monkeypatch.setenv("CUDA_HOME", str(ctk_root))
+    monkeypatch.setattr(
+        pathfinder,
+        "process_wide_compatibility_guard_rails",
+        CompatibilityGuardRails(driver_cuda_version=_driver_cuda_version(13000)),
+    )
+
+    public_find = getattr(pathfinder, public_find_name)
+    public_locate = getattr(pathfinder, public_locate_name)
+
+    public_locate.cache_clear()
+    get_cuda_path_or_home.cache_clear()
+    try:
+        assert public_find(name) == abs_path
+        assert public_find(name) == abs_path
+        assert try_with_cuda_home_calls == [name]
+
+        public_find.cache_clear()
+        assert public_locate(name).abs_path == abs_path
+        assert try_with_cuda_home_calls == [name, name]
+    finally:
+        public_find.cache_clear()
+        public_locate.cache_clear()
+        get_cuda_path_or_home.cache_clear()

cuda_pathfinder/tests/test_find_bitcode_lib.py

@@ -29,10 +29,10 @@ def _bitcode_lib_filename(libname: str) -> str:
 
 @pytest.fixture
 def clear_find_bitcode_lib_cache():
-    find_bitcode_lib_module.find_bitcode_lib.cache_clear()
+    find_bitcode_lib_module.locate_bitcode_lib.cache_clear()
     get_cuda_path_or_home.cache_clear()
     yield
-    find_bitcode_lib_module.find_bitcode_lib.cache_clear()
+    find_bitcode_lib_module.locate_bitcode_lib.cache_clear()
     get_cuda_path_or_home.cache_clear()
 
 
@@ -124,11 +124,13 @@ def find_expected_sub_dir(sub_dir):
     assert located_lib.abs_path == site_packages_path
     assert located_lib.found_via == "site-packages"
     os.remove(site_packages_path)
+    find_bitcode_lib_module.locate_bitcode_lib.cache_clear()
 
     located_lib = locate_bitcode_lib(libname)
     assert located_lib.abs_path == conda_path
     assert located_lib.found_via == "conda"
     os.remove(conda_path)
+    find_bitcode_lib_module.locate_bitcode_lib.cache_clear()
 
     located_lib = locate_bitcode_lib(libname)
     assert located_lib.abs_path == cuda_home_path

cuda_pathfinder/tests/test_find_static_lib.py

@@ -24,10 +24,10 @@
 
 @pytest.fixture
 def clear_find_static_lib_cache():
-    find_static_lib_module.find_static_lib.cache_clear()
+    find_static_lib_module.locate_static_lib.cache_clear()
     get_cuda_path_or_home.cache_clear()
     yield
-    find_static_lib_module.find_static_lib.cache_clear()
+    find_static_lib_module.locate_static_lib.cache_clear()
     get_cuda_path_or_home.cache_clear()
 
 
@@ -106,11 +106,13 @@ def test_locate_static_lib_search_order(monkeypatch, tmp_path):
     assert located_lib.abs_path == site_packages_path
     assert located_lib.found_via == "site-packages"
     os.remove(site_packages_path)
+    find_static_lib_module.locate_static_lib.cache_clear()
 
     located_lib = locate_static_lib("cudadevrt")
     assert located_lib.abs_path == conda_path
     assert located_lib.found_via == "conda"
     os.remove(conda_path)
+    find_static_lib_module.locate_static_lib.cache_clear()
 
     located_lib = locate_static_lib("cudadevrt")
     assert located_lib.abs_path == cuda_home_path