Skip to content
Merged
Show file tree
Hide file tree
Changes from 13 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ class LoadedDL:
abs_path: Optional[str]
was_already_loaded_from_elsewhere: bool
_handle_uint: int # Platform-agnostic unsigned pointer value
foundvia: Optional[str] = None
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Two trivial independent change here would be great:

  1. found_via with underscore. I think you only have to change this and the one reference child_load_nvidia_dynamic_lib_helper.py
  2. Could you please try if pre-commit run --all-files still passes after removing Optional[] here? That'd be a promise that we're always populating the field.



def load_dependencies(libname: str, load_func: Callable[[str], LoadedDL]) -> None:
Expand Down
10 changes: 6 additions & 4 deletions cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_linux.py
Original file line number Diff line number Diff line change
Expand Up @@ -138,7 +138,9 @@ def check_if_already_loaded_from_elsewhere(libname: str, _have_abs_path: bool) -
except OSError:
continue
else:
return LoadedDL(abs_path_for_dynamic_library(libname, handle), True, handle._handle)
return LoadedDL(
abs_path_for_dynamic_library(libname, handle), True, handle._handle, "already-loaded-from-elsewhere"
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you please add was- here, so it's consistent with the existing was_already_loaded_from_elsewhere (so that people don't wonder "why the difference")?

But I think it's great to use minus signs instead of underscores, as you do.

)
return None


Expand Down Expand Up @@ -170,7 +172,7 @@ def load_with_system_search(libname: str) -> Optional[LoadedDL]:
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(abs_path, False, handle._handle)
return LoadedDL(abs_path, False, handle._handle, "system-search")
return None


Expand All @@ -193,7 +195,7 @@ def _work_around_known_bugs(libname: str, found_path: str) -> None:
ctypes.CDLL(dep_path, CDLL_MODE)


def load_with_abs_path(libname: str, found_path: str) -> LoadedDL:
def load_with_abs_path(libname: str, found_path: str, found_via: Optional[str] = None) -> LoadedDL:
"""Load a dynamic library from the given path.

Args:
Expand All @@ -211,4 +213,4 @@ def load_with_abs_path(libname: str, found_path: str) -> LoadedDL:
handle = _load_lib(libname, found_path)
except OSError as e:
raise RuntimeError(f"Failed to dlopen {found_path}: {e}") from e
return LoadedDL(found_path, False, handle._handle)
return LoadedDL(found_path, False, handle._handle, found_via)
13 changes: 8 additions & 5 deletions cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_windows.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,10 @@ def abs_path_for_dynamic_library(libname: str, handle: ctypes.wintypes.HMODULE)
return buffer.value


def check_if_already_loaded_from_elsewhere(libname: str, have_abs_path: bool) -> Optional[LoadedDL]:
def check_if_already_loaded_from_elsewhere(
libname: str,
have_abs_path: bool,
) -> Optional[LoadedDL]:
for dll_name in SUPPORTED_WINDOWS_DLLS.get(libname, ()):
handle = kernel32.GetModuleHandleW(dll_name)
if handle:
Expand All @@ -110,7 +113,7 @@ def check_if_already_loaded_from_elsewhere(libname: str, have_abs_path: bool) ->
# load_with_abs_path(). To make the side-effect more deterministic,
# activate it even if the library was already loaded from elsewhere.
add_dll_directory(abs_path)
return LoadedDL(abs_path, True, ctypes_handle_to_unsigned_int(handle))
return LoadedDL(abs_path, True, ctypes_handle_to_unsigned_int(handle), "already-loaded-from-elsewhere")
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

was- here, too.

return None


Expand All @@ -128,12 +131,12 @@ def load_with_system_search(libname: str) -> Optional[LoadedDL]:
handle = kernel32.LoadLibraryExW(dll_name, None, 0)
if handle:
abs_path = abs_path_for_dynamic_library(libname, handle)
return LoadedDL(abs_path, False, ctypes_handle_to_unsigned_int(handle))
return LoadedDL(abs_path, False, ctypes_handle_to_unsigned_int(handle), "system-search")

return None


def load_with_abs_path(libname: str, found_path: str) -> LoadedDL:
def load_with_abs_path(libname: str, found_path: str, found_via: Optional[str] = None) -> LoadedDL:
"""Load a dynamic library from the given path.

Args:
Expand All @@ -156,4 +159,4 @@ def load_with_abs_path(libname: str, found_path: str) -> LoadedDL:
error_code = ctypes.GetLastError() # type: ignore[attr-defined]
raise RuntimeError(f"Failed to load DLL at {found_path}: Windows error {error_code}")

return LoadedDL(found_path, False, ctypes_handle_to_unsigned_int(handle))
return LoadedDL(found_path, False, ctypes_handle_to_unsigned_int(handle), found_via)
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,13 @@
def _load_lib_no_cache(libname: str) -> LoadedDL:
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would be nice to move this below the class _LoadNvidiaDynamicLib code (it's more in line with how I organized the rest of the pathfinder code).

Also, this could be a one-liner:

    return _LoadNvidiaDynamicLib(libname).load_lib()

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I suspect it'll be simpler (less code) to wrap this function in a helper that adds found_via into LoadedDL.

def _load_lib_no_cache_impl(libname: str) -> LoadedDL, str
    existing code, but we also keep track of `found_via` and return (loaded, found_via)
def _load_lib_no_cache(libname: str) -> LoadedDL:
    loaded, found_via = _load_lib_no_cache_impl(libname)
    loaded.found_via = FoundVia(found_via)
    return loaded

Alternatively, we could pass found_via down along with abs_path, so that the LoadedDL instances can be constructed with it, rather than retrofitting it in later. That's a bit more surgical, not sure how it'll shake out, but it would seem cleanest.

finder = _FindNvidiaDynamicLib(libname)
abs_path = finder.try_site_packages()
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor but would be nice: keep the found_via assignment closer the the source of the abs_path:

    abs_path = finder.try_site_packages()
    if abs_path is not None:
        found_via = "site-packages"
    else:
        abs_path = finder.try_with_conda_prefix()
        if abs_path is not None:
            found_via = "conda"


if abs_path is None:
abs_path = finder.try_with_conda_prefix()
if abs_path is not None:
found_via = "conda"
else:
found_via = "site-packages"

# If the library was already loaded by someone else, reproduce any OS-specific
# side-effects we would have applied on a direct absolute-path load (e.g.,
Expand All @@ -49,8 +54,10 @@ def _load_lib_no_cache(libname: str) -> LoadedDL:
abs_path = finder.try_with_cuda_home()
if abs_path is None:
finder.raise_not_found_error()
else:
found_via = "CUDA_HOME"

return load_with_abs_path(libname, abs_path)
return load_with_abs_path(libname, abs_path, found_via)


@functools.cache
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ def child_process_func(libname):
if loaded_dl_fresh.was_already_loaded_from_elsewhere:
raise RuntimeError("loaded_dl_fresh.was_already_loaded_from_elsewhere")
validate_abs_path(loaded_dl_fresh.abs_path)
assert loaded_dl_fresh.foundvia is not None

loaded_dl_from_cache = load_nvidia_dynamic_lib(libname)
if loaded_dl_from_cache is not loaded_dl_fresh:
Expand Down