pytorch

7ef76ce6 - Preloads more nvidia pypi library for multi arch distributions (#94355)

Preloads more nvidia pypi library for multi arch distributions (#94355) Following the same logic of preloading cudnn and cublas from the pypi folder in multi-arch disributions, where Pure-lib vs Plat-lib matters, this PR adds the logic for the rest of the cuda pypi libraries that were integrated. I have tested this PR by running the code block locally and installing/uninstalling nvidia pypi libraries: ``` import sys import os def _preload_cuda_deps(): """Preloads cudnn/cublas deps if they could not be found otherwise.""" # Should only be called on Linux if default path resolution have failed cuda_libs = { 'cublas': 'libcublas.so.11', 'cudnn': 'libcudnn.so.8', 'cuda_nvrtc': 'libnvrtc.so.11.2', 'cuda_runtime': 'libcudart.so.11.0', 'cuda_cupti': 'libcupti.so.11.7', 'cufft': 'libcufft.so.10', 'curand': 'libcurand.so.10', 'cusolver': 'libcusolver.so.11', 'cusparse': 'libcusparse.so.11', 'nccl': 'libnccl.so.2', 'nvtx': 'libnvToolsExt.so.1', } cuda_libs_paths = {lib_folder: None for lib_folder in cuda_libs.keys()} for path in sys.path: nvidia_path = os.path.join(path, 'nvidia') if not os.path.exists(nvidia_path): continue for lib_folder, lib_name in cuda_libs.items(): candidate_path = os.path.join(nvidia_path, lib_folder, 'lib', lib_name) if os.path.exists(candidate_path) and not cuda_libs_paths[lib_folder]: cuda_libs_paths[lib_folder] = candidate_path if all(cuda_libs_paths.values()): break if not all(cuda_libs_paths.values()): none_libs = [lib for lib in cuda_libs_paths if not cuda_libs_paths[lib]] raise ValueError(f"{', '.join(none_libs)} not found in the system path {sys.path}") _preload_cuda_deps() ``` I don't have access to a multi-arch environment, so if somebody could verify a wheel with this patch on a multi-arch distribution, that would be great! Pull Request resolved: https://github.com/pytorch/pytorch/pull/94355 Approved by: https://github.com/atalman