Add LoadedDL.found_via

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/find_nvidia_dynamic_lib.py

@@ -7,6 +7,7 @@
 from typing import Optional
 
 from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError
+from cuda.pathfinder._dynamic_libs.load_dl_linux import get_candidate_sonames, _load_lib, abs_path_for_dynamic_library
 from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import (
     SITE_PACKAGES_LIBDIRS_LINUX,
     SITE_PACKAGES_LIBDIRS_WINDOWS,
@@ -16,6 +17,10 @@
 from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages
 from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
 
+if IS_WINDOWS:
+  from cuda.pathfinder._dynamic_libs.load_dl_windows import find_with_system_search_windows as find_with_system_search  
+else:
+    from cuda.pathfinder._dynamic_libs.load_dl_linux import find_with_system_search_linux as find_with_system_search
 
 def _no_such_file_in_sub_dirs(
     sub_dirs: Sequence[str], file_wild: str, error_messages: list[str], attachments: list[str]
@@ -163,6 +168,7 @@ def __init__(self, libname: str):
         self.error_messages: list[str] = []
         self.attachments: list[str] = []
         self.abs_path: Optional[str] = None
+        self.distribution: Optional[str] = None
 
     def try_site_packages(self) -> Optional[str]:
         if IS_WINDOWS:
@@ -186,6 +192,10 @@ def try_with_conda_prefix(self) -> Optional[str]:
     def try_with_cuda_home(self) -> Optional[str]:
         return self._find_using_lib_dir(_find_lib_dir_using_cuda_home(self.libname))
 
+    def try_with_system_search(self) -> Optional[str]:
+        return find_with_system_search(self.libname)
+
+
     def _find_using_lib_dir(self, lib_dir: Optional[str]) -> Optional[str]:
         if lib_dir is None:
             return None
@@ -203,6 +213,8 @@ def _find_using_lib_dir(self, lib_dir: Optional[str]) -> Optional[str]:
                 self.error_messages,
                 self.attachments,
             )
+        if self.abs_path is not None:
+            self.distribution = "CUDA_HOME"
 
     def raise_not_found_error(self) -> None:
         err = ", ".join(self.error_messages)

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_common.py

@@ -10,13 +10,17 @@
 class DynamicLibNotFoundError(RuntimeError):
     pass
 
+@dataclass
+class Distribution:
+    name: str
+    version: str
 
 @dataclass
 class LoadedDL:
     abs_path: Optional[str]
     was_already_loaded_from_elsewhere: bool
     _handle_uint: int  # Platform-agnostic unsigned pointer value
-
+    distribution: Optional[Distribution] = None
 
 def load_dependencies(libname: str, load_func: Callable[[str], LoadedDL]) -> None:
     for dep in DIRECT_DEPENDENCIES.get(libname, ()):

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_linux.py

@@ -7,7 +7,7 @@
 import os
 from typing import Optional, cast
 
-from cuda.pathfinder._dynamic_libs.load_dl_common import LoadedDL
+from cuda.pathfinder._dynamic_libs.load_dl_common import LoadedDL, Distribution
 from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import (
     LIBNAMES_REQUIRING_RTLD_DEEPBIND,
     SUPPORTED_LINUX_SONAMES,
@@ -170,9 +170,21 @@ def load_with_system_search(libname: str) -> Optional[LoadedDL]:
             abs_path = abs_path_for_dynamic_library(libname, handle)
             if abs_path is None:
                 raise RuntimeError(f"No expected symbol for {libname=!r}")
-            return LoadedDL(abs_path, False, handle._handle)
+            return LoadedDL(abs_path, False, handle._handle, Distribution("system", None))
     return None
 
+def find_with_system_search_linux(libname: str) -> Optional[str]:
+    for soname in get_candidate_sonames(libname):
+        try:
+            handle = _load_lib(libname, soname)
+        except OSError:
+            pass
+        else:
+            abs_path = abs_path_for_dynamic_library(libname, handle)
+            if abs_path is None:
+                raise RuntimeError(f"No expected symbol for {libname=!r}")
+            return abs_path
+    return None
 
 def _work_around_known_bugs(libname: str, found_path: str) -> None:
     if libname == "nvrtc":
@@ -194,21 +206,11 @@ def _work_around_known_bugs(libname: str, found_path: str) -> None:
 
 
 def load_with_abs_path(libname: str, found_path: str) -> LoadedDL:
-    """Load a dynamic library from the given path.
 
-    Args:
-        libname: The name of the library to load
-        found_path: The absolute path to the library file
-
-    Returns:
-        A LoadedDL object representing the loaded library
-
-    Raises:
-        RuntimeError: If the library cannot be loaded
-    """
     _work_around_known_bugs(libname, found_path)
     try:
         handle = _load_lib(libname, found_path)
     except OSError as e:
         raise RuntimeError(f"Failed to dlopen {found_path}: {e}") from e
     return LoadedDL(found_path, False, handle._handle)
+

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_windows.py

@@ -132,6 +132,24 @@ def load_with_system_search(libname: str) -> Optional[LoadedDL]:
 
     return None
 
+def find_with_system_search_windows(libname: str) -> Optional[str]:
+    """Try to find a DLL using system search paths.
+
+    Args:
+        libname: The name of the library to find
+
+    Returns:
+        The absolute path to the DLL if found, None otherwise
+    """
+    # Reverse tabulated names to achieve new → old search order.
+    for dll_name in reversed(SUPPORTED_WINDOWS_DLLS.get(libname, ())):
+        handle = kernel32.GetModuleHandleW(dll_name)
+        if handle:
+            abs_path = abs_path_for_dynamic_library(libname, handle)
+            return abs_path
+
+    return None
+
 
 def load_with_abs_path(libname: str, found_path: str) -> LoadedDL:
     """Load a dynamic library from the given path.

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py

@@ -6,51 +6,27 @@
 import sys
 
 from cuda.pathfinder._dynamic_libs.find_nvidia_dynamic_lib import _FindNvidiaDynamicLib
-from cuda.pathfinder._dynamic_libs.load_dl_common import LoadedDL, load_dependencies
+from cuda.pathfinder._dynamic_libs.load_dl_common import Distribution, LoadedDL, load_dependencies
 from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
 
 if IS_WINDOWS:
     from cuda.pathfinder._dynamic_libs.load_dl_windows import (
         check_if_already_loaded_from_elsewhere,
         load_with_abs_path,
-        load_with_system_search,
     )
 else:
     from cuda.pathfinder._dynamic_libs.load_dl_linux import (
         check_if_already_loaded_from_elsewhere,
         load_with_abs_path,
-        load_with_system_search,
     )
 
+from typing import Callable, Optional
 
-def _load_lib_no_cache(libname: str) -> LoadedDL:
-    finder = _FindNvidiaDynamicLib(libname)
-    abs_path = finder.try_site_packages()
-    if abs_path is None:
-        abs_path = finder.try_with_conda_prefix()
-
-    # If the library was already loaded by someone else, reproduce any OS-specific
-    # side-effects we would have applied on a direct absolute-path load (e.g.,
-    # AddDllDirectory on Windows for libs that require it).
-    loaded = check_if_already_loaded_from_elsewhere(libname, abs_path is not None)
-
-    # Load dependencies regardless of who loaded the primary lib first.
-    # Doing this *after* the side-effect ensures dependencies resolve consistently
-    # relative to the actually loaded location.
-    load_dependencies(libname, load_nvidia_dynamic_lib)
-
-    if loaded is not None:
-        return loaded
-
-    if abs_path is None:
-        loaded = load_with_system_search(libname)
-        if loaded is not None:
-            return loaded
-        abs_path = finder.try_with_cuda_home()
-        if abs_path is None:
-            finder.raise_not_found_error()
 
-    return load_with_abs_path(libname, abs_path)
+def _load_lib_no_cache(libname: str) -> LoadedDL:
+    loader = _LoadNvidiaDynamicLib(libname)
+    dl = loader.load_lib()
+    return dl
 
 
 @functools.cache
@@ -130,3 +106,75 @@ def load_nvidia_dynamic_lib(libname: str) -> LoadedDL:
             f" {sys.version_info.major}.{sys.version_info.minor}"
         )
     return _load_lib_no_cache(libname)
+
+
+class _LoadNvidiaDynamicLib:
+    def __init__(self, libname: str):
+        self.finder = _FindNvidiaDynamicLib(libname)
+        self.libname = self.finder.libname
+
+    def _load_with_dependencies(
+        self, get_path_func: Callable[[], Optional[LoadedDL]], dist_name: str
+    ) -> Optional[LoadedDL]:
+        # If the library was already loaded by someone else, reproduce any OS-specific
+        # side-effects we would have applied on a direct absolute-path load (e.g.,
+        # AddDllDirectory on Windows for libs that require it).
+        loaded = check_if_already_loaded_from_elsewhere(self.libname, True)
+
+        # Load dependencies regardless of who loaded the primary lib first.
+        # Doing this *after* the side-effect ensures dependencies resolve consistently
+        # relative to the actually loaded location.
+        load_dependencies(self.libname, load_nvidia_dynamic_lib)
+
+        abs_path = get_path_func()
+        if abs_path is None:
+            return None
+
+        dist = Distribution(name=dist_name, version="unknown")
+        if loaded is not None:
+            loaded.distribution = dist
+            return loaded
+
+        dl = load_with_abs_path(self.libname, abs_path)
+        dl.distribution = dist
+        return dl
+
+    def _load_simple(self, get_path_func: Callable[[], Optional[LoadedDL]], dist_name: str) -> Optional[LoadedDL]:
+        abs_path = get_path_func()
+        if abs_path is None:
+            return None
+
+        dl = load_with_abs_path(self.libname, abs_path)
+        dl.distribution = Distribution(name=dist_name, version="unknown")
+        return dl
+
+    def load_with_site_packages(self) -> Optional[LoadedDL]:
+        return self._load_simple(self.finder.try_site_packages, "site-packages")
+
+    def load_with_conda_prefix(self) -> Optional[LoadedDL]:
+        return self._load_simple(self.finder.try_with_conda_prefix, "conda")
+
+    def load_with_system_search(self) -> Optional[LoadedDL]:
+        return self._load_with_dependencies(self.finder.try_with_system_search, "system")
+
+    def load_with_cuda_home(self) -> Optional[LoadedDL]:
+        return self._load_with_dependencies(self.finder.try_with_cuda_home, "CUDA_HOME")
+
+    def load_lib(self) -> LoadedDL:
+        dl = self.load_with_site_packages()
+        if dl is not None:
+            return dl
+
+        dl = self.load_with_conda_prefix()
+        if dl is not None:
+            return dl
+
+        dl = self.load_with_system_search()
+        if dl is not None:
+            return dl
+
+        dl = self.load_with_cuda_home()
+        if dl is not None:
+            return dl
+
+        self.finder.raise_not_found_error()