Skip to content

Commit 1acddf2

Browse files
edgchen1edgchen1
authored
Split out Linux CUDA Python package builds into separate stages (#27490)
### Description <!-- Describe your changes. --> Split out Linux CUDA Python package builds into separate stages. ### Motivation and Context <!-- - Why is this change required? What problem does it solve? - If it fixes an open issue, please link to the issue here. --> Reduce overall packaging pipeline time by running Python version builds in separate stages, allowing them to run in parallel. Example build: https://aiinfra.visualstudio.com/Lotus/_build/results?buildId=1102253&view=results Reduced time from ~3h30m to 1h38m.
1 parent 4612613 commit 1acddf2

File tree

3 files changed

+108
-20
lines changed

3 files changed

+108
-20
lines changed

‎tools/ci_build/github/azure-pipelines/build-perf-test-binaries-pipeline.yml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,3 +33,7 @@ stages:
3333
cmake_build_type: Release
3434
cuda_version: 12.8
3535
docker_base_image: onnxruntimebuildcache.azurecr.io/internal/azureml/onnxruntime/build/cuda12_x64_almalinux8_gcc14:20251017.1
36+
python_version: '3.14'
37+
docker_python_exe_path: '/opt/python/cp314-cp314/bin/python3.14'
38+
stage_name: Linux_py_GPU_Wheels_x86_64_3_14
39+
wheel_artifact_name: onnxruntime_gpu_3_14

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

Lines changed: 57 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,23 @@ parameters:
3737
- '3.13'
3838
- '3.14'
3939

40+
- name: LinuxPythonConfigurations
41+
type: object
42+
displayName: 'Linux Python build configurations'
43+
default:
44+
- python_version: '3.11'
45+
docker_python_exe_path: '/opt/python/cp311-cp311/bin/python3.11'
46+
- python_version: '3.12'
47+
docker_python_exe_path: '/opt/python/cp312-cp312/bin/python3.12'
48+
- python_version: '3.13'
49+
docker_python_exe_path: '/opt/python/cp313-cp313/bin/python3.13'
50+
- python_version: '3.13t'
51+
docker_python_exe_path: '/opt/python/cp313-cp313t/bin/python3.13'
52+
- python_version: '3.14'
53+
docker_python_exe_path: '/opt/python/cp314-cp314/bin/python3.14'
54+
- python_version: '3.14t'
55+
docker_python_exe_path: '/opt/python/cp314-cp314t/bin/python3.14'
56+
4057
- name: cmake_cuda_archs
4158
type: string
4259
displayName: 'CMAKE_CUDA_ARCHITECTURES for Windows'
@@ -67,11 +84,49 @@ stages:
6784
EP_BUILD_FLAGS: --enable_lto --use_cuda --cuda_home=$(Agent.TempDirectory)\v${{ parameters.cuda_version }} --cmake_extra_defines "CMAKE_CUDA_ARCHITECTURES=${{ parameters.cmake_cuda_archs }}"
6885
use_tensorrt: True
6986

70-
- template: py-linux-gpu-stage.yml
71-
parameters:
87+
# Linux: one parallel stage per Python version
88+
- ${{ each config in parameters.LinuxPythonConfigurations }}:
89+
- template: py-linux-gpu-stage.yml
90+
parameters:
91+
stage_name: Linux_py_GPU_Wheels_x86_64_${{ replace(config.python_version, '.', '_') }}
7292
arch: 'x86_64'
7393
machine_pool: 'onnxruntime-Ubuntu2404-AMD-CPU'
7494
extra_build_arg: ${{ parameters.build_py_parameters }}
7595
cmake_build_type: ${{ parameters.cmake_build_type }}
7696
cuda_version: ${{ parameters.cuda_version }}
7797
docker_base_image: ${{ parameters.docker_base_image }}
98+
python_version: ${{ config.python_version }}
99+
docker_python_exe_path: ${{ config.docker_python_exe_path }}
100+
wheel_artifact_name: onnxruntime_gpu_${{ replace(config.python_version, '.', '_') }}
101+
# Only publish the large build intermediates artifact for Python 3.12. That's the only version we test now.
102+
${{ if eq(config.python_version, '3.12') }}:
103+
build_intermediates_artifact_name: linux_gpu_wheel_x86_64
104+
105+
# Merge per-version Linux wheel artifacts into a single combined artifact for downstream consumers
106+
- stage: Linux_py_GPU_Wheels_Merge_Artifacts
107+
dependsOn:
108+
- ${{ each config in parameters.LinuxPythonConfigurations }}:
109+
- Linux_py_GPU_Wheels_x86_64_${{ replace(config.python_version, '.', '_') }}
110+
jobs:
111+
- job: Linux_py_GPU_Wheels_Merge_Artifacts
112+
workspace:
113+
clean: all
114+
pool:
115+
name: 'onnxruntime-Ubuntu2404-AMD-CPU'
116+
os: linux
117+
templateContext:
118+
outputs:
119+
- output: pipelineArtifact
120+
targetPath: $(Build.ArtifactStagingDirectory)/onnxruntime_gpu
121+
artifactName: onnxruntime_gpu
122+
steps:
123+
- checkout: self
124+
clean: true
125+
submodules: none
126+
127+
- ${{ each config in parameters.LinuxPythonConfigurations }}:
128+
- task: DownloadPipelineArtifact@2
129+
displayName: 'Download wheel - Python ${{ config.python_version }}'
130+
inputs:
131+
artifact: onnxruntime_gpu_${{ replace(config.python_version, '.', '_') }}
132+
targetPath: $(Build.ArtifactStagingDirectory)/onnxruntime_gpu

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

Lines changed: 47 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -28,11 +28,35 @@ parameters:
2828
- 12.8
2929
- 13.0
3030

31+
- name: python_version
32+
type: string
33+
displayName: 'Python version label (e.g. 3.12, 3.13t)'
34+
35+
- name: docker_python_exe_path
36+
type: string
37+
displayName: 'Full path to the Python executable inside the Docker image'
38+
39+
- name: stage_name
40+
type: string
41+
displayName: 'Unique stage name for this build'
42+
43+
- name: wheel_artifact_name
44+
type: string
45+
displayName: 'Name for the published wheel artifact'
46+
47+
- name: build_intermediates_artifact_name
48+
type: string
49+
default: ''
50+
displayName: |
51+
Name for the published build intermediates artifact.
52+
If the name is empty, no build intermediates artifact will be published.
53+
The build intermediates can be used for testing.
54+
3155
stages:
32-
- stage: Linux_py_GPU_Wheels_${{ parameters.arch }}
56+
- stage: ${{ parameters.stage_name }}
3357
dependsOn: []
3458
jobs:
35-
- job: Linux_py_GPU_Wheels_${{ parameters.arch }}
59+
- job: ${{ parameters.stage_name }}
3660
timeoutInMinutes: 360
3761
workspace:
3862
clean: all
@@ -43,15 +67,16 @@ stages:
4367
outputs:
4468
- output: pipelineArtifact
4569
targetPath: $(Build.ArtifactStagingDirectory)/dist
46-
artifactName: onnxruntime_gpu
47-
- output: pipelineArtifact
48-
targetPath: $(Build.ArtifactStagingDirectory)/${{ parameters.cmake_build_type }}
49-
artifactName: linux_gpu_wheel_${{ parameters.arch }}
70+
artifactName: ${{ parameters.wheel_artifact_name }}
71+
- ${{ if ne(parameters.build_intermediates_artifact_name, '') }}:
72+
- output: pipelineArtifact
73+
targetPath: $(Build.ArtifactStagingDirectory)/${{ parameters.cmake_build_type }}
74+
artifactName: ${{ parameters.build_intermediates_artifact_name }}
5075
variables:
5176
- name: extra_build_args
5277
${{ if ne(parameters.extra_build_arg, '') }}:
5378
value: '-x "${{ parameters.extra_build_arg }}"'
54-
${{ if eq(parameters.extra_build_arg, '') }}:
79+
${{ else }}:
5580
value: ''
5681
- template: ../templates/common-variables.yml
5782
- name: trt_version
@@ -75,24 +100,28 @@ stages:
75100

76101

77102
- task: Bash@3
78-
displayName: 'Build Python Wheel'
103+
displayName: 'Build Python Wheel - ${{ parameters.python_version }}'
79104
inputs:
80105
targetType: filePath
81106
filePath: tools/ci_build/github/linux/run_python_dockerbuild.sh
82-
arguments: -i onnxruntimecuda${{ replace(parameters.cuda_version, '.', '') }}xtrt86build${{ parameters.arch }} -d "GPU" -c ${{ parameters.cmake_build_type }} $(extra_build_args)
107+
arguments: >-
108+
-i onnxruntimecuda${{ replace(parameters.cuda_version, '.', '') }}xtrt86build${{ parameters.arch }}
109+
-d "GPU"
110+
-c ${{ parameters.cmake_build_type }}
111+
-p "${{ parameters.docker_python_exe_path }}"
112+
$(extra_build_args)
83113
env:
84114
CUDA_VERSION: ${{ parameters.cuda_version }}
85115

86116
- script: |
87117
set -e -x
88-
mv $(Build.BinariesDirectory)/${{ parameters.cmake_build_type }} ./${{ parameters.cmake_build_type }}
89118
mv $(Build.BinariesDirectory)/dist ./dist
90-
pushd dist
91-
find . -name \*.whl -exec unzip -qq -o {} \;
92-
rm -r onnxruntime
93-
popd
94-
pushd ${{ parameters.cmake_build_type }}
95-
find . -name \*.whl -exec unzip -qq -o {} \;
96-
popd
97119
workingDirectory: '$(Build.ArtifactStagingDirectory)'
98-
displayName: 'Move files'
120+
displayName: 'Move wheel files'
121+
122+
- ${{ if ne(parameters.build_intermediates_artifact_name, '') }}:
123+
- script: |
124+
set -e -x
125+
mv $(Build.BinariesDirectory)/${{ parameters.cmake_build_type }} ./${{ parameters.cmake_build_type }}
126+
workingDirectory: '$(Build.ArtifactStagingDirectory)'
127+
displayName: 'Move build intermediates'

0 commit comments

Comments
 (0)