Adding optional CUDA DLLs when installing onnxruntime_gpu

onnxruntime/python/onnxruntime_cuda_temp_env.py

@@ -0,0 +1,57 @@
+import os
+import platform
+import site
+
+
+class TemporaryEnv:
+    def __init__(self, updates):
+        self.original_env = os.environ.copy()
+        os.environ.update(updates)
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        os.environ.clear()
+        os.environ.update(self.original_env)
+
+
+def get_nvidia_dll_paths() -> str:
+    # Get the site-packages path where nvidia packages are installed
+    site_packages_path = site.getsitepackages()[0]
+    nvidia_path = os.path.join(site_packages_path, "nvidia")
+
+    # Collect all directories under site-packages/nvidia that contain .dll files (for Windows)
+    dll_paths = []
+    for root, files in os.walk(nvidia_path):
+        if any(file.endswith(".dll") for file in files):
+            dll_paths.append(root)
+    return os.pathsep.join(dll_paths)
+
+
+def get_nvidia_so_paths() -> str:
+    # Get the site-packages path where nvidia packages are installed
+    site_packages_path = site.getsitepackages()[0]
+    nvidia_path = os.path.join(site_packages_path, "nvidia")
+
+    # Collect all directories under site-packages/nvidia that contain .so files (for Linux)
+    so_paths = []
+    for root, files in os.walk(nvidia_path):
+        if any(file.endswith(".so") for file in files):
+            so_paths.append(root)
+    return os.pathsep.join(so_paths)
+
+
+def setup_temp_env_for_ort_cuda():
+    # Determine platform and set up the environment accordingly
+    if platform.system() == "Windows":  # Windows
+        nvidia_dlls_path = get_nvidia_dll_paths()
+        if nvidia_dlls_path:
+            return TemporaryEnv({"PATH": nvidia_dlls_path + os.pathsep + os.environ.get("PATH")})
+        else:
+            return TemporaryEnv({"PATH": os.environ.get("PATH")})
+    elif platform.system() == "Linux":
+        nvidia_so_paths = get_nvidia_so_paths()
+        if nvidia_so_paths:
+            return TemporaryEnv({"LD_LIBRARY_PATH": nvidia_so_paths + os.pathsep + os.environ.get("LD_LIBRARY_PATH")})
+        else:
+            return TemporaryEnv({"LD_LIBRARY_PATH": os.environ.get("LD_LIBRARY_PATH")})
+    else:
+        return None

onnxruntime/python/onnxruntime_inference_collection.py

@@ -439,6 +439,9 @@ def __init__(
         means execute a node using `CUDAExecutionProvider`
         if capable, otherwise execute using `CPUExecutionProvider`.
         """
+        from .onnxruntime_cuda_temp_env import setup_temp_env_for_ort_cuda
+
+        self.env_manager = setup_temp_env_for_ort_cuda()
         super().__init__()
 
         if isinstance(path_or_bytes, (str, os.PathLike)):
@@ -575,6 +578,11 @@ def _register_ep_custom_ops(self, session_options, providers, provider_options,
             ):
                 C.register_tensorrt_plugins_as_custom_ops(session_options, providers[i][1])
 
+    def __exit__(self, exc_type, exc_value, traceback):
+        if self.env_manager is not None:
+            self.env_manager.__exit__()
+        return False
+
 
 class IOBinding:
     """

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,20 @@ 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-cudnn-cu{cuda_version_major}",
+                f"nvidia-cufft-cu{cuda_version_major}",
+                f"nvidia-cuda-nvrtc-cu{cuda_version_major}",
+            ]
+        }
 
 cmd_classes = {}
 if bdist_wheel is not None:
@@ -783,4 +799,5 @@ def save_build_and_package_info(package_name, version_number, cuda_version, rocm
         ]
     },
     classifiers=classifiers,
+    extras_require=extras_require,
 )