Remove os.add_dll_directory() and os.environ["PATH"] manipulations from find_nvidia_dynamic_library.py. Add supported_libs.LIBNAMES_REQUIRING_OS_ADD_DLL_DIRECTORY and use from load_nvidia_dynamic_library().

Author: rwgk

cuda_bindings/cuda/bindings/_path_finder/find_nvidia_dynamic_library.py

@@ -38,14 +38,8 @@ def _find_so_using_nvidia_lib_dirs(libname, so_basename, error_messages, attachm
     return None
 
 
-def _append_to_os_environ_path(dirpath):
-    curr_path = os.environ.get("PATH")
-    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\
@@ -56,19 +50,12 @@ def _find_dll_under_dir(dirpath, file_wild):
         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
+    return dll_name
 
 
 def _find_dll_using_nvidia_bin_dirs(libname, error_messages, attachments):

cuda_bindings/cuda/bindings/_path_finder/load_nvidia_dynamic_library.py

@@ -4,6 +4,7 @@
 
 import ctypes
 import functools
+import os
 import sys
 from typing import Optional, Tuple
 
@@ -19,7 +20,6 @@
 
 else:
     import ctypes.util
-    import os
 
     _LINUX_CDLL_MODE = os.RTLD_NOW | os.RTLD_GLOBAL
 
@@ -38,7 +38,23 @@ class Dl_info(ctypes.Structure):
 
 
 from .find_nvidia_dynamic_library import _find_nvidia_dynamic_library
-from .supported_libs import DIRECT_DEPENDENCIES, EXPECTED_LIB_SYMBOLS, SUPPORTED_LINUX_SONAMES, SUPPORTED_WINDOWS_DLLS
+from .supported_libs import (
+    DIRECT_DEPENDENCIES,
+    EXPECTED_LIB_SYMBOLS,
+    LIBNAMES_REQUIRING_OS_ADD_DLL_DIRECTORY,
+    SUPPORTED_LINUX_SONAMES,
+    SUPPORTED_WINDOWS_DLLS,
+)
+
+
+def _add_dll_directory(dll_abs_path):
+    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))
 
 
 @functools.cache
@@ -137,6 +153,8 @@ def load_nvidia_dynamic_library(libname: str) -> int:
         found.raise_if_abs_path_is_None()
 
     if sys.platform == "win32":
+        if libname in LIBNAMES_REQUIRING_OS_ADD_DLL_DIRECTORY:
+            _add_dll_directory(found.abs_path)
         flags = _WINBASE_LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | _WINBASE_LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR
         try:
             handle = win32api.LoadLibraryEx(found.abs_path, 0, flags)

cuda_bindings/cuda/bindings/_path_finder/supported_libs.py

@@ -314,6 +314,11 @@
     ),
 }
 
+LIBNAMES_REQUIRING_OS_ADD_DLL_DIRECTORY = (
+    "cufft",
+    "nvrtc",
+)
+
 # Based on nm output for Linux x86_64 /usr/local/cuda (12.8.1)
 EXPECTED_LIB_SYMBOLS = {
     "nvJitLink": ("nvJitLinkVersion",),

cuda_bindings/tests/test_path_finder.py

@@ -22,6 +22,10 @@ def test_all_libnames_windows_dlls_consistency():
     assert tuple(sorted(ALL_LIBNAMES)) == tuple(sorted(supported_libs.SUPPORTED_WINDOWS_DLLS.keys()))
 
 
+def test_all_libnames_libnames_requiring_os_add_dll_directory_consistency():
+    assert not (set(supported_libs.LIBNAMES_REQUIRING_OS_ADD_DLL_DIRECTORY) - set(ALL_LIBNAMES))
+
+
 def test_all_libnames_expected_lib_symbols_consistency():
     assert tuple(sorted(ALL_LIBNAMES)) == tuple(sorted(supported_libs.EXPECTED_LIB_SYMBOLS.keys()))