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,18 @@ 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-cudnn-cu12~=9.0",
+            "nvidia-cufft-cu12~=11.0",
+            "nvidia-curand-cu12~=10.0",
+        ]
+    }
+
 # Setup
 setup(
     name=package_name,
@@ -771,6 +785,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/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_version {{ parameters.cuda_version }} --cuda_home=$(Agent.TempDirectory)\v${{ parameters.cuda_version }} --cudnn_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 }}

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

@@ -33,7 +33,7 @@ parameters:
    - Release
    - RelWithDebInfo
    - MinSizeRel
-   
+
 - name: use_tensorrt
   type: boolean
   default: false
@@ -195,7 +195,6 @@ stages:
             TMPDIR: "$(Agent.TempDirectory)"
 
         - powershell: |
-
             python -m pip uninstall -y onnxruntime onnxruntime-gpu -qq
             Get-ChildItem -Path $(Build.ArtifactStagingDirectory)/*cp${{ replace(parameters.PYTHON_VERSION,'.','') }}*.whl | foreach {pip --disable-pip-version-check install --upgrade $_.fullname tabulate}
             mkdir -p $(Agent.TempDirectory)\ort_test_data