Skip to content

Commit 1137e15

Browse files
rwgkAndy-Jost
andauthored
Support non-CTK Nvidia libraries, add general fallback for unsupported libs under site-packages (#864)
* Add find_site_packages_so.py * Use find_all_so_files_under_all_site_packages() from _find_so_using_nvidia_lib_dirs() * Bump cuda-pathfinder version to `1.1.1a3` * Limit site-packages search to ("nvidia", "nvpl") subdirs. * Replace find_all_so_files_under_all_site_packages → find_all_so_files_via_metadata * Add find_site_packages_dll.py and use from _find_dll_using_nvidia_bin_dirs() * Add mathdx, cufftMp DIRECT_DEPENDENCIES * Add LIBNAMES_REQUIRING_RTLD_DEEPBIND feature (for cufftMp) * pyproject.toml: add libmathdx, cufftmpm nvshmem, nvpl-fft wheels for testing. * Add SITE_PACKAGES_LIBDIRS_LINUX * Add make_site_packages_libdirs_linux.py * Use SITE_PACKAGES_LIBDIRS_LINUX in _find_so_using_nvidia_lib_dirs, keep find_all_so_files_via_metadata as fallback * Add SITE_PACKAGES_LIBDIRS_WINDOWS and toolshed/make_site_packages_libdirs_windows.py * chmod 755 make_site_packages_libdirs_windows.py * Adds paths for the CUDA static library based on CUDA_HOME (#608). * Removes LIB and LIBRARY_PATH environment variables from the build-wheel workflow. * Updates Linux install to search both lib and lib64 directories for CUDA libraries. * Removes LIBRARY_PATH environment variable from installation docs (no longer needed due to resolution of #608). * Use SITE_PACKAGES_LIBDIRS_WINDOWS in _find_dll_using_nvidia_bin_dirs, keep find_all_dll_files_via_metadata as fallback * Factor out SITE_PACKAGES_LIBDIRS_*_CTK, add test_supported_libnames_*_site_packages_libdirs_ctk_consistency * Also exercise "other" (non-CTK) libnames in test_load_nvidia_dynamic_lib.py Factor out tests/child_load_nvidia_dynamic_lib_helper.py to significantly improve performance. * Exercise fallback code path using pygit2 wheel. * Add other_wheels,foreign_wheels to pip install nvidia_wheels_cu13 * Add toolshed/collect_site_packages_so_files.sh, with terse Usage comment. * Add toolshed/collect_site_packages_dll_files.ps1 with terse Usage comment. * Add pygit2 comments. * Replace special-case workaround in tests/child_load_nvidia_dynamic_lib_helper.py with a more general approach. * Add anticipated CTK 13 paths for mathdx in SITE_PACKAGES_LIBDIRS_LINUX_OTHER, SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER * Rename other_wheels → nvidia_wheels_host * WIP * Restore _no_such_file_in_sub_dirs error reporting * Use `pip install -v ".[nvidia_wheels_cu${TEST_CUDA_MAJOR},nvidia_wheels_host,foreign_wheels]"` to also test pathfinder with `nvidia_wheels_cu12` * Export TEST_CUDA_MAJOR to the GITHUB_ENV * Fix existing (on main) pre-commit error * Do not install nvidia-cufftmp-cu12 on Windows (it is only a placeholder package) * Leo's --only-binary=:all: suggestions * Leo's --only-binary=:all: suggestions (toolshed scripts) * Remove fallback code paths in _find_so_using_nvidia_lib_dirs, _find_dll_using_nvidia_bin_dirs and associated foreign_wheels unit test * Consolidate make_site_packages_libdirs_linux.py + make_site_packages_libdirs_windows.py → make_site_packages_libdirs.py --------- Co-authored-by: Andy Jost <[email protected]>
1 parent 333367e commit 1137e15

15 files changed

+494
-84
lines changed

.github/workflows/test-wheel-linux.yml

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -321,16 +321,15 @@ jobs:
321321
pip install $(ls cuda_python*.whl)[all]
322322
fi
323323
324-
- name: Install cuda.pathfinder nvidia_wheels_cu13
325-
if: startsWith(matrix.CUDA_VER, '13.')
324+
- name: Install cuda.pathfinder extra wheels for testing
326325
run: |
326+
set -euo pipefail
327327
pushd cuda_pathfinder
328-
pip install -v .[nvidia_wheels_cu13]
329-
pip freeze
328+
pip install --only-binary=:all: -v ".[nvidia_wheels_cu${TEST_CUDA_MAJOR},nvidia_wheels_host]"
329+
pip list
330330
popd
331331
332332
- name: Run cuda.pathfinder tests with all_must_work
333-
if: startsWith(matrix.CUDA_VER, '13.')
334333
env:
335334
CUDA_PATHFINDER_TEST_LOAD_NVIDIA_DYNAMIC_LIB_STRICTNESS: all_must_work
336335
run: run-tests pathfinder

.github/workflows/test-wheel-windows.yml

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -288,17 +288,15 @@ jobs:
288288
pip install "$((Get-ChildItem -Filter cuda_python*.whl).FullName)[all]"
289289
}
290290
291-
- name: Install cuda.pathfinder nvidia_wheels_cu13
292-
if: startsWith(matrix.CUDA_VER, '13.')
291+
- name: Install cuda.pathfinder extra wheels for testing
293292
shell: bash --noprofile --norc -xeuo pipefail {0}
294293
run: |
295294
pushd cuda_pathfinder
296-
pip install -v .[nvidia_wheels_cu13]
297-
pip freeze
295+
pip install --only-binary=:all: -v ".[nvidia_wheels_cu${TEST_CUDA_MAJOR},nvidia_wheels_host]"
296+
pip list
298297
popd
299298
300299
- name: Run cuda.pathfinder tests with all_must_work
301-
if: startsWith(matrix.CUDA_VER, '13.')
302300
env:
303301
CUDA_PATHFINDER_TEST_LOAD_NVIDIA_DYNAMIC_LIB_STRICTNESS: all_must_work
304302
shell: bash --noprofile --norc -xeuo pipefail {0}

ci/tools/env-vars

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ elif [[ "${1}" == "test" ]]; then
6969
echo "SETUP_SANITIZER=${SETUP_SANITIZER}" >> $GITHUB_ENV
7070
echo "SKIP_CUDA_BINDINGS_TEST=${SKIP_CUDA_BINDINGS_TEST}" >> $GITHUB_ENV
7171
echo "SKIP_CYTHON_TEST=${SKIP_CYTHON_TEST}" >> $GITHUB_ENV
72+
echo "TEST_CUDA_MAJOR=${TEST_CUDA_MAJOR}" >> $GITHUB_ENV
7273
fi
7374

7475
echo "CUDA_BINDINGS_ARTIFACT_BASENAME=${CUDA_BINDINGS_ARTIFACT_BASENAME}" >> $GITHUB_ENV

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/find_nvidia_dynamic_lib.py

Lines changed: 31 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@
1010
from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError
1111
from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import (
1212
IS_WINDOWS,
13+
SITE_PACKAGES_LIBDIRS_LINUX,
14+
SITE_PACKAGES_LIBDIRS_WINDOWS,
1315
is_suppressed_dll_file,
1416
)
1517
from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs, find_sub_dirs_all_sitepackages
@@ -28,22 +30,25 @@ def _no_such_file_in_sub_dirs(
2830
def _find_so_using_nvidia_lib_dirs(
2931
libname: str, so_basename: str, error_messages: list[str], attachments: list[str]
3032
) -> Optional[str]:
31-
file_wild = so_basename + "*"
32-
nvidia_sub_dirs_list: list[tuple[str, ...]] = [("nvidia", "*", "lib")] # works also for CTK 13 nvvm
33-
if libname == "nvvm":
34-
nvidia_sub_dirs_list.append(("nvidia", "*", "nvvm", "lib64")) # CTK 12
35-
for nvidia_sub_dirs in nvidia_sub_dirs_list:
36-
for lib_dir in find_sub_dirs_all_sitepackages(nvidia_sub_dirs):
37-
# First look for an exact match
38-
so_name = os.path.join(lib_dir, so_basename)
39-
if os.path.isfile(so_name):
40-
return so_name
41-
# Look for a versioned library
42-
# Using sort here mainly to make the result deterministic.
43-
for so_name in sorted(glob.glob(os.path.join(lib_dir, file_wild))):
33+
rel_dirs = SITE_PACKAGES_LIBDIRS_LINUX.get(libname)
34+
if rel_dirs is not None:
35+
sub_dirs_searched = []
36+
file_wild = so_basename + "*"
37+
for rel_dir in rel_dirs:
38+
sub_dir = tuple(rel_dir.split(os.path.sep))
39+
for abs_dir in find_sub_dirs_all_sitepackages(sub_dir):
40+
# First look for an exact match
41+
so_name = os.path.join(abs_dir, so_basename)
4442
if os.path.isfile(so_name):
4543
return so_name
46-
_no_such_file_in_sub_dirs(nvidia_sub_dirs, file_wild, error_messages, attachments)
44+
# Look for a versioned library
45+
# Using sort here mainly to make the result deterministic.
46+
for so_name in sorted(glob.glob(os.path.join(abs_dir, file_wild))):
47+
if os.path.isfile(so_name):
48+
return so_name
49+
sub_dirs_searched.append(sub_dir)
50+
for sub_dir in sub_dirs_searched:
51+
_no_such_file_in_sub_dirs(sub_dir, file_wild, error_messages, attachments)
4752
return None
4853

4954

@@ -59,18 +64,18 @@ def _find_dll_under_dir(dirpath: str, file_wild: str) -> Optional[str]:
5964
def _find_dll_using_nvidia_bin_dirs(
6065
libname: str, lib_searched_for: str, error_messages: list[str], attachments: list[str]
6166
) -> Optional[str]:
62-
nvidia_sub_dirs_list: list[tuple[str, ...]] = [
63-
("nvidia", "*", "bin"), # CTK 12
64-
("nvidia", "*", "bin", "*"), # CTK 13, e.g. site-packages\nvidia\cu13\bin\x86_64\
65-
]
66-
if libname == "nvvm":
67-
nvidia_sub_dirs_list.append(("nvidia", "*", "nvvm", "bin")) # Only for CTK 12
68-
for nvidia_sub_dirs in nvidia_sub_dirs_list:
69-
for bin_dir in find_sub_dirs_all_sitepackages(nvidia_sub_dirs):
70-
dll_name = _find_dll_under_dir(bin_dir, lib_searched_for)
71-
if dll_name is not None:
72-
return dll_name
73-
_no_such_file_in_sub_dirs(nvidia_sub_dirs, lib_searched_for, error_messages, attachments)
67+
rel_dirs = SITE_PACKAGES_LIBDIRS_WINDOWS.get(libname)
68+
if rel_dirs is not None:
69+
sub_dirs_searched = []
70+
for rel_dir in rel_dirs:
71+
sub_dir = tuple(rel_dir.split(os.path.sep))
72+
for abs_dir in find_sub_dirs_all_sitepackages(sub_dir):
73+
dll_name = _find_dll_under_dir(abs_dir, lib_searched_for)
74+
if dll_name is not None:
75+
return dll_name
76+
sub_dirs_searched.append(sub_dir)
77+
for sub_dir in sub_dirs_searched:
78+
_no_such_file_in_sub_dirs(sub_dir, lib_searched_for, error_messages, attachments)
7479
return None
7580

7681

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_linux.py

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,10 @@
88
from typing import Optional, cast
99

1010
from cuda.pathfinder._dynamic_libs.load_dl_common import LoadedDL
11-
from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import SUPPORTED_LINUX_SONAMES
11+
from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import (
12+
LIBNAMES_REQUIRING_RTLD_DEEPBIND,
13+
SUPPORTED_LINUX_SONAMES,
14+
)
1215

1316
CDLL_MODE = os.RTLD_NOW | os.RTLD_GLOBAL
1417

@@ -138,6 +141,13 @@ def check_if_already_loaded_from_elsewhere(libname: str, _have_abs_path: bool) -
138141
return None
139142

140143

144+
def _load_lib(libname: str, filename: str) -> ctypes.CDLL:
145+
cdll_mode = CDLL_MODE
146+
if libname in LIBNAMES_REQUIRING_RTLD_DEEPBIND:
147+
cdll_mode |= os.RTLD_DEEPBIND
148+
return ctypes.CDLL(filename, cdll_mode)
149+
150+
141151
def load_with_system_search(libname: str) -> Optional[LoadedDL]:
142152
"""Try to load a library using system search paths.
143153
@@ -152,13 +162,14 @@ def load_with_system_search(libname: str) -> Optional[LoadedDL]:
152162
"""
153163
for soname in get_candidate_sonames(libname):
154164
try:
155-
handle = ctypes.CDLL(soname, CDLL_MODE)
165+
handle = _load_lib(libname, soname)
166+
except OSError:
167+
pass
168+
else:
156169
abs_path = abs_path_for_dynamic_library(libname, handle)
157170
if abs_path is None:
158171
raise RuntimeError(f"No expected symbol for {libname=!r}")
159172
return LoadedDL(abs_path, False, handle._handle)
160-
except OSError:
161-
pass
162173
return None
163174

164175

@@ -196,7 +207,7 @@ def load_with_abs_path(libname: str, found_path: str) -> LoadedDL:
196207
"""
197208
_work_around_known_bugs(libname, found_path)
198209
try:
199-
handle = ctypes.CDLL(found_path, CDLL_MODE)
210+
handle = _load_lib(libname, found_path)
200211
except OSError as e:
201212
raise RuntimeError(f"Failed to dlopen {found_path}: {e}") from e
202213
return LoadedDL(found_path, False, handle._handle)

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/supported_nvidia_libs.py

Lines changed: 93 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@
6363
SUPPORTED_LIBNAMES = SUPPORTED_LIBNAMES_WINDOWS if IS_WINDOWS else SUPPORTED_LIBNAMES_LINUX
6464

6565
# Based on ldd output for Linux x86_64 nvidia-*-cu12 wheels (12.8.1)
66-
DIRECT_DEPENDENCIES = {
66+
DIRECT_DEPENDENCIES_CTK = {
6767
"cublas": ("cublasLt",),
6868
"cufftw": ("cufft",),
6969
# "cufile_rdma": ("cufile",),
@@ -82,6 +82,10 @@
8282
"npps": ("nppc",),
8383
"nvblas": ("cublas", "cublasLt"),
8484
}
85+
DIRECT_DEPENDENCIES = DIRECT_DEPENDENCIES_CTK | {
86+
"mathdx": ("nvrtc",),
87+
"cufftMp": ("nvshmem_host",),
88+
}
8589

8690
# Based on these released files:
8791
# cuda_11.0.3_450.51.06_linux.run
@@ -104,7 +108,7 @@
104108
# cuda_12.9.1_575.57.08_linux.run
105109
# cuda_13.0.0_580.65.06_linux.run
106110
# Generated with toolshed/build_pathfinder_sonames.py
107-
SUPPORTED_LINUX_SONAMES = {
111+
SUPPORTED_LINUX_SONAMES_CTK = {
108112
"cublas": (
109113
"libcublas.so.11",
110114
"libcublas.so.12",
@@ -232,6 +236,13 @@
232236
"libnvvm.so.4",
233237
),
234238
}
239+
SUPPORTED_LINUX_SONAMES_OTHER = {
240+
"cufftMp": ("libcufftMp.so.11",),
241+
"mathdx": ("libmathdx.so.0",),
242+
"nvpl_fftw": ("libnvpl_fftw.so.0",),
243+
"nvshmem_host": ("libnvshmem_host.so.3",),
244+
}
245+
SUPPORTED_LINUX_SONAMES = SUPPORTED_LINUX_SONAMES_CTK | SUPPORTED_LINUX_SONAMES_OTHER
235246

236247
# Based on these released files:
237248
# cuda_11.0.3_451.82_win10.exe
@@ -254,7 +265,7 @@
254265
# cuda_12.9.1_576.57_windows.exe
255266
# cuda_13.0.0_windows.exe
256267
# Generated with toolshed/build_pathfinder_dlls.py
257-
SUPPORTED_WINDOWS_DLLS = {
268+
SUPPORTED_WINDOWS_DLLS_CTK = {
258269
"cublas": (
259270
"cublas64_11.dll",
260271
"cublas64_12.dll",
@@ -384,12 +395,91 @@
384395
"nvvm70.dll",
385396
),
386397
}
398+
SUPPORTED_WINDOWS_DLLS_OTHER = {
399+
"mathdx": ("mathdx64_0.dll",),
400+
}
401+
SUPPORTED_WINDOWS_DLLS = SUPPORTED_WINDOWS_DLLS_CTK | SUPPORTED_WINDOWS_DLLS_OTHER
387402

388403
LIBNAMES_REQUIRING_OS_ADD_DLL_DIRECTORY = (
389404
"cufft",
390405
"nvrtc",
391406
)
392407

408+
LIBNAMES_REQUIRING_RTLD_DEEPBIND = ("cufftMp",)
409+
410+
# Based on output of toolshed/make_site_packages_libdirs_linux.py
411+
SITE_PACKAGES_LIBDIRS_LINUX_CTK = {
412+
"cublas": ("nvidia/cu13/lib", "nvidia/cublas/lib"),
413+
"cublasLt": ("nvidia/cu13/lib", "nvidia/cublas/lib"),
414+
"cudart": ("nvidia/cu13/lib", "nvidia/cuda_runtime/lib"),
415+
"cufft": ("nvidia/cu13/lib", "nvidia/cufft/lib"),
416+
"cufftw": ("nvidia/cu13/lib", "nvidia/cufft/lib"),
417+
"cufile": ("nvidia/cu13/lib", "nvidia/cufile/lib"),
418+
# "cufile_rdma": ("nvidia/cu13/lib", "nvidia/cufile/lib"),
419+
"curand": ("nvidia/cu13/lib", "nvidia/curand/lib"),
420+
"cusolver": ("nvidia/cu13/lib", "nvidia/cusolver/lib"),
421+
"cusolverMg": ("nvidia/cu13/lib", "nvidia/cusolver/lib"),
422+
"cusparse": ("nvidia/cu13/lib", "nvidia/cusparse/lib"),
423+
"nppc": ("nvidia/cu13/lib", "nvidia/npp/lib"),
424+
"nppial": ("nvidia/cu13/lib", "nvidia/npp/lib"),
425+
"nppicc": ("nvidia/cu13/lib", "nvidia/npp/lib"),
426+
"nppidei": ("nvidia/cu13/lib", "nvidia/npp/lib"),
427+
"nppif": ("nvidia/cu13/lib", "nvidia/npp/lib"),
428+
"nppig": ("nvidia/cu13/lib", "nvidia/npp/lib"),
429+
"nppim": ("nvidia/cu13/lib", "nvidia/npp/lib"),
430+
"nppist": ("nvidia/cu13/lib", "nvidia/npp/lib"),
431+
"nppisu": ("nvidia/cu13/lib", "nvidia/npp/lib"),
432+
"nppitc": ("nvidia/cu13/lib", "nvidia/npp/lib"),
433+
"npps": ("nvidia/cu13/lib", "nvidia/npp/lib"),
434+
"nvJitLink": ("nvidia/cu13/lib", "nvidia/nvjitlink/lib"),
435+
"nvblas": ("nvidia/cu13/lib", "nvidia/cublas/lib"),
436+
"nvfatbin": ("nvidia/cu13/lib", "nvidia/nvfatbin/lib"),
437+
"nvjpeg": ("nvidia/cu13/lib", "nvidia/nvjpeg/lib"),
438+
"nvrtc": ("nvidia/cu13/lib", "nvidia/cuda_nvrtc/lib"),
439+
"nvvm": ("nvidia/cu13/lib", "nvidia/cuda_nvcc/nvvm/lib64"),
440+
}
441+
SITE_PACKAGES_LIBDIRS_LINUX_OTHER = {
442+
"cufftMp": ("nvidia/cufftmp/cu12/lib",),
443+
"mathdx": ("nvidia/cu13/lib", "nvidia/cu12/lib"),
444+
"nvpl_fftw": ("nvpl/lib",),
445+
"nvshmem_host": ("nvidia/nvshmem/lib",),
446+
}
447+
SITE_PACKAGES_LIBDIRS_LINUX = SITE_PACKAGES_LIBDIRS_LINUX_CTK | SITE_PACKAGES_LIBDIRS_LINUX_OTHER
448+
449+
# Based on output of toolshed/make_site_packages_libdirs_windows.py
450+
SITE_PACKAGES_LIBDIRS_WINDOWS_CTK = {
451+
"cublas": ("nvidia/cu13/bin/x86_64", "nvidia/cublas/bin"),
452+
"cublasLt": ("nvidia/cu13/bin/x86_64", "nvidia/cublas/bin"),
453+
"cudart": ("nvidia/cu13/bin/x86_64", "nvidia/cuda_runtime/bin"),
454+
"cufft": ("nvidia/cu13/bin/x86_64", "nvidia/cufft/bin"),
455+
"cufftw": ("nvidia/cu13/bin/x86_64", "nvidia/cufft/bin"),
456+
"curand": ("nvidia/cu13/bin/x86_64", "nvidia/curand/bin"),
457+
"cusolver": ("nvidia/cu13/bin/x86_64", "nvidia/cusolver/bin"),
458+
"cusolverMg": ("nvidia/cu13/bin/x86_64", "nvidia/cusolver/bin"),
459+
"cusparse": ("nvidia/cu13/bin/x86_64", "nvidia/cusparse/bin"),
460+
"nppc": ("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"),
461+
"nppial": ("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"),
462+
"nppicc": ("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"),
463+
"nppidei": ("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"),
464+
"nppif": ("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"),
465+
"nppig": ("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"),
466+
"nppim": ("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"),
467+
"nppist": ("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"),
468+
"nppisu": ("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"),
469+
"nppitc": ("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"),
470+
"npps": ("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"),
471+
"nvJitLink": ("nvidia/cu13/bin/x86_64", "nvidia/nvjitlink/bin"),
472+
"nvblas": ("nvidia/cu13/bin/x86_64", "nvidia/cublas/bin"),
473+
"nvfatbin": ("nvidia/cu13/bin/x86_64", "nvidia/nvfatbin/bin"),
474+
"nvjpeg": ("nvidia/cu13/bin/x86_64", "nvidia/nvjpeg/bin"),
475+
"nvrtc": ("nvidia/cu13/bin/x86_64", "nvidia/cuda_nvrtc/bin"),
476+
"nvvm": ("nvidia/cu13/bin/x86_64", "nvidia/cuda_nvcc/nvvm/bin"),
477+
}
478+
SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER = {
479+
"mathdx": ("nvidia/cu13/bin/x86_64", "nvidia/cu12/bin"),
480+
}
481+
SITE_PACKAGES_LIBDIRS_WINDOWS = SITE_PACKAGES_LIBDIRS_WINDOWS_CTK | SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER
482+
393483

394484
def is_suppressed_dll_file(path_basename: str) -> bool:
395485
if path_basename.startswith("nvrtc"):
Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
import collections
5+
import functools
6+
import importlib.metadata
7+
8+
9+
@functools.cache
10+
def find_all_dll_files_via_metadata() -> dict[str, tuple[str, ...]]:
11+
results: collections.defaultdict[str, list[str]] = collections.defaultdict(list)
12+
13+
# sort dists for deterministic output
14+
for dist in sorted(importlib.metadata.distributions(), key=lambda d: (d.metadata.get("Name", ""), d.version)):
15+
files = dist.files
16+
if not files:
17+
continue
18+
for relpath in sorted(files, key=lambda p: str(p)): # deterministic
19+
relname = relpath.name.lower()
20+
if not relname.endswith(".dll"):
21+
continue
22+
abs_path = str(dist.locate_file(relpath))
23+
results[relname].append(abs_path)
24+
25+
# plain dicts; sort inner list for stability
26+
return {k: tuple(sorted(v)) for k, v in results.items()}
Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
import collections
5+
import functools
6+
import importlib.metadata
7+
import re
8+
9+
_SO_RE = re.compile(r"\.so(?:$|\.)") # matches libfoo.so or libfoo.so.1.2.3
10+
11+
12+
def split_so_version_suffix(so_filename: str) -> tuple[str, str]:
13+
idx = so_filename.rfind(".so")
14+
assert idx > 0, so_filename
15+
idx += 3
16+
return (so_filename[:idx], so_filename[idx:])
17+
18+
19+
@functools.cache
20+
def find_all_so_files_via_metadata() -> dict[str, dict[str, tuple[str, ...]]]:
21+
results: collections.defaultdict[str, collections.defaultdict[str, list[str]]] = collections.defaultdict(
22+
lambda: collections.defaultdict(list)
23+
)
24+
25+
# sort dists for deterministic output
26+
for dist in sorted(importlib.metadata.distributions(), key=lambda d: (d.metadata.get("Name", ""), d.version)):
27+
files = dist.files
28+
if not files:
29+
continue
30+
for relpath in sorted(files, key=lambda p: str(p)): # deterministic
31+
relname = relpath.name
32+
if not _SO_RE.search(relname):
33+
continue
34+
so_basename, so_version_suffix = split_so_version_suffix(relname)
35+
abs_path = str(dist.locate_file(relpath))
36+
results[so_basename][so_version_suffix].append(abs_path)
37+
38+
# plain dicts; sort inner lists for stability
39+
return {k: {kk: tuple(sorted(vv)) for kk, vv in v.items()} for k, v in results.items()}
Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
11
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
# SPDX-License-Identifier: Apache-2.0
33

4-
__version__ = "1.1.1a2"
4+
__version__ = "1.1.1a3"

0 commit comments

Comments
 (0)