Adding optional CUDA DLLs when installing onnxruntime_gpu

onnxruntime/__init__.py

@@ -76,3 +76,80 @@
     __version__ = version
 
 onnxruntime_validation.check_distro_info()
+
+# Load nvidia libraries from site-packages/nvidia if the package is onnxruntime-gpu
+if (
+    __package__ == "onnxruntime-gpu"
+    # Just in case we rename the package name in the future
+    or __package__ == "onnxruntime-cuda"
+    or __package__ == "onnxruntime_gpu"
+    or __package__ == "onnxruntime_cuda"
+):
+    import ctypes
+    import os
+    import platform
+    import re
+    import site
+    import logging
+    # Get the site-packages path where nvidia packages are installed
+    site_packages_path = site.getsitepackages()[-1]
+    nvidia_path = os.path.join(site_packages_path, "nvidia")
+    # Traverse the directory and subdirectories
+    if platform.system() == "Windows":  #
+        # Define the list of DLL patterns, curand and nvJitLink are not included for Windows
+        cuda_libs = (
+            "cublas",
+            "cublasLt",
+            "cudnn",
+            "cudart",
+            "cufft",
+            # "curand",
+            # "nvJitLink",
+        )
+        # Construct a regex pattern for each library name with optional parts
+        # Pattern explanation:
+        # - `libname`: Match the base library name (e.g., "cudart")
+        # - `(?:64)?`: Optionally match "64"
+        # - `(?:_\d+)*`: Match zero or more occurrences of "_n" where "n" is one or more digits
+        # - `.dll$`: End with ".dll" ignoring case
+        lib_pattern = {lib: re.compile(rf"{lib}(?:64)?(?:_\d+)*\.dll$", re.IGNORECASE) for lib in cuda_libs}
+        # Collect all directories under site-packages/nvidia that contain .dll files (for Windows)
+        for root, _, files in os.walk(nvidia_path):
+            # Add the current directory to the DLL search path
+
+            with os.add_dll_directory(root):
+                # Find all .dll files in the current directory
+                for file in files:
+                    for pattern in lib_pattern.items().values():
+                        if pattern.match(file):
+                            dll_path = os.path.join(root, file)
+                            try:
+                                _ = ctypes.CDLL(dll_path)
+                            except Exception as e:
+                                logging.error(f"Failed to load {dll_path}: {e}")
+    elif platform.system() == "Linux":
+        # Define the patterns with optional version number and case-insensitivity
+        cuda_libs = (
+            "libcublas.so",
+            "libcublasLt.so",
+            "libcudnn.so",
+            "libcudart.so",
+            "libnvrtc.so",
+            "libcufft.so",
+            "libcurand.so",
+            "libnvJitLink.so",
+        )
+
+        # Regular expression to match .so files with optional versioning (e.g., .so, .so.1, .so.2.3)
+        lib_pattern = {pattern: re.compile(rf"{re.escape(pattern)}(\.\d+)*$", re.IGNORECASE) for pattern in cuda_libs}
+
+        # Traverse the directory and subdirectories
+        for root, _, files in os.walk(nvidia_path):
+            for file in files:
+                # Check if the file matches the .so pattern
+                for regex in lib_pattern.items().values():
+                    if regex.match(file):  # Check if the file matches the pattern
+                        so_path = os.path.join(root, file)
+                        _ = ctypes.CDLL(so_path)
+    else:
+        pass

setup.py

@@ -7,6 +7,7 @@
 import datetime
 import logging
 import platform
+import re
 import shlex
 import subprocess
 import sys
@@ -54,6 +55,7 @@ def parse_arg_remove_string(argv, arg_name_equal):
 wheel_name_suffix = parse_arg_remove_string(sys.argv, "--wheel_name_suffix=")
 
 cuda_version = None
+cuda_version_major = None
 rocm_version = None
 is_migraphx = False
 is_rocm = False
@@ -63,6 +65,11 @@ def parse_arg_remove_string(argv, arg_name_equal):
 if wheel_name_suffix == "gpu":
     # TODO: how to support multiple CUDA versions?
     cuda_version = parse_arg_remove_string(sys.argv, "--cuda_version=")
+    if cuda_version is not None:
+        if not bool(re.match(r"^\d+\.\d+(\.\d+)?$", cuda_version)):
+            logger.error("CUDA version must be in format 'x.y' or  'x.y.z'")
+            sys.exit(1)
+        cuda_version_major = cuda_version.split(".")[0]
 elif parse_arg_remove_boolean(sys.argv, "--use_rocm"):
     is_rocm = True
     rocm_version = parse_arg_remove_string(sys.argv, "--rocm_version=")
@@ -705,11 +712,22 @@ def reformat_run_count(count_str):
     version_number = version_number + local_version
     if is_rocm and enable_rocm_profiling:
         version_number = version_number + ".profiling"
-
+extras_require = {}
 if wheel_name_suffix:
     if not (enable_training and wheel_name_suffix == "gpu"):
         # for training packages, local version is used to indicate device types
         package_name = f"{package_name}-{wheel_name_suffix}"
+    if wheel_name_suffix == "gpu" and cuda_version_major is not None:
+        extras_require = {
+            # Optional 'cuda_dlls' dependencies
+            "cuda_dlls": [
+                f"nvidia-cuda-nvrtc-cu{cuda_version_major}",
+                f"nvidia-cuda-runtime-cu{cuda_version_major}",
+                f"nvidia-cudnn-cu{cuda_version_major}",
+                f"nvidia-cufft-cu{cuda_version_major}",
+                f"nvidia-curand-cu{cuda_version_major}",
+            ]
+        }
 
 cmd_classes = {}
 if bdist_wheel is not None:
@@ -783,4 +801,5 @@ def save_build_and_package_info(package_name, version_number, cuda_version, rocm
         ]
     },
     classifiers=classifiers,
+    extras_require=extras_require,
 )