Skip to content

Commit 24a887e

Browse files
tianleiwugithub-actions[bot]
authored andcommitted
Add extra requires for cuda/cudnn DLLs to onnxruntime-gpu python package (#23659)
### Description Add extra requires for cuda/cudnn DLLs to onnxruntime-gpu python package. During building wheel, make sure to add cuda version parameters to build command line like `--cuda_version 12.8`. Note that we only add extra requires for cuda 12 for now. If a package is built with cuda 11, no extra requires will be added. Examples to install extra DLLs from wheel: ``` pip install onnxruntime_gpu-1.21.0-cp310-cp310-linux_x86_64.whl[cuda,cudnn] ``` If install cudnn DLLs but not cuda DLLs: ``` pip install onnxruntime_gpu-1.21.0-cp310-cp310-linux_x86_64.whl[cudnn] ``` Example section in METADATA file of dist-info: ``` Provides-Extra: cuda Requires-Dist: nvidia-cuda-nvrtc-cu12~=12.0; extra == "cuda" Requires-Dist: nvidia-cuda-runtime-cu12~=12.0; extra == "cuda" Requires-Dist: nvidia-cufft-cu12~=11.0; extra == "cuda" Requires-Dist: nvidia-curand-cu12~=10.0; extra == "cuda" Provides-Extra: cudnn Requires-Dist: nvidia-cudnn-cu12~=9.0; extra == "cudnn" ... ``` ### Motivation and Context Jian had a PR: #22506. This adds only part of the change. Extra change include updating the windows gpu python packaging pipeline to pass cuda version to the build command line. --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
1 parent 550c046 commit 24a887e

File tree

3 files changed

+43
-3
lines changed

3 files changed

+43
-3
lines changed

setup.py

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ def parse_arg_remove_string(argv, arg_name_equal):
5454
wheel_name_suffix = parse_arg_remove_string(sys.argv, "--wheel_name_suffix=")
5555

5656
cuda_version = None
57+
is_cuda_version_12 = False
5758
rocm_version = None
5859
is_migraphx = False
5960
is_rocm = False
@@ -63,6 +64,8 @@ def parse_arg_remove_string(argv, arg_name_equal):
6364
if wheel_name_suffix == "gpu":
6465
# TODO: how to support multiple CUDA versions?
6566
cuda_version = parse_arg_remove_string(sys.argv, "--cuda_version=")
67+
if cuda_version:
68+
is_cuda_version_12 = cuda_version.startswith("12.")
6669
elif parse_arg_remove_boolean(sys.argv, "--use_rocm"):
6770
is_rocm = True
6871
rocm_version = parse_arg_remove_string(sys.argv, "--rocm_version=")
@@ -721,7 +724,6 @@ def reformat_run_count(count_str):
721724
with open(requirements_path) as f:
722725
install_requires = f.read().splitlines()
723726

724-
725727
if enable_training:
726728

727729
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
754756

755757
save_build_and_package_info(package_name, version_number, cuda_version, rocm_version)
756758

759+
extras_require = {}
760+
if package_name == "onnxruntime-gpu" and is_cuda_version_12:
761+
extras_require = {
762+
"cuda": [
763+
"nvidia-cuda-nvrtc-cu12~=12.0",
764+
"nvidia-cuda-runtime-cu12~=12.0",
765+
"nvidia-cufft-cu12~=11.0",
766+
"nvidia-curand-cu12~=10.0",
767+
],
768+
"cudnn": [
769+
"nvidia-cudnn-cu12~=9.0",
770+
],
771+
}
772+
757773
# Setup
758774
setup(
759775
name=package_name,
@@ -771,6 +787,7 @@ def save_build_and_package_info(package_name, version_number, cuda_version, rocm
771787
download_url="https://github.com/microsoft/onnxruntime/tags",
772788
data_files=data_files,
773789
install_requires=install_requires,
790+
extras_require=extras_require,
774791
python_requires=">=3.10",
775792
keywords="onnx machine learning",
776793
entry_points={

tools/ci_build/build.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2366,11 +2366,32 @@ def run_nodejs_tests(nodejs_binding_dir):
23662366
run_subprocess(args, cwd=nodejs_binding_dir)
23672367

23682368

2369+
def parse_cuda_version_from_json(cuda_home):
2370+
version_file_path = os.path.join(cuda_home, "version.json")
2371+
if not os.path.exists(version_file_path):
2372+
print(f"version.json not found in {cuda_home}.")
2373+
else:
2374+
try:
2375+
with open(version_file_path) as version_file:
2376+
version_data = json.load(version_file)
2377+
cudart_info = version_data.get("cuda")
2378+
if cudart_info and "version" in cudart_info:
2379+
parts = cudart_info["version"].split(".")
2380+
return ".".join(parts[:2])
2381+
except FileNotFoundError:
2382+
print(f"version.json not found in {cuda_home}.")
2383+
except json.JSONDecodeError:
2384+
print(f"Error decoding JSON from version.json in {cuda_home}.")
2385+
2386+
return ""
2387+
2388+
23692389
def build_python_wheel(
23702390
source_dir,
23712391
build_dir,
23722392
configs,
23732393
use_cuda,
2394+
cuda_home,
23742395
cuda_version,
23752396
use_rocm,
23762397
use_migraphx,
@@ -2418,6 +2439,7 @@ def build_python_wheel(
24182439
if use_cuda:
24192440
# The following line assumes no other EP is enabled
24202441
args.append("--wheel_name_suffix=gpu")
2442+
cuda_version = cuda_version or parse_cuda_version_from_json(cuda_home)
24212443
if cuda_version:
24222444
args.append(f"--cuda_version={cuda_version}")
24232445
elif use_rocm:
@@ -3075,6 +3097,7 @@ def main():
30753097
build_dir,
30763098
configs,
30773099
args.use_cuda,
3100+
cuda_home,
30783101
args.cuda_version,
30793102
args.use_rocm,
30803103
args.use_migraphx,

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ stages:
5656
PYTHON_VERSION: ${{ python_version }}
5757
EP_NAME: gpu
5858
CudaVersion: ${{ parameters.cuda_version }}
59-
EP_BUILD_FLAGS: --enable_lto --cuda_home=$(Agent.TempDirectory)\v${{ parameters.cuda_version }} --cmake_extra_defines "CMAKE_CUDA_ARCHITECTURES=75;80;90"
59+
EP_BUILD_FLAGS: --enable_lto --use_cuda --cuda_home=$(Agent.TempDirectory)\v${{ parameters.cuda_version }} --cmake_extra_defines "CMAKE_CUDA_ARCHITECTURES=75;80;90"
6060
use_tensorrt: True
6161

6262
- ${{ if eq(parameters.enable_linux_cuda, true) }}:
@@ -80,4 +80,4 @@ stages:
8080
PYTHON_VERSION: ${{ python_version }}
8181
EP_BUILD_FLAGS: --use_dml --cmake_extra_defines CMAKE_SYSTEM_VERSION=10.0.18362.0 --enable_wcos
8282
EP_NAME: directml
83-
cmake_build_type: ${{ parameters.cmake_build_type }}
83+
cmake_build_type: ${{ parameters.cmake_build_type }}

0 commit comments

Comments
 (0)