[pathfinder] Respect site-packages root priority across wheel layouts

cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binary_utility.py

@@ -8,7 +8,7 @@
 from cuda.pathfinder._binaries import supported_nvidia_binaries, windows_nsight
 from cuda.pathfinder._utils.ctk_root_canary import CTK_ROOT_CANARY_ANCHOR_LIBNAMES
 from cuda.pathfinder._utils.env_vars import get_cuda_path_or_home
-from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages
+from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages_ordered
 from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
 
 
@@ -193,10 +193,8 @@ def find_nvidia_binary_utility(utility_name: str) -> str | None:
 
     # 1. Search in site-packages (NVIDIA wheels)
     candidate_dirs = supported_nvidia_binaries.SITE_PACKAGES_BINDIRS.get(utility_name, ())
-    dirs = []
-
-    for sub_dir in candidate_dirs:
-        dirs.extend(find_sub_dirs_all_sitepackages(sub_dir.split(os.sep)))
+    patterns = tuple(tuple(sub_dir.split(os.sep)) for sub_dir in candidate_dirs)
+    dirs = find_sub_dirs_all_sitepackages_ordered(patterns) if patterns else []
 
     # 2. Search in Conda environment
     if (conda_prefix := os.environ.get("CONDA_PREFIX")) is not None:

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py

@@ -47,7 +47,7 @@ def for_arch(self, target_arch: str) -> tuple[str, ...]:
 
 
 def _ctk_windows_wheel_dirs(cuda13_bin_dir: str, cuda12_dir: str) -> WindowsSearchDirs:
-    """Search CUDA 13 first, with the x64-only CUDA 12 wheel as an x64 fallback."""
+    """Within each site root, prefer CUDA 13, with CUDA 12 as an x64 fallback."""
     cuda13_bin_path = PurePosixPath(cuda13_bin_dir)
     return WindowsSearchDirs(
         x64=((cuda13_bin_path / "x86_64").as_posix(), cuda12_dir),

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py

@@ -18,7 +18,10 @@
 
 from cuda.pathfinder._dynamic_libs.lib_descriptor import LibDescriptor
 from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import is_suppressed_dll_file
-from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages
+from cuda.pathfinder._utils.find_sub_dirs import (
+    find_sub_dirs_all_sitepackages,
+    find_sub_dirs_all_sitepackages_ordered,
+)
 from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
 from cuda.pathfinder._utils.windows_arch import windows_pe_matches_arch, windows_python_arch
 
@@ -48,14 +51,11 @@ def _find_so_in_rel_dirs(
     error_messages: list[str],
     attachments: list[str],
 ) -> str | None:
-    sub_dirs_searched: list[tuple[str, ...]] = []
-    for rel_dir in rel_dirs:
-        sub_dir = PurePath(rel_dir).parts
-        for abs_dir in find_sub_dirs_all_sitepackages(sub_dir):
-            so_path = _find_descriptor_so_under_dir(abs_dir, desc)
-            if so_path is not None:
-                return so_path
-        sub_dirs_searched.append(sub_dir)
+    sub_dirs_searched = tuple(PurePath(rel_dir).parts for rel_dir in rel_dirs)
+    for abs_dir in find_sub_dirs_all_sitepackages_ordered(sub_dirs_searched):
+        so_path = _find_descriptor_so_under_dir(abs_dir, desc)
+        if so_path is not None:
+            return so_path
     for sub_dir in sub_dirs_searched:
         _no_such_file_in_sub_dirs(sub_dir, file_description, error_messages, attachments)
     return None
@@ -88,15 +88,12 @@ def _find_dll_in_rel_dirs(
     error_messages: list[str],
     attachments: list[str],
 ) -> str | None:
-    sub_dirs_searched: list[tuple[str, ...]] = []
+    sub_dirs_searched = tuple(PurePath(rel_dir).parts for rel_dir in rel_dirs)
     checked_arch = target_arch if desc.requires_windows_binary_arch_check else None
-    for rel_dir in rel_dirs:
-        sub_dir = PurePath(rel_dir).parts
-        for abs_dir in find_sub_dirs_all_sitepackages(sub_dir):
-            dll_name = _find_descriptor_dll_under_dir(abs_dir, desc, checked_arch)
-            if dll_name is not None:
-                return dll_name
-        sub_dirs_searched.append(sub_dir)
+    for abs_dir in find_sub_dirs_all_sitepackages_ordered(sub_dirs_searched):
+        dll_name = _find_descriptor_dll_under_dir(abs_dir, desc, checked_arch)
+        if dll_name is not None:
+            return dll_name
     for sub_dir in sub_dirs_searched:
         _no_such_file_in_sub_dirs(sub_dir, lib_searched_for, error_messages, attachments)
     return None

cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py

@@ -22,7 +22,7 @@
 )
 from cuda.pathfinder._utils.ctk_root_canary import CTK_ROOT_CANARY_ANCHOR_LIBNAMES
 from cuda.pathfinder._utils.env_vars import get_cuda_path_or_home
-from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages
+from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages_ordered
 from cuda.pathfinder._utils.path_sort import numeric_aware_path_sort_key
 
 if TYPE_CHECKING:
@@ -81,11 +81,10 @@ def _locate_in_anchor_layout(desc: HeaderDescriptor, anchor_point: str) -> str |
 
 def find_in_site_packages(desc: HeaderDescriptor) -> LocatedHeaderDir | None:
     """Search pip wheel install locations."""
-    for sub_dir in desc.site_packages_dirs:
-        hdr_dir: str  # help mypy
-        for hdr_dir in find_sub_dirs_all_sitepackages(tuple(sub_dir.split("/"))):
-            if _joined_isfile(hdr_dir, desc.header_basename):
-                return LocatedHeaderDir(abs_path=hdr_dir, found_via="site-packages")
+    patterns = tuple(tuple(sub_dir.split("/")) for sub_dir in desc.site_packages_dirs)
+    for hdr_dir in find_sub_dirs_all_sitepackages_ordered(patterns):
+        if _joined_isfile(hdr_dir, desc.header_basename):
+            return LocatedHeaderDir(abs_path=hdr_dir, found_via="site-packages")
     return None
 
 

cuda_pathfinder/cuda/pathfinder/_static_libs/find_bitcode_lib.py

@@ -8,7 +8,7 @@
 from typing import NoReturn, TypedDict
 
 from cuda.pathfinder._utils.env_vars import get_cuda_path_or_home
-from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages
+from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages_ordered
 from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
 
 
@@ -90,12 +90,11 @@ def __init__(self, name: str) -> None:
         self.attachments: list[str] = []
 
     def try_site_packages(self) -> Path | None:
-        for rel_dir in self.site_packages_dirs:
-            sub_dir = tuple(rel_dir.split("/"))
-            for abs_dir in find_sub_dirs_all_sitepackages(sub_dir):
-                file_path = Path(abs_dir, self.filename)
-                if file_path.is_file():
-                    return file_path
+        patterns = tuple(tuple(rel_dir.split("/")) for rel_dir in self.site_packages_dirs)
+        for abs_dir in find_sub_dirs_all_sitepackages_ordered(patterns):
+            file_path = Path(abs_dir, self.filename)
+            if file_path.is_file():
+                return file_path
         return None
 
     def try_with_conda_prefix(self) -> Path | None:

cuda_pathfinder/cuda/pathfinder/_static_libs/find_static_lib.py

@@ -8,7 +8,7 @@
 from typing import NoReturn, TypedDict
 
 from cuda.pathfinder._utils.env_vars import get_cuda_path_or_home
-from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages
+from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages_ordered
 from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
 from cuda.pathfinder._utils.windows_arch import windows_python_arch
 
@@ -87,12 +87,11 @@ def __init__(self, name: str) -> None:
         self.attachments: list[str] = []
 
     def try_site_packages(self) -> Path | None:
-        for rel_dir in self.site_packages_dirs:
-            sub_dir = tuple(rel_dir.split("/"))
-            for abs_dir in find_sub_dirs_all_sitepackages(sub_dir):
-                file_path = Path(abs_dir, self.filename)
-                if file_path.is_file():
-                    return file_path
+        patterns = tuple(tuple(rel_dir.split("/")) for rel_dir in self.site_packages_dirs)
+        for abs_dir in find_sub_dirs_all_sitepackages_ordered(patterns):
+            file_path = Path(abs_dir, self.filename)
+            if file_path.is_file():
+                return file_path
         return None
 
     def try_with_conda_prefix(self) -> Path | None:

cuda_pathfinder/cuda/pathfinder/_utils/find_sub_dirs.py

@@ -49,7 +49,7 @@ def find_sub_dirs_sys_path(sub_dirs: Sequence[str]) -> list[str]:
     return find_sub_dirs(sys.path, sub_dirs)
 
 
-def find_sub_dirs_all_sitepackages(sub_dirs: Sequence[str]) -> list[str]:
+def _sitepackages_in_order() -> list[str]:
     parent_dirs = list(site.getsitepackages())
     if site.ENABLE_USER_SITE:
         user_site = site.getusersitepackages()
@@ -62,4 +62,18 @@ def find_sub_dirs_all_sitepackages(sub_dirs: Sequence[str]) -> list[str]:
             # Detect venv by checking if sys.prefix differs from sys.base_prefix
             insert_idx = 1 if sys.prefix != sys.base_prefix else 0
             parent_dirs.insert(insert_idx, user_site)
-    return find_sub_dirs(parent_dirs, sub_dirs)
+    return parent_dirs
+
+
+def find_sub_dirs_all_sitepackages(sub_dirs: Sequence[str]) -> list[str]:
+    return find_sub_dirs(_sitepackages_in_order(), sub_dirs)
+
+
+def find_sub_dirs_all_sitepackages_ordered(sub_dir_patterns: Sequence[Sequence[str]]) -> list[str]:
+    """Search site roots in priority order, then the wheel layouts within each root."""
+    return [
+        directory
+        for parent in _sitepackages_in_order()
+        for pattern in sub_dir_patterns
+        for directory in find_sub_dirs((parent,), pattern)
+    ]

cuda_pathfinder/tests/test_find_bitcode_lib.py

@@ -105,15 +105,13 @@ def test_locate_bitcode_lib_search_order(monkeypatch, tmp_path, libname):
         tuple(rel_dir.split("/")) for rel_dir in _bitcode_lib_info(libname)["site_packages_dirs"]
     )
 
-    def find_expected_sub_dir(sub_dir):
-        assert sub_dir in site_packages_sub_dirs
-        if sub_dir == site_packages_sub_dirs[0]:
-            return [str(site_packages_lib_dir)]
-        return []
+    def find_expected_sub_dir(patterns):
+        assert patterns == site_packages_sub_dirs
+        return [str(site_packages_lib_dir)]
 
     monkeypatch.setattr(
         find_bitcode_lib_module,
-        "find_sub_dirs_all_sitepackages",
+        "find_sub_dirs_all_sitepackages_ordered",
         find_expected_sub_dir,
     )
     monkeypatch.setenv("CONDA_PREFIX", str(conda_prefix))
@@ -145,7 +143,7 @@ def test_find_bitcode_lib_not_found_error_includes_cuda_home_directory_listing(m
 
     monkeypatch.setattr(
         find_bitcode_lib_module,
-        "find_sub_dirs_all_sitepackages",
+        "find_sub_dirs_all_sitepackages_ordered",
         lambda _sub_dir: [],
     )
     monkeypatch.delenv("CONDA_PREFIX", raising=False)
@@ -166,7 +164,7 @@ def test_find_bitcode_lib_not_found_error_includes_cuda_home_directory_listing(m
 def test_find_bitcode_lib_not_found_error_without_cuda_home(monkeypatch):
     monkeypatch.setattr(
         find_bitcode_lib_module,
-        "find_sub_dirs_all_sitepackages",
+        "find_sub_dirs_all_sitepackages_ordered",
         lambda _sub_dir: [],
     )
     monkeypatch.delenv("CONDA_PREFIX", raising=False)