-
Notifications
You must be signed in to change notification settings - Fork 214
Add conda-specific search into load_nvidia_dynamic_lib()
#1003
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
5174ea2
8da27a7
37cb325
3e5aa1d
5b3c5eb
233ad08
2642d35
edff6ec
5ba032c
cc7d84d
894c6b5
bd5e2ae
ce2ca30
9e0b4a8
9cd31bb
5fd4801
7498392
2841734
c83fa0d
22999d1
3e7ca0e
917803e
f1d39d9
37414a5
1bde30e
5056e46
2c336c9
735ca85
3d6ea31
ab2a349
44ed3e6
45fb2cd
d19d3ec
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,21 +1,20 @@ | ||
# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
# SPDX-License-Identifier: Apache-2.0 | ||
|
||
import functools | ||
import glob | ||
import os | ||
from collections.abc import Sequence | ||
from typing import Optional | ||
|
||
from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError | ||
from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import ( | ||
IS_WINDOWS, | ||
SITE_PACKAGES_LIBDIRS_LINUX, | ||
SITE_PACKAGES_LIBDIRS_WINDOWS, | ||
is_suppressed_dll_file, | ||
) | ||
from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path | ||
from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs, find_sub_dirs_all_sitepackages | ||
from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages | ||
from cuda.pathfinder._utils.platform_aware import IS_WINDOWS | ||
|
||
|
||
def _no_such_file_in_sub_dirs( | ||
|
@@ -80,37 +79,50 @@ def _find_dll_using_nvidia_bin_dirs( | |
return None | ||
|
||
|
||
def _find_lib_dir_using_cuda_home(libname: str) -> Optional[str]: | ||
cuda_home = get_cuda_home_or_path() | ||
if cuda_home is None: | ||
return None | ||
subdirs_list: tuple[tuple[str, ...], ...] | ||
def _find_lib_dir_using_anchor_point(libname: str, anchor_point: str, linux_lib_dir: str) -> Optional[str]: | ||
# Resolve paths for the four cases: | ||
# Windows/Linux x nvvm yes/no | ||
if IS_WINDOWS: | ||
if libname == "nvvm": # noqa: SIM108 | ||
subdirs_list = ( | ||
("nvvm", "bin", "*"), # CTK 13 | ||
("nvvm", "bin"), # CTK 12 | ||
) | ||
rel_paths = [ | ||
"nvvm/bin/*", # CTK 13 | ||
"nvvm/bin", # CTK 12 | ||
] | ||
else: | ||
subdirs_list = ( | ||
("bin", "x64"), # CTK 13 | ||
("bin",), # CTK 12 | ||
) | ||
rel_paths = [ | ||
"bin/x64", # CTK 13 | ||
leofang marked this conversation as resolved.
Show resolved
Hide resolved
|
||
"bin", # CTK 12 | ||
] | ||
else: | ||
if libname == "nvvm": # noqa: SIM108 | ||
subdirs_list = (("nvvm", "lib64"),) | ||
rel_paths = ["nvvm/lib64"] | ||
else: | ||
subdirs_list = ( | ||
("lib64",), # CTK | ||
("lib",), # Conda | ||
) | ||
for sub_dirs in subdirs_list: | ||
dirname: str # work around bug in mypy | ||
for dirname in find_sub_dirs((cuda_home,), sub_dirs): | ||
return dirname | ||
rel_paths = [linux_lib_dir] | ||
|
||
for rel_path in rel_paths: | ||
for dirname in sorted(glob.glob(os.path.join(anchor_point, rel_path))): | ||
if os.path.isdir(dirname): | ||
return dirname | ||
|
||
return None | ||
|
||
|
||
def _find_lib_dir_using_cuda_home(libname: str) -> Optional[str]: | ||
cuda_home = get_cuda_home_or_path() | ||
if cuda_home is None: | ||
return None | ||
return _find_lib_dir_using_anchor_point(libname, anchor_point=cuda_home, linux_lib_dir="lib64") | ||
|
||
|
||
def _find_lib_dir_using_conda_prefix(libname: str) -> Optional[str]: | ||
conda_prefix = os.environ.get("CONDA_PREFIX") | ||
if not conda_prefix: | ||
return None | ||
return _find_lib_dir_using_anchor_point( | ||
libname, anchor_point=os.path.join(conda_prefix, "Library") if IS_WINDOWS else conda_prefix, linux_lib_dir="lib" | ||
) | ||
kkraus14 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
||
|
||
def _find_so_using_lib_dir( | ||
lib_dir: str, so_basename: str, error_messages: list[str], attachments: list[str] | ||
) -> Optional[str]: | ||
|
@@ -144,55 +156,55 @@ def _find_dll_using_lib_dir( | |
class _FindNvidiaDynamicLib: | ||
def __init__(self, libname: str): | ||
self.libname = libname | ||
if IS_WINDOWS: | ||
self.lib_searched_for = f"{libname}*.dll" | ||
else: | ||
self.lib_searched_for = f"lib{libname}.so" | ||
self.error_messages: list[str] = [] | ||
self.attachments: list[str] = [] | ||
self.abs_path = None | ||
self.abs_path: Optional[str] = None | ||
|
||
def try_site_packages(self) -> Optional[str]: | ||
if IS_WINDOWS: | ||
self.lib_searched_for = f"{libname}*.dll" | ||
if self.abs_path is None: | ||
self.abs_path = _find_dll_using_nvidia_bin_dirs( | ||
libname, | ||
self.lib_searched_for, | ||
self.error_messages, | ||
self.attachments, | ||
) | ||
return _find_dll_using_nvidia_bin_dirs( | ||
self.libname, | ||
self.lib_searched_for, | ||
self.error_messages, | ||
self.attachments, | ||
) | ||
else: | ||
self.lib_searched_for = f"lib{libname}.so" | ||
if self.abs_path is None: | ||
self.abs_path = _find_so_using_nvidia_lib_dirs( | ||
libname, | ||
self.lib_searched_for, | ||
self.error_messages, | ||
self.attachments, | ||
) | ||
|
||
def retry_with_cuda_home_priority_last(self) -> None: | ||
cuda_home_lib_dir = _find_lib_dir_using_cuda_home(self.libname) | ||
if cuda_home_lib_dir is not None: | ||
if IS_WINDOWS: | ||
self.abs_path = _find_dll_using_lib_dir( | ||
cuda_home_lib_dir, | ||
self.libname, | ||
self.error_messages, | ||
self.attachments, | ||
) | ||
else: | ||
self.abs_path = _find_so_using_lib_dir( | ||
cuda_home_lib_dir, | ||
self.lib_searched_for, | ||
self.error_messages, | ||
self.attachments, | ||
) | ||
|
||
def raise_if_abs_path_is_None(self) -> str: # noqa: N802 | ||
if self.abs_path: | ||
return self.abs_path | ||
return _find_so_using_nvidia_lib_dirs( | ||
self.libname, | ||
self.lib_searched_for, | ||
self.error_messages, | ||
self.attachments, | ||
) | ||
|
||
def try_with_conda_prefix(self) -> Optional[str]: | ||
return self._find_using_lib_dir(_find_lib_dir_using_conda_prefix(self.libname)) | ||
|
||
def try_with_cuda_home(self) -> Optional[str]: | ||
return self._find_using_lib_dir(_find_lib_dir_using_cuda_home(self.libname)) | ||
|
||
def _find_using_lib_dir(self, lib_dir: Optional[str]) -> Optional[str]: | ||
if lib_dir is None: | ||
return None | ||
if IS_WINDOWS: | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I know this is pre-existing, but it seems like this is begging for some sort of interface that you can isolate into modules and then import at the top level based on platform, similar to how to the There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I tried this twice before but backed off both times after seeing where it's going, and determining it's a net loss in terms of code organization and readability. In contrast, it was clearly a net win for the rest of the code underneath |
||
return _find_dll_using_lib_dir( | ||
lib_dir, | ||
self.libname, | ||
self.error_messages, | ||
self.attachments, | ||
) | ||
else: | ||
return _find_so_using_lib_dir( | ||
lib_dir, | ||
self.lib_searched_for, | ||
self.error_messages, | ||
self.attachments, | ||
) | ||
|
||
def raise_not_found_error(self) -> None: | ||
err = ", ".join(self.error_messages) | ||
att = "\n".join(self.attachments) | ||
raise DynamicLibNotFoundError(f'Failure finding "{self.lib_searched_for}": {err}\n{att}') | ||
|
||
|
||
@functools.cache | ||
def find_nvidia_dynamic_lib(libname: str) -> str: | ||
return _FindNvidiaDynamicLib(libname).raise_if_abs_path_is_None() |
leofang marked this conversation as resolved.
Show resolved
Hide resolved
|
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,18 @@ | ||
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
# SPDX-License-Identifier: Apache-2.0 | ||
|
||
import sys | ||
|
||
IS_WINDOWS = sys.platform == "win32" | ||
|
||
|
||
def quote_for_shell(s: str) -> str: | ||
if IS_WINDOWS: | ||
# This is a relatively heavy import; keep pathfinder if possible. | ||
from subprocess import list2cmdline # nosec B404 | ||
|
||
return list2cmdline([s]) | ||
else: | ||
import shlex | ||
|
||
return shlex.quote(s) |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,4 +1,4 @@ | ||
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
# SPDX-License-Identifier: Apache-2.0 | ||
|
||
__version__ = "1.2.3" | ||
__version__ = "1.3.0" |
Uh oh!
There was an error while loading. Please reload this page.