Handle both Windows conda static-lib layouts. (#2015)

Look for cudadevrt under both Library/lib/x64 and Library/lib so CUDA 12 conda environments resolve the real static library instead of falling through to a misleading CUDA_PATH error.

Made-with: Cursor

Author: rwgk

cuda_pathfinder/cuda/pathfinder/_static_libs/find_static_lib.py

@@ -28,15 +28,15 @@ class LocatedStaticLib:
 class _StaticLibInfo(TypedDict):
     filename: str
     ctk_rel_paths: tuple[str, ...]
-    conda_rel_path: str
+    conda_rel_paths: tuple[str, ...]
     site_packages_dirs: tuple[str, ...]
 
 
 _SUPPORTED_STATIC_LIBS_INFO: dict[str, _StaticLibInfo] = {
     "cudadevrt": {
         "filename": "cudadevrt.lib" if IS_WINDOWS else "libcudadevrt.a",
         "ctk_rel_paths": (os.path.join("lib", "x64"),) if IS_WINDOWS else ("lib64", "lib"),
-        "conda_rel_path": os.path.join("lib", "x64") if IS_WINDOWS else "lib",
+        "conda_rel_paths": ((os.path.join("lib", "x64"), "lib") if IS_WINDOWS else ("lib",)),
         "site_packages_dirs": (
             ("nvidia/cu13/lib/x64", "nvidia/cuda_runtime/lib/x64")
             if IS_WINDOWS
@@ -66,7 +66,7 @@ def __init__(self, name: str) -> None:
         self.config: _StaticLibInfo = _SUPPORTED_STATIC_LIBS_INFO[name]
         self.filename: str = self.config["filename"]
         self.ctk_rel_paths: tuple[str, ...] = self.config["ctk_rel_paths"]
-        self.conda_rel_path: str = self.config["conda_rel_path"]
+        self.conda_rel_paths: tuple[str, ...] = self.config["conda_rel_paths"]
         self.site_packages_dirs: tuple[str, ...] = self.config["site_packages_dirs"]
         self.error_messages: list[str] = []
         self.attachments: list[str] = []
@@ -86,9 +86,10 @@ def try_with_conda_prefix(self) -> str | None:
             return None
 
         anchor = os.path.join(conda_prefix, "Library") if IS_WINDOWS else conda_prefix
-        file_path = os.path.join(anchor, self.conda_rel_path, self.filename)
-        if os.path.isfile(file_path):
-            return file_path
+        for rel_path in self.conda_rel_paths:
+            file_path = os.path.join(anchor, rel_path, self.filename)
+            if os.path.isfile(file_path):
+                return file_path
         return None
 
     def try_with_cuda_home(self) -> str | None:

cuda_pathfinder/tests/test_find_static_lib.py

@@ -78,7 +78,7 @@ def test_locate_static_lib(info_summary_append, libname):
 @pytest.mark.usefixtures("clear_find_static_lib_cache")
 def test_locate_static_lib_search_order(monkeypatch, tmp_path):
     filename = CUDADEVRT_INFO["filename"]
-    conda_rel_path = CUDADEVRT_INFO["conda_rel_path"]
+    conda_rel_path = CUDADEVRT_INFO["conda_rel_paths"][0]
 
     site_pkg_rel = CUDADEVRT_INFO["site_packages_dirs"][0]
     site_packages_lib_dir = tmp_path / "site-packages" / Path(site_pkg_rel.replace("/", os.sep))
@@ -117,6 +117,32 @@ def test_locate_static_lib_search_order(monkeypatch, tmp_path):
     assert located_lib.found_via == "CUDA_PATH"
 
 
+@pytest.mark.usefixtures("clear_find_static_lib_cache")
+def test_locate_static_lib_conda_rel_path_fallback(monkeypatch, tmp_path):
+    filename = CUDADEVRT_INFO["filename"]
+    conda_rel_paths = CUDADEVRT_INFO["conda_rel_paths"]
+    if len(conda_rel_paths) == 1:
+        monkeypatch.setitem(CUDADEVRT_INFO, "conda_rel_paths", ("missing-first", conda_rel_paths[0]))
+        conda_rel_paths = CUDADEVRT_INFO["conda_rel_paths"]
+
+    conda_prefix = tmp_path / "conda-prefix"
+    conda_lib_dir = _conda_anchor(conda_prefix) / Path(conda_rel_paths[1])
+    conda_path = _make_static_lib_file(conda_lib_dir, filename)
+
+    monkeypatch.setattr(
+        find_static_lib_module,
+        "find_sub_dirs_all_sitepackages",
+        lambda _sub_dir: [],
+    )
+    monkeypatch.setenv("CONDA_PREFIX", str(conda_prefix))
+    monkeypatch.delenv("CUDA_HOME", raising=False)
+    monkeypatch.delenv("CUDA_PATH", raising=False)
+
+    located_lib = locate_static_lib("cudadevrt")
+    assert located_lib.abs_path == conda_path
+    assert located_lib.found_via == "conda"
+
+
 @pytest.mark.usefixtures("clear_find_static_lib_cache")
 def test_find_static_lib_not_found_error_includes_cuda_home_directory_listing(monkeypatch, tmp_path):
     filename = CUDADEVRT_INFO["filename"]