Add extra requires for cuda/cudnn DLLs to onnxruntime-gpu python package

setup.py

@@ -54,6 +54,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
+is_cuda_version_12 = False
 rocm_version = None
 is_migraphx = False
 is_rocm = False
@@ -63,6 +64,8 @@ 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_cuda_version_12 = cuda_version.startswith("12.")
 elif parse_arg_remove_boolean(sys.argv, "--use_rocm"):
     is_rocm = True
     rocm_version = parse_arg_remove_string(sys.argv, "--rocm_version=")
@@ -721,7 +724,6 @@ def reformat_run_count(count_str):
 with open(requirements_path) as f:
     install_requires = f.read().splitlines()
 
-
 if enable_training:
 
     def save_build_and_package_info(package_name, version_number, cuda_version, rocm_version):
@@ -754,6 +756,20 @@ def save_build_and_package_info(package_name, version_number, cuda_version, rocm
 
     save_build_and_package_info(package_name, version_number, cuda_version, rocm_version)
 
+extras_require = {}
+if package_name == "onnxruntime-gpu" and is_cuda_version_12:
+    extras_require = {
+        "cuda": [
+            "nvidia-cuda-nvrtc-cu12~=12.0",
+            "nvidia-cuda-runtime-cu12~=12.0",
+            "nvidia-cufft-cu12~=11.0",
+            "nvidia-curand-cu12~=10.0",
+        ],
+        "cudnn": [
+            "nvidia-cudnn-cu12~=9.0",
+        ],
+    }
+
 # Setup
 setup(
     name=package_name,
@@ -771,6 +787,7 @@ def save_build_and_package_info(package_name, version_number, cuda_version, rocm
     download_url="https://github.com/microsoft/onnxruntime/tags",
     data_files=data_files,
     install_requires=install_requires,
+    extras_require=extras_require,
     python_requires=">=3.10",
     keywords="onnx machine learning",
     entry_points={

tools/ci_build/build.py

@@ -2366,11 +2366,32 @@ def run_nodejs_tests(nodejs_binding_dir):
     run_subprocess(args, cwd=nodejs_binding_dir)
 
 
+def parse_cuda_version_from_json(cuda_home):
+    version_file_path = os.path.join(cuda_home, "version.json")
+    if not os.path.exists(version_file_path):
+        print(f"version.json not found in {cuda_home}.")
+    else:
+        try:
+            with open(version_file_path) as version_file:
+                version_data = json.load(version_file)
+                cudart_info = version_data.get("cuda")
+                if cudart_info and "version" in cudart_info:
+                    parts = cudart_info["version"].split(".")
+                    return ".".join(parts[:2])
+        except FileNotFoundError:
+            print(f"version.json not found in {cuda_home}.")
+        except json.JSONDecodeError:
+            print(f"Error decoding JSON from version.json in {cuda_home}.")
+
+    return ""
+
+
 def build_python_wheel(
     source_dir,
     build_dir,
     configs,
     use_cuda,
+    cuda_home,
     cuda_version,
     use_rocm,
     use_migraphx,
@@ -2418,6 +2439,7 @@ def build_python_wheel(
         if use_cuda:
             # The following line assumes no other EP is enabled
             args.append("--wheel_name_suffix=gpu")
+            cuda_version = cuda_version or parse_cuda_version_from_json(cuda_home)
             if cuda_version:
                 args.append(f"--cuda_version={cuda_version}")
         elif use_rocm:
@@ -3075,6 +3097,7 @@ def main():
                 build_dir,
                 configs,
                 args.use_cuda,
+                cuda_home,
                 args.cuda_version,
                 args.use_rocm,
                 args.use_migraphx,

tools/ci_build/github/azure-pipelines/stages/py-gpu-packaging-stage.yml

@@ -56,7 +56,7 @@ stages:
           PYTHON_VERSION: ${{ python_version }}
           EP_NAME: gpu
           CudaVersion: ${{ parameters.cuda_version }}
-          EP_BUILD_FLAGS: --enable_lto --cuda_home=$(Agent.TempDirectory)\v${{ parameters.cuda_version }} --cmake_extra_defines "CMAKE_CUDA_ARCHITECTURES=75;80;90"
+          EP_BUILD_FLAGS: --enable_lto --use_cuda --cuda_home=$(Agent.TempDirectory)\v${{ parameters.cuda_version }} --cmake_extra_defines "CMAKE_CUDA_ARCHITECTURES=75;80;90"
           use_tensorrt: True
 
   - ${{ if eq(parameters.enable_linux_cuda, true) }}:
@@ -80,4 +80,4 @@ stages:
           PYTHON_VERSION: ${{ python_version }}
           EP_BUILD_FLAGS: --use_dml --cmake_extra_defines CMAKE_SYSTEM_VERSION=10.0.18362.0 --enable_wcos
           EP_NAME: directml
-          cmake_build_type: ${{ parameters.cmake_build_type }}
+          cmake_build_type: ${{ parameters.cmake_build_type }}