Save result of factoring out load_dl_common.py, load_dl_linux.py, load_dl_windows.py with the help of Cursor.

Author: rwgk

cuda_bindings/cuda/bindings/_path_finder/load_dl_common.py

@@ -1,13 +1,22 @@
 # Copyright 2025 NVIDIA Corporation.  All rights reserved.
-#
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
 from dataclasses import dataclass
-from typing import Optional
+from typing import Callable, Optional
+
+from .supported_libs import DIRECT_DEPENDENCIES
 
 
 @dataclass
 class LoadedDL:
+    """Represents a loaded dynamic library.
+
+    Attributes:
+        handle: The library handle (can be converted to void* in Cython)
+        abs_path: The absolute path to the library file
+        was_already_loaded_from_elsewhere: Whether the library was already loaded
+    """
+
     # ATTENTION: To convert `handle` back to `void*` in cython:
     #     Linux:   `cdef void* ptr = <void*><uintptr_t>`
     #     Windows: `cdef void* ptr = <void*><intptr_t>`
@@ -17,13 +26,35 @@ class LoadedDL:
 
 
 def add_dll_directory(dll_abs_path: str) -> None:
-    """Add a DLL directory to the search path and update PATH environment variable."""
+    """Add a DLL directory to the search path and update PATH environment variable.
+
+    Args:
+        dll_abs_path: Absolute path to the DLL file
+
+    Raises:
+        AssertionError: If the directory containing the DLL does not exist
+    """
     import os
-    
+
     dirpath = os.path.dirname(dll_abs_path)
     assert os.path.isdir(dirpath), dll_abs_path
     # Add the DLL directory to the search path
     os.add_dll_directory(dirpath)
     # Update PATH as a fallback for dependent DLL resolution
     curr_path = os.environ.get("PATH")
-    os.environ["PATH"] = dirpath if curr_path is None else os.pathsep.join((curr_path, dirpath)) 
+    os.environ["PATH"] = dirpath if curr_path is None else os.pathsep.join((curr_path, dirpath))
+
+
+def load_dependencies(libname: str, load_func: Callable[[str], LoadedDL]) -> None:
+    """Load all dependencies for a given library.
+
+    Args:
+        libname: The name of the library whose dependencies should be loaded
+        load_func: The function to use for loading libraries (e.g. load_nvidia_dynamic_library)
+
+    Example:
+        >>> load_dependencies("cudart", load_nvidia_dynamic_library)
+        # This will load all dependencies of cudart using the provided loading function
+    """
+    for dep in DIRECT_DEPENDENCIES.get(libname, ()):
+        load_func(dep)

cuda_bindings/cuda/bindings/_path_finder/load_dl_linux.py

@@ -1,23 +1,24 @@
 # Copyright 2025 NVIDIA Corporation.  All rights reserved.
-#
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
 import ctypes
 import ctypes.util
 import os
-from typing import Optional, Tuple
+from typing import Optional
 
 from .load_dl_common import LoadedDL
 
-_LINUX_CDLL_MODE = os.RTLD_NOW | os.RTLD_GLOBAL
+CDLL_MODE = os.RTLD_NOW | os.RTLD_GLOBAL
 
-_LIBDL_PATH = ctypes.util.find_library("dl") or "libdl.so.2"
-_LIBDL = ctypes.CDLL(_LIBDL_PATH)
-_LIBDL.dladdr.argtypes = [ctypes.c_void_p, ctypes.c_void_p]
-_LIBDL.dladdr.restype = ctypes.c_int
+LIBDL_PATH = ctypes.util.find_library("dl") or "libdl.so.2"
+LIBDL = ctypes.CDLL(LIBDL_PATH)
+LIBDL.dladdr.argtypes = [ctypes.c_void_p, ctypes.c_void_p]
+LIBDL.dladdr.restype = ctypes.c_int
 
 
 class Dl_info(ctypes.Structure):
+    """Structure used by dladdr to return information about a loaded symbol."""
+
     _fields_ = [
         ("dli_fname", ctypes.c_char_p),  # path to .so
         ("dli_fbase", ctypes.c_void_p),
@@ -27,50 +28,98 @@ class Dl_info(ctypes.Structure):
 
 
 def abs_path_for_dynamic_library(libname: str, handle: int) -> Optional[str]:
-    """Get the absolute path of a loaded dynamic library on Linux."""
+    """Get the absolute path of a loaded dynamic library on Linux.
+
+    Args:
+        libname: The name of the library
+        handle: The library handle
+
+    Returns:
+        The absolute path to the library file, or None if no expected symbol is found
+
+    Raises:
+        OSError: If dladdr fails to get information about the symbol
+    """
     from .supported_libs import EXPECTED_LIB_SYMBOLS
-    
+
     for symbol_name in EXPECTED_LIB_SYMBOLS[libname]:
         symbol = getattr(handle, symbol_name, None)
         if symbol is not None:
             break
     else:
         return None
-        
+
     addr = ctypes.cast(symbol, ctypes.c_void_p)
     info = Dl_info()
-    if _LIBDL.dladdr(addr, ctypes.byref(info)) == 0:
+    if LIBDL.dladdr(addr, ctypes.byref(info)) == 0:
         raise OSError(f"dladdr failed for {libname=!r}")
     return info.dli_fname.decode()
 
 
-def load_and_report_path(libname: str, soname: str) -> Tuple[int, str]:
-    """Load a dynamic library and return its handle and absolute path."""
-    handle = ctypes.CDLL(soname, _LINUX_CDLL_MODE)
-    abs_path = abs_path_for_dynamic_library(libname, handle)
-    if abs_path is None:
-        raise RuntimeError(f"No expected symbol for {libname=!r}")
-    return handle._handle, abs_path
+def check_if_already_loaded(libname: str) -> Optional[LoadedDL]:
+    """Check if the library is already loaded in the process.
 
+    Args:
+        libname: The name of the library to check
 
-def load_dynamic_library(libname: str, found_path: str) -> LoadedDL:
-    """Load a dynamic library from the given path."""
-    try:
-        handle = ctypes.CDLL(found_path, _LINUX_CDLL_MODE)
-    except OSError as e:
-        raise RuntimeError(f"Failed to dlopen {found_path}: {e}") from e
-    return LoadedDL(handle._handle, found_path, False)
-
+    Returns:
+        A LoadedDL object if the library is already loaded, None otherwise
 
-def check_if_already_loaded(libname: str) -> Optional[LoadedDL]:
-    """Check if the library is already loaded in the process."""
+    Example:
+        >>> loaded = check_if_already_loaded("cudart")
+        >>> if loaded is not None:
+        ...     print(f"Library already loaded from {loaded.abs_path}")
+    """
     from .supported_libs import SUPPORTED_LINUX_SONAMES
-    
+
     for soname in SUPPORTED_LINUX_SONAMES.get(libname, ()):
         try:
             handle = ctypes.CDLL(soname, mode=os.RTLD_NOLOAD)
         except OSError:
             continue
         else:
             return LoadedDL(handle._handle, abs_path_for_dynamic_library(libname, handle), True)
-    return None 
+    return None
+
+
+def load_with_system_search(libname: str, soname: str) -> Optional[LoadedDL]:
+    """Try to load a library using system search paths.
+
+    Args:
+        libname: The name of the library to load
+        soname: The soname to search for
+
+    Returns:
+        A LoadedDL object if successful, None if the library cannot be loaded
+
+    Raises:
+        RuntimeError: If the library is loaded but no expected symbol is found
+    """
+    try:
+        handle = ctypes.CDLL(soname, CDLL_MODE)
+        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(handle._handle, abs_path, False)
+    except OSError:
+        return 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
+    """
+    try:
+        handle = ctypes.CDLL(found_path, CDLL_MODE)
+    except OSError as e:
+        raise RuntimeError(f"Failed to dlopen {found_path}: {e}") from e
+    return LoadedDL(handle._handle, found_path, False)

cuda_bindings/cuda/bindings/_path_finder/load_dl_windows.py

@@ -1,23 +1,33 @@
 # Copyright 2025 NVIDIA Corporation.  All rights reserved.
-#
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
 import ctypes
 import ctypes.wintypes
-from typing import Optional, Tuple
+import functools
+from typing import Optional
 
 import pywintypes
 import win32api
 
-from .load_dl_common import LoadedDL
+from .load_dl_common import LoadedDL, add_dll_directory
 
 # Mirrors WinBase.h (unfortunately not defined already elsewhere)
-_WINBASE_LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR = 0x00000100
-_WINBASE_LOAD_LIBRARY_SEARCH_DEFAULT_DIRS = 0x00001000
+WINBASE_LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR = 0x00000100
+WINBASE_LOAD_LIBRARY_SEARCH_DEFAULT_DIRS = 0x00001000
 
 
 def abs_path_for_dynamic_library(handle: int) -> str:
-    """Get the absolute path of a loaded dynamic library on Windows."""
+    """Get the absolute path of a loaded dynamic library on Windows.
+
+    Args:
+        handle: The library handle
+
+    Returns:
+        The absolute path to the DLL file
+
+    Raises:
+        OSError: If GetModuleFileNameW fails
+    """
     buf = ctypes.create_unicode_buffer(260)
     n_chars = ctypes.windll.kernel32.GetModuleFileNameW(ctypes.wintypes.HMODULE(handle), buf, len(buf))
     if n_chars == 0:
@@ -27,7 +37,14 @@ def abs_path_for_dynamic_library(handle: int) -> str:
 
 @functools.cache
 def cuDriverGetVersion() -> int:
-    """Get the CUDA driver version."""
+    """Get the CUDA driver version.
+
+    Returns:
+        The CUDA driver version number
+
+    Raises:
+        AssertionError: If the driver version cannot be obtained
+    """
     handle = win32api.LoadLibrary("nvcuda.dll")
 
     kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
@@ -45,50 +62,80 @@ def cuDriverGetVersion() -> int:
     return driver_ver.value
 
 
-@functools.cache
-def load_with_dll_basename(name: str) -> Tuple[Optional[int], Optional[str]]:
-    """Try to load a DLL by its basename."""
+def check_if_already_loaded(libname: str) -> Optional[LoadedDL]:
+    """Check if the library is already loaded in the process.
+
+    Args:
+        libname: The name of the library to check
+
+    Returns:
+        A LoadedDL object if the library is already loaded, None otherwise
+
+    Example:
+        >>> loaded = check_if_already_loaded("cudart")
+        >>> if loaded is not None:
+        ...     print(f"Library already loaded from {loaded.abs_path}")
+    """
+    from .supported_libs import SUPPORTED_WINDOWS_DLLS
+
+    for dll_name in SUPPORTED_WINDOWS_DLLS.get(libname, ()):
+        try:
+            handle = win32api.GetModuleHandle(dll_name)
+        except pywintypes.error:
+            continue
+        else:
+            return LoadedDL(handle, abs_path_for_dynamic_library(handle), True)
+    return None
+
+
+def load_with_system_search(name: str, _unused: str) -> Optional[LoadedDL]:
+    """Try to load a DLL using system search paths.
+
+    Args:
+        name: The name of the library to load
+        _unused: Unused parameter (kept for interface consistency)
+
+    Returns:
+        A LoadedDL object if successful, None if the library cannot be loaded
+    """
     from .supported_libs import SUPPORTED_WINDOWS_DLLS
-    
+
     driver_ver = cuDriverGetVersion()
     del driver_ver  # Keeping this here because it will probably be needed in the future.
 
     dll_names = SUPPORTED_WINDOWS_DLLS.get(name)
     if dll_names is None:
-        return None, None
+        return None
 
     for dll_name in dll_names:
         handle = ctypes.windll.kernel32.LoadLibraryW(ctypes.c_wchar_p(dll_name))
         if handle:
-            return handle, abs_path_for_dynamic_library(handle)
+            return LoadedDL(handle, abs_path_for_dynamic_library(handle), False)
+
+    return None
 
-    return None, None
 
+def load_with_abs_path(libname: str, found_path: str) -> LoadedDL:
+    """Load a dynamic library from the given path.
 
-def load_dynamic_library(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 DLL file
+
+    Returns:
+        A LoadedDL object representing the loaded library
+
+    Raises:
+        RuntimeError: If the DLL cannot be loaded
+    """
     from .supported_libs import LIBNAMES_REQUIRING_OS_ADD_DLL_DIRECTORY
-    
+
     if libname in LIBNAMES_REQUIRING_OS_ADD_DLL_DIRECTORY:
         add_dll_directory(found_path)
-        
-    flags = _WINBASE_LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | _WINBASE_LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR
+
+    flags = WINBASE_LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | WINBASE_LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR
     try:
         handle = win32api.LoadLibraryEx(found_path, 0, flags)
     except pywintypes.error as e:
         raise RuntimeError(f"Failed to load DLL at {found_path}: {e}") from e
     return LoadedDL(handle, found_path, False)
-
-
-def check_if_already_loaded(libname: str) -> Optional[LoadedDL]:
-    """Check if the library is already loaded in the process."""
-    from .supported_libs import SUPPORTED_WINDOWS_DLLS
-    
-    for dll_name in SUPPORTED_WINDOWS_DLLS.get(libname, ()):
-        try:
-            handle = win32api.GetModuleHandle(dll_name)
-        except pywintypes.error:
-            continue
-        else:
-            return LoadedDL(handle, abs_path_for_dynamic_library(handle), True)
-    return None