77import ctypes .wintypes
88import os
99import struct
10+ from collections .abc import Iterator
1011from typing import TYPE_CHECKING
1112
1213from cuda .pathfinder ._dynamic_libs .load_dl_common import LoadedDL
1516 from cuda .pathfinder ._dynamic_libs .lib_descriptor import LibDescriptor
1617
1718# Mirrors WinBase.h (unfortunately not defined already elsewhere)
18- WINBASE_LOAD_WITH_ALTERED_SEARCH_PATH = 0x00000008
1919WINBASE_LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR = 0x00000100
2020WINBASE_LOAD_LIBRARY_SEARCH_DEFAULT_DIRS = 0x00001000
2121
4848kernel32 .AddDllDirectory .argtypes = [ctypes .wintypes .LPCWSTR ]
4949kernel32 .AddDllDirectory .restype = ctypes .c_void_p # DLL_DIRECTORY_COOKIE
5050
51- # SearchPathW - find a file in the system search path
52- kernel32 .SearchPathW .argtypes = [
53- ctypes .wintypes .LPCWSTR , # lpPath (NULL to use standard search)
54- ctypes .wintypes .LPCWSTR , # lpFileName
55- ctypes .wintypes .LPCWSTR , # lpExtension
56- ctypes .wintypes .DWORD , # nBufferLength
57- ctypes .wintypes .LPWSTR , # lpBuffer
58- ctypes .POINTER (ctypes .wintypes .LPWSTR ), # lpFilePart
59- ]
60- kernel32 .SearchPathW .restype = ctypes .wintypes .DWORD
61-
6251
6352def ctypes_handle_to_unsigned_int (handle : ctypes .wintypes .HMODULE ) -> int :
6453 """Convert ctypes HMODULE to unsigned int."""
@@ -127,29 +116,55 @@ def check_if_already_loaded_from_elsewhere(desc: LibDescriptor, have_abs_path: b
127116 return None
128117
129118
130- def _search_path_for_dll (dll_name : str ) -> str | None :
131- """Search for a DLL using Windows SearchPathW.
132-
133- Args:
134- dll_name: The name of the DLL to find
119+ def _iter_env_path_directories (path_value : str | None ) -> Iterator [str ]:
120+ """Yield normalized directories from PATH without consulting the current directory."""
121+ seen : set [str ] = set ()
122+ if not path_value :
123+ return
124+
125+ for raw_entry in path_value .split (os .pathsep ):
126+ entry = os .path .expandvars (raw_entry .strip ().strip ('"' ))
127+ if not entry :
128+ continue
129+ if not os .path .isabs (entry ):
130+ # Relative PATH entries would implicitly consult the current
131+ # directory, which we explicitly avoid for DLL lookup.
132+ continue
133+ if not os .path .isdir (entry ):
134+ continue
135+
136+ normalized_entry = os .path .normcase (os .path .normpath (entry ))
137+ if normalized_entry in seen :
138+ continue
139+ seen .add (normalized_entry )
140+ yield entry
141+
142+
143+ def _find_dll_on_env_path (dll_name : str ) -> str | None :
144+ """Locate a DLL by scanning PATH entries explicitly."""
145+ for dirpath in _iter_env_path_directories (os .environ .get ("PATH" )):
146+ candidate = os .path .join (dirpath , dll_name )
147+ if os .path .isfile (candidate ):
148+ return candidate
149+ return None
135150
136- Returns:
137- The absolute path to the DLL if found, None otherwise
138- """
139- buffer = ctypes .create_unicode_buffer (260 ) # MAX_PATH
140- length = kernel32 .SearchPathW (None , dll_name , None , len (buffer ), buffer , None )
141151
142- if length == 0 :
152+ def _try_load_with_process_dll_search (desc : LibDescriptor , dll_name : str ) -> LoadedDL | None :
153+ """Try the process DLL search path configured by CPython/Windows."""
154+ handle = kernel32 .LoadLibraryExW (dll_name , None , 0 )
155+ if not handle :
143156 return None
144157
145- # If buffer was too small, try with larger buffer
146- if length > len (buffer ):
147- buffer = ctypes .create_unicode_buffer (length )
148- length = kernel32 .SearchPathW (None , dll_name , None , len (buffer ), buffer , None )
149- if length == 0 :
150- return None
158+ abs_path = abs_path_for_dynamic_library (desc .name , handle )
159+ return LoadedDL (abs_path , False , ctypes_handle_to_unsigned_int (handle ), "system-search" )
151160
152- return buffer .value
161+
162+ def _try_load_with_env_path_fallback (desc : LibDescriptor , dll_name : str ) -> LoadedDL | None :
163+ """Fallback for CTK-style installs exposed only via PATH."""
164+ found_path = _find_dll_on_env_path (dll_name )
165+ if found_path is None :
166+ return None
167+ return load_with_abs_path (desc , found_path , "system-search" )
153168
154169
155170def load_with_system_search (desc : LibDescriptor ) -> LoadedDL | None :
@@ -161,16 +176,25 @@ def load_with_system_search(desc: LibDescriptor) -> LoadedDL | None:
161176 Returns:
162177 A LoadedDL object if successful, None if the library cannot be loaded
163178 """
164- # Reverse tabulated names to achieve new -> old search order.
165- for dll_name in reversed (desc .windows_dlls ):
166- # SearchPathW bypasses Python 3.8+'s SetDefaultDllDirectories restriction.
167- found_path = _search_path_for_dll (dll_name )
168- if found_path :
169- # LOAD_WITH_ALTERED_SEARCH_PATH additionally ensures dependencies
170- # are resolved from the DLL's directory.
171- handle = kernel32 .LoadLibraryExW (found_path , None , WINBASE_LOAD_WITH_ALTERED_SEARCH_PATH )
172- if handle :
173- return LoadedDL (found_path , False , ctypes_handle_to_unsigned_int (handle ), "system-search" )
179+ dll_names = tuple (reversed (desc .windows_dlls ))
180+
181+ # Phase 1: preserve the native process DLL search path (application dir,
182+ # system32, AddDllDirectory user dirs, loaded-module list).
183+ for dll_name in dll_names :
184+ loaded = _try_load_with_process_dll_search (desc , dll_name )
185+ if loaded is not None :
186+ return loaded
187+
188+ if desc .packaged_with == "driver" :
189+ return None
190+
191+ # Phase 2: explicit PATH fallback for CTK-style installs only. Avoid
192+ # SearchPathW because its search semantics differ from LoadLibraryExW and
193+ # can consult the current directory.
194+ for dll_name in dll_names :
195+ loaded = _try_load_with_env_path_fallback (desc , dll_name )
196+ if loaded is not None :
197+ return loaded
174198
175199 return None
176200
0 commit comments