Factor out _find_dll_under_dir(dirpath, file_wild) and reuse from _find_dll_using_nvidia_bin_dirs(), _find_dll_using_cudalib_dir() (to fix loading nvrtc64_120_0.dll from local CTK)

Author: rwgk

cuda_bindings/cuda/bindings/_path_finder/find_nvidia_dynamic_library.py

@@ -43,39 +43,43 @@ def _append_to_os_environ_path(dirpath):
     os.environ["PATH"] = dirpath if curr_path is None else os.pathsep.join((curr_path, dirpath))
 
 
+def _find_dll_under_dir(dirpath, file_wild):
+    dll_name = None
+    have_builtins = False
+    for path in sorted(glob.glob(os.path.join(dirpath, file_wild))):
+        # nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-win_amd64.whl:
+        #     nvidia\cuda_nvrtc\bin\
+        #         nvrtc-builtins64_128.dll
+        #         nvrtc64_120_0.alt.dll
+        #         nvrtc64_120_0.dll
+        node = os.path.basename(path)
+        if node.endswith(".alt.dll"):
+            continue
+        if "-builtins" in node:
+            have_builtins = True
+            continue
+        if dll_name is not None:
+            continue
+        if os.path.isfile(path):
+            dll_name = path
+    if dll_name is not None:
+        if have_builtins:
+            # Add the DLL directory to the search path
+            os.add_dll_directory(dirpath)
+            # Update PATH as a fallback for dependent DLL resolution
+            _append_to_os_environ_path(dirpath)
+        return dll_name
+
+
 def _find_dll_using_nvidia_bin_dirs(libname, error_messages, attachments):
     if libname == "nvvm":  # noqa: SIM108
         nvidia_sub_dirs = ("nvidia", "*", "nvvm", "bin")
     else:
         nvidia_sub_dirs = ("nvidia", "*", "bin")
     file_wild = libname + "*.dll"
     for bin_dir in sys_path_find_sub_dirs(nvidia_sub_dirs):
-        dll_name = None
-        have_builtins = False
-        for path in sorted(glob.glob(os.path.join(bin_dir, file_wild))):
-            # nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-win_amd64.whl:
-            #     nvidia\cuda_nvrtc\bin\
-            #         nvrtc-builtins64_128.dll
-            #         nvrtc64_120_0.alt.dll
-            #         nvrtc64_120_0.dll
-            # See also:
-            # https://github.com/NVIDIA/cuda-python/pull/563#discussion_r2054427641
-            node = os.path.basename(path)
-            if node.endswith(".alt.dll"):
-                continue
-            if "-builtins" in node:
-                have_builtins = True
-                continue
-            if dll_name is not None:
-                continue
-            if os.path.isfile(path):
-                dll_name = path
+        dll_name = _find_dll_under_dir(bin_dir, file_wild)
         if dll_name is not None:
-            if have_builtins:
-                # Add the DLL directory to the search path
-                os.add_dll_directory(bin_dir)
-                # Update PATH as a fallback for dependent DLL resolution
-                _append_to_os_environ_path(bin_dir)
             return dll_name
     _no_such_file_in_sub_dirs(nvidia_sub_dirs, file_wild, error_messages, attachments)
     return None
@@ -124,9 +128,9 @@ def _find_dll_using_cudalib_dir(libname, error_messages, attachments):
     if cudalib_dir is None:
         return None
     file_wild = libname + "*.dll"
-    for dll_name in sorted(glob.glob(os.path.join(cudalib_dir, file_wild))):
-        if os.path.isfile(dll_name):
-            return dll_name
+    dll_name = _find_dll_under_dir(cudalib_dir, file_wild)
+    if dll_name is not None:
+        return dll_name
     error_messages.append(f"No such file: {file_wild}")
     attachments.append(f'  listdir("{cudalib_dir}"):')
     for node in sorted(os.listdir(cudalib_dir)):