Introduce dataclass LoadedDL as return type for load_nvidia_dynamic_library()

Author: rwgk

cuda_bindings/cuda/bindings/_bindings/cynvrtc.pyx.in

@@ -56,7 +56,7 @@ cdef int cuPythonInit() except -1 nogil:
 
     {{if 'Windows' == platform.system()}}
     with gil:
-        handle = path_finder.load_nvidia_dynamic_library("nvrtc")
+        handle = path_finder.load_nvidia_dynamic_library("nvrtc").handle
         {{if 'nvrtcGetErrorString' in found_functions}}
         try:
             global __nvrtcGetErrorString
@@ -242,7 +242,7 @@ cdef int cuPythonInit() except -1 nogil:
 
     {{else}}
     with gil:
-        handle = <void*><uintptr_t>path_finder.load_nvidia_dynamic_library("nvrtc")
+        handle = <void*><uintptr_t>path_finder.load_nvidia_dynamic_library("nvrtc").handle
     {{if 'nvrtcGetErrorString' in found_functions}}
     global __nvrtcGetErrorString
     __nvrtcGetErrorString = dlfcn.dlsym(handle, 'nvrtcGetErrorString')

cuda_bindings/cuda/bindings/_internal/nvjitlink_linux.pyx

@@ -53,7 +53,7 @@ cdef void* __nvJitLinkVersion = NULL
 
 
 cdef void* load_library(int driver_ver) except* with gil:
-    cdef uintptr_t handle = path_finder.load_nvidia_dynamic_library("nvJitLink")
+    cdef uintptr_t handle = path_finder.load_nvidia_dynamic_library("nvJitLink").handle
     return <void*>handle
 
 

cuda_bindings/cuda/bindings/_internal/nvjitlink_windows.pyx

@@ -40,7 +40,7 @@ cdef void* __nvJitLinkVersion = NULL
 
 
 cdef void* load_library(int driver_ver) except* with gil:
-    cdef intptr_t handle = path_finder.load_nvidia_dynamic_library("nvJitLink")
+    cdef intptr_t handle = path_finder.load_nvidia_dynamic_library("nvJitLink").handle
     return <void*>handle
 
 

cuda_bindings/cuda/bindings/_internal/nvvm_linux.pyx

@@ -51,7 +51,7 @@ cdef void* __nvvmGetProgramLog = NULL
 
 
 cdef void* load_library(const int driver_ver) except* with gil:
-    cdef uintptr_t handle = path_finder.load_nvidia_dynamic_library("nvvm")
+    cdef uintptr_t handle = path_finder.load_nvidia_dynamic_library("nvvm").handle
     return <void*>handle
 
 

cuda_bindings/cuda/bindings/_internal/nvvm_windows.pyx

@@ -38,7 +38,7 @@ cdef void* __nvvmGetProgramLog = NULL
 
 
 cdef void* load_library(int driver_ver) except* with gil:
-    cdef intptr_t handle = path_finder.load_nvidia_dynamic_library("nvvm")
+    cdef intptr_t handle = path_finder.load_nvidia_dynamic_library("nvvm").handle
     return <void*>handle
 
 

cuda_bindings/cuda/bindings/_path_finder/load_nvidia_dynamic_library.py

@@ -6,6 +6,7 @@
 import functools
 import os
 import sys
+from dataclasses import dataclass
 from typing import Optional, Tuple
 
 if sys.platform == "win32":
@@ -114,19 +115,29 @@ def _load_and_report_path_linux(libname, soname: str) -> Tuple[int, str]:
     return handle, info.dli_fname.decode()
 
 
+@dataclass
+class LoadedDL:
+    # ATTENTION: To convert `handle` back to `void*` in cython:
+    #     Linux:   `cdef void* ptr = <void*><uintptr_t>`
+    #     Windows: `cdef void* ptr = <void*><intptr_t>`
+    handle: int
+    abs_path: Optional[str]
+    was_already_loaded_from_elsewhere: bool
+
+
 @functools.cache
-def load_nvidia_dynamic_library(libname: str) -> int:
+def load_nvidia_dynamic_library(libname: str) -> LoadedDL:
     # Detect if the library was loaded already in some other way (i.e. not via this function).
     if sys.platform == "win32":
-        for dll_name in SUPPORTED_WINDOWS_DLLS.get(libname):
+        for dll_name in SUPPORTED_WINDOWS_DLLS.get(libname, ()):
             try:
-                return win32api.GetModuleHandle(dll_name)
+                return LoadedDL(win32api.GetModuleHandle(dll_name), None, True)
             except pywintypes.error:
                 pass
     else:
-        for soname in SUPPORTED_LINUX_SONAMES.get(libname):
+        for soname in SUPPORTED_LINUX_SONAMES.get(libname, ()):
             try:
-                return ctypes.CDLL(soname, mode=os.RTLD_NOLOAD)
+                return LoadedDL(ctypes.CDLL(soname, mode=os.RTLD_NOLOAD), None, True)
             except OSError:
                 pass
 
@@ -138,18 +149,14 @@ def load_nvidia_dynamic_library(libname: str) -> int:
         if sys.platform == "win32":
             handle, abs_path = _windows_load_with_dll_basename(libname)
             if handle:
-                # Use `cdef void* ptr = <void*><intptr_t>` in cython to convert back to void*
-                print(f"SYSTEM ABS_PATH for {libname=!r}: {abs_path}", flush=True)
-                return handle
+                return LoadedDL(handle, abs_path, False)
         else:
             try:
                 handle, abs_path = _load_and_report_path_linux(libname, found.lib_searched_for)
-            except OSError as e:
-                print(f"SYSTEM OSError for {libname=!r}: {e!r}", flush=True)
+            except OSError:
+                pass
             else:
-                # Use `cdef void* ptr = <void*><uintptr_t>` in cython to convert back to void*
-                print(f"SYSTEM ABS_PATH for {libname=!r}: {abs_path}", flush=True)
-                return handle._handle  # C unsigned int
+                return LoadedDL(handle._handle, abs_path, False)
         found.raise_if_abs_path_is_None()
 
     if sys.platform == "win32":
@@ -160,14 +167,10 @@ def load_nvidia_dynamic_library(libname: str) -> int:
             handle = win32api.LoadLibraryEx(found.abs_path, 0, flags)
         except pywintypes.error as e:
             raise RuntimeError(f"Failed to load DLL at {found.abs_path}: {e}") from e
-        # Use `cdef void* ptr = <void*><intptr_t>` in cython to convert back to void*
-        print(f"FOUND ABS_PATH for {libname=!r}: {found.abs_path}", flush=True)
-        return handle  # C signed int, matches win32api.GetProcAddress
+        return LoadedDL(handle, found.abs_path, False)
     else:
         try:
             handle = ctypes.CDLL(found.abs_path, _LINUX_CDLL_MODE)
         except OSError as e:
             raise RuntimeError(f"Failed to dlopen {found.abs_path}: {e}") from e
-        # Use `cdef void* ptr = <void*><uintptr_t>` in cython to convert back to void*
-        print(f"FOUND ABS_PATH for {libname=!r}: {found.abs_path}", flush=True)
-        return handle._handle  # C unsigned int
+        return LoadedDL(handle._handle, found.abs_path, False)

cuda_bindings/tests/path_finder.py

@@ -1,8 +1,8 @@
 import sys
+import traceback
 
 from cuda.bindings import path_finder
 from cuda.bindings._path_finder import cuda_paths, supported_libs
-from cuda.bindings._path_finder.find_nvidia_dynamic_library import find_nvidia_dynamic_library
 
 ALL_LIBNAMES = path_finder.SUPPORTED_LIBNAMES + supported_libs.PARTIALLY_SUPPORTED_LIBNAMES
 
@@ -11,26 +11,20 @@ def run(args):
     assert len(args) == 0
 
     paths = cuda_paths.get_cuda_paths()
-
     for k, v in paths.items():
         print(f"{k}: {v}", flush=True)
     print()
 
-    for libname in supported_libs.SUPPORTED_WINDOWS_DLLS:
-        if libname not in ALL_LIBNAMES:
-            print(f"MISSING IN SUPPORTED_LIBNAMES: {libname}")
-
     for libname in ALL_LIBNAMES:
-        print(libname)
-        dlls = supported_libs.SUPPORTED_WINDOWS_DLLS.get(libname)
-        if dlls is None:
-            print(f"MISSING IN SUPPORTED_WINDOWS_DLLS: {libname}")
-        for fun in (find_nvidia_dynamic_library, path_finder.load_nvidia_dynamic_library):
-            try:
-                out = fun(libname)
-            except Exception as e:
-                out = f"EXCEPTION: {type(e)} {str(e)}"
-            print(out)
+        print(f"{libname=}")
+        try:
+            loaded_dl = path_finder.load_nvidia_dynamic_library(libname)
+        except Exception as e:
+            print(f"EXCEPTION for {libname=}:")
+            traceback.print_exc(file=sys.stdout)
+        else:
+            print(f"    {loaded_dl.abs_path=!r}")
+            print(f"    {loaded_dl.was_already_loaded_from_elsewhere=!r}")
         print()