File size: 7,630 Bytes
76f9669 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 | # SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
import glob
import os
from collections.abc import Sequence
from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError
from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import (
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_all_sitepackages
from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
def _no_such_file_in_sub_dirs(
sub_dirs: Sequence[str], file_wild: str, error_messages: list[str], attachments: list[str]
) -> None:
error_messages.append(f"No such file: {file_wild}")
for sub_dir in find_sub_dirs_all_sitepackages(sub_dirs):
attachments.append(f' listdir("{sub_dir}"):')
for node in sorted(os.listdir(sub_dir)):
attachments.append(f" {node}")
def _find_so_using_nvidia_lib_dirs(
libname: str, so_basename: str, error_messages: list[str], attachments: list[str]
) -> str | None:
rel_dirs = SITE_PACKAGES_LIBDIRS_LINUX.get(libname)
if rel_dirs is not None:
sub_dirs_searched = []
file_wild = so_basename + "*"
for rel_dir in rel_dirs:
sub_dir = tuple(rel_dir.split(os.path.sep))
for abs_dir in find_sub_dirs_all_sitepackages(sub_dir):
# First look for an exact match
so_name = os.path.join(abs_dir, so_basename)
if os.path.isfile(so_name):
return so_name
# Look for a versioned library
# Using sort here mainly to make the result deterministic.
for so_name in sorted(glob.glob(os.path.join(abs_dir, file_wild))):
if os.path.isfile(so_name):
return so_name
sub_dirs_searched.append(sub_dir)
for sub_dir in sub_dirs_searched:
_no_such_file_in_sub_dirs(sub_dir, file_wild, error_messages, attachments)
return None
def _find_dll_under_dir(dirpath: str, file_wild: str) -> str | None:
for path in sorted(glob.glob(os.path.join(dirpath, file_wild))):
if not os.path.isfile(path):
continue
if not is_suppressed_dll_file(os.path.basename(path)):
return path
return None
def _find_dll_using_nvidia_bin_dirs(
libname: str, lib_searched_for: str, error_messages: list[str], attachments: list[str]
) -> str | None:
rel_dirs = SITE_PACKAGES_LIBDIRS_WINDOWS.get(libname)
if rel_dirs is not None:
sub_dirs_searched = []
for rel_dir in rel_dirs:
sub_dir = tuple(rel_dir.split(os.path.sep))
for abs_dir in find_sub_dirs_all_sitepackages(sub_dir):
dll_name = _find_dll_under_dir(abs_dir, lib_searched_for)
if dll_name is not None:
return dll_name
sub_dirs_searched.append(sub_dir)
for sub_dir in sub_dirs_searched:
_no_such_file_in_sub_dirs(sub_dir, lib_searched_for, error_messages, attachments)
return None
def _find_lib_dir_using_anchor_point(libname: str, anchor_point: str, linux_lib_dir: str) -> str | None:
# Resolve paths for the four cases:
# Windows/Linux x nvvm yes/no
if IS_WINDOWS:
if libname == "nvvm": # noqa: SIM108
rel_paths = [
"nvvm/bin/*", # CTK 13
"nvvm/bin", # CTK 12
]
else:
rel_paths = [
"bin/x64", # CTK 13
"bin", # CTK 12
]
else:
if libname == "nvvm": # noqa: SIM108
rel_paths = ["nvvm/lib64"]
else:
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) -> str | None:
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) -> str | None:
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"
)
def _find_so_using_lib_dir(
lib_dir: str, so_basename: str, error_messages: list[str], attachments: list[str]
) -> str | None:
so_name = os.path.join(lib_dir, so_basename)
if os.path.isfile(so_name):
return so_name
error_messages.append(f"No such file: {so_name}")
attachments.append(f' listdir("{lib_dir}"):')
if not os.path.isdir(lib_dir):
attachments.append(" DIRECTORY DOES NOT EXIST")
else:
for node in sorted(os.listdir(lib_dir)):
attachments.append(f" {node}")
return None
def _find_dll_using_lib_dir(
lib_dir: str, libname: str, error_messages: list[str], attachments: list[str]
) -> str | None:
file_wild = libname + "*.dll"
dll_name = _find_dll_under_dir(lib_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("{lib_dir}"):')
for node in sorted(os.listdir(lib_dir)):
attachments.append(f" {node}")
return None
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: str | None = None
def try_site_packages(self) -> str | None:
if IS_WINDOWS:
return _find_dll_using_nvidia_bin_dirs(
self.libname,
self.lib_searched_for,
self.error_messages,
self.attachments,
)
else:
return _find_so_using_nvidia_lib_dirs(
self.libname,
self.lib_searched_for,
self.error_messages,
self.attachments,
)
def try_with_conda_prefix(self) -> str | None:
return self._find_using_lib_dir(_find_lib_dir_using_conda_prefix(self.libname))
def try_with_cuda_home(self) -> str | None:
return self._find_using_lib_dir(_find_lib_dir_using_cuda_home(self.libname))
def _find_using_lib_dir(self, lib_dir: str | None) -> str | None:
if lib_dir is None:
return None
if IS_WINDOWS:
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}')
|