Skip to content

Commit 5b3fb01

Browse files
lanluo-nvidialanluo-nvidia
authored
[CI] Add manylinux auditwheel repair (#778)
1 parent f06fba6 commit 5b3fb01

File tree

16 files changed

+791
-33
lines changed

16 files changed

+791
-33
lines changed

‎.github/workflows/mlir-tensorrt-build-test.yml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ on:
44
workflow_call:
55
inputs:
66
channel:
7-
description: 'Channel, valid values are "nightly", "test", or "release"'
7+
description: 'Channel, valid values are "nightly", "test", "release" or "pypi-release"'
88
default: "test"
99
type: string
1010
build-matrix:
@@ -26,7 +26,7 @@ jobs:
2626
env:
2727
# eg. TENSORRT_VERSION: 10.12 or 10.13
2828
MLIR_TRT_DOWNLOAD_TENSORRT_VERSION: ${{ matrix.trt }}
29-
# eg. CHANNEL: nightly, test or release
29+
# eg. CHANNEL: nightly, test, release or pypi-release
3030
CHANNEL: ${{ inputs.channel }}
3131
ARCH: ${{ matrix.arch }}
3232
CMAKE_PRESET: ${{ matrix.cmake_preset }}

‎.github/workflows/mlir-tensorrt-pypi.yml‎

Lines changed: 252 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,252 @@
1+
name: MLIR-TensorRT PyPI Release CI
2+
3+
on:
4+
workflow_dispatch:
5+
inputs:
6+
confirm_publish_to_testpypi:
7+
description: 'I confirm that I have verified version to be published to test.pypi.org'
8+
required: false
9+
type: boolean
10+
default: false
11+
confirm_publish_to_pypi:
12+
description: 'I confirm that I have verified version to be published to pypi.org'
13+
required: false
14+
type: boolean
15+
default: false
16+
17+
defaults:
18+
run:
19+
shell: bash
20+
21+
jobs:
22+
generate-matrix:
23+
name: Generate Build Matrix (pypi-release)
24+
runs-on: ubuntu-latest
25+
outputs:
26+
matrix: ${{ steps.generate-matrix.outputs.matrix }}
27+
steps:
28+
- uses: actions/checkout@v6
29+
with:
30+
fetch-depth: 5
31+
- name: Generate Build Matrix
32+
id: generate-matrix
33+
run: |
34+
set -euo pipefail
35+
set -x
36+
MATRIX_BLOB="$(python3 ./.github/workflows/mlir-tensorrt/generate-matrix.py --channel pypi-release)"
37+
echo "${MATRIX_BLOB}"
38+
echo "matrix=${MATRIX_BLOB}" >> "${GITHUB_OUTPUT}"
39+
40+
pypi-release-wheels-build:
41+
name: ${{ matrix.arch }} - Build PyPI Release Wheels
42+
needs:
43+
- generate-matrix
44+
permissions:
45+
id-token: write
46+
packages: write
47+
contents: read
48+
strategy:
49+
matrix: ${{ fromJSON(needs.generate-matrix.outputs.matrix) }}
50+
runs-on: ${{ matrix.github_runner }}
51+
env:
52+
MLIR_TRT_DOWNLOAD_TENSORRT_VERSION: ${{ matrix.trt }}
53+
CMAKE_PRESET: ${{ matrix.cmake_preset }}
54+
CCACHE_RESTORE_KEY: mlir-tensorrt-ccache-v1-${{ matrix.arch }}-${{ matrix.cmake_preset }}
55+
CPM_RESTORE_KEY: mlir-tensorrt-cpm-v1
56+
timeout-minutes: 120
57+
container:
58+
# pypi audit wheel repair requires rockylinux8
59+
image: ${{ matrix.docker_image }}
60+
options: >-
61+
--gpus all
62+
--shm-size=1g
63+
steps:
64+
- name: Checkout TensorRT-Incubator
65+
uses: actions/checkout@v6
66+
with:
67+
fetch-depth: 5
68+
69+
- name: Create Cache Folders
70+
run: |
71+
set -euo pipefail
72+
set -x
73+
export CPM_SOURCE_CACHE=${GITHUB_WORKSPACE}/mlir-tensorrt/.cache.cpm
74+
export CCACHE_DIR=${GITHUB_WORKSPACE}/mlir-tensorrt/ccache
75+
76+
echo "CPM_SOURCE_CACHE=$CPM_SOURCE_CACHE" >> "$GITHUB_ENV"
77+
echo "CCACHE_DIR=$CCACHE_DIR" >> "$GITHUB_ENV"
78+
79+
mkdir -p ${CCACHE_DIR}
80+
mkdir -p ${CPM_SOURCE_CACHE}
81+
82+
- name: Compute CCache Key
83+
id: ccache-key
84+
run: |
85+
hash=$( (find mlir-tensorrt/compiler \
86+
mlir-tensorrt/common \
87+
mlir-tensorrt/kernel \
88+
mlir-tensorrt/tensorrt \
89+
mlir-tensorrt/integrations \
90+
mlir-tensorrt/executor \
91+
-type f \( -name '*.cpp' -o -name '*.h' \) \
92+
-exec sha256sum {} \; ; \
93+
sha256sum mlir-tensorrt/DependencyProvider.cmake \
94+
mlir-tensorrt/CMakeLists.txt) \
95+
| sort | sha256sum | cut -d' ' -f1)
96+
echo "key=${{ env.CCACHE_RESTORE_KEY }}-${hash}" >> $GITHUB_OUTPUT
97+
98+
- name: Compute CPM Key
99+
id: cpm-key
100+
run: |
101+
hash=$(sha256sum mlir-tensorrt/DependencyProvider.cmake | cut -d' ' -f1)
102+
echo "key=${{ env.CPM_RESTORE_KEY }}-${hash}" >> $GITHUB_OUTPUT
103+
104+
- name: Restore CCache
105+
id: restore-ccache
106+
uses: actions/cache/restore@v4
107+
with:
108+
key: ${{ steps.ccache-key.outputs.key }}
109+
restore-keys: |
110+
${{ env.CCACHE_RESTORE_KEY }}
111+
path: |
112+
${{ env.CCACHE_DIR }}
113+
114+
- name: Restore CPM cache
115+
id: restore-cpm
116+
uses: actions/cache/restore@v4
117+
with:
118+
key: ${{ steps.cpm-key.outputs.key }}
119+
enableCrossOsArchive: true
120+
restore-keys: |
121+
${{ env.CPM_RESTORE_KEY }}
122+
path: |
123+
mlir-tensorrt/.cache.cpm/*
124+
!mlir-tensorrt/.cache.cpm/tensorrt
125+
!mlir-tensorrt/.cache.cpm/tensorrt/**
126+
127+
- name: Build Wheels With CUDA:${{ matrix.cuda }} + TensorRT:${{ matrix.trt }}
128+
env:
129+
MLIR_TRT_DOWNLOAD_TENSORRT_VERSION: ${{ matrix.trt }}
130+
ARCH: ${{ matrix.arch }}
131+
CMAKE_PRESET: distribution-wheels
132+
run: |
133+
set -euo pipefail
134+
set -x
135+
cd mlir-tensorrt
136+
# Build only pjrt wheels for PyPI upload, with auditwheel repair
137+
# if you want to build for post version, set MLIR_TRT_PYPI_POST_VERSION=1
138+
# MLIR_TRT_PYPI=1 MLIR_TRT_PYPI_POST_VERSION=1 PACKAGES="pjrt" VERBOSE=1 ./build_tools/scripts/cicd-build-wheels.sh
139+
MLIR_TRT_PYPI=1 PACKAGES="pjrt" VERBOSE=1 ./build_tools/scripts/cicd-build-wheels.sh
140+
141+
- name: Upload Wheels
142+
uses: actions/upload-artifact@v4
143+
with:
144+
name: release-wheels-${{ matrix.arch }}-cu${{ matrix.cuda }}-trt${{ matrix.trt }}
145+
path: mlir-tensorrt/dist
146+
if-no-files-found: error
147+
148+
pypi-release-wheels-smoke-test:
149+
name: ${{ matrix.arch }} - Smoke Test PyPI Release Wheels
150+
needs:
151+
- generate-matrix
152+
- pypi-release-wheels-build
153+
permissions:
154+
id-token: write
155+
packages: write
156+
contents: read
157+
strategy:
158+
matrix: ${{ fromJSON(needs.generate-matrix.outputs.matrix) }}
159+
runs-on: ${{ matrix.github_runner }}
160+
timeout-minutes: 120
161+
container:
162+
image: ghcr.io/nvidia/tensorrt-incubator/mlir-tensorrt:cuda13.0-ubuntu24.04-0.1
163+
options: >-
164+
--gpus all
165+
--shm-size=1g
166+
steps:
167+
- name: Checkout TensorRT-Incubator
168+
uses: actions/checkout@v6
169+
with:
170+
fetch-depth: 5
171+
- name: Download built wheels
172+
uses: actions/download-artifact@v4
173+
with:
174+
pattern: release-wheels-*
175+
merge-multiple: true
176+
path: mlir-tensorrt/dist
177+
178+
- name: Smoke Test Wheels With CUDA:${{ matrix.cuda }} + TensorRT:${{ matrix.trt }}
179+
run: |
180+
set -euo pipefail
181+
set -x
182+
cd mlir-tensorrt
183+
uv sync --python 3.12 --extra cu13 --extra flashinfer-cu13
184+
source .venv/bin/activate
185+
ls -lart dist/
186+
suffix="manylinux_2_28_${{ matrix.arch }}.whl"
187+
whl_file=$(find dist/ -name "mlir_tensorrt_jax-*-cp312-cp312-*${suffix}")
188+
echo "whl_file: ${whl_file}"
189+
uv pip install ${whl_file}
190+
191+
export JAX_PLATFORMS="mlir_tensorrt,cpu"
192+
python -m pytest integrations/PJRT/test/JAX/ -v
193+
194+
test-pypi-release-wheels-publish:
195+
name: Publish to TestPyPI
196+
if: ${{ inputs.confirm_publish_to_testpypi }}
197+
needs:
198+
- pypi-release-wheels-build
199+
- pypi-release-wheels-smoke-test
200+
runs-on: ubuntu-latest
201+
environment:
202+
name: testpypi
203+
url: https://test.pypi.org/project/mlir-tensorrt-jax/
204+
permissions:
205+
id-token: write
206+
packages: write
207+
contents: read
208+
steps:
209+
- name: Download built wheels
210+
uses: actions/download-artifact@v4
211+
with:
212+
pattern: release-wheels-*
213+
merge-multiple: true
214+
path: dist
215+
216+
- name: Publish to TestPyPI
217+
uses: pypa/gh-action-pypi-publish@release/v1
218+
with:
219+
skip-existing: true
220+
verbose: true
221+
repository-url: https://test.pypi.org/legacy/
222+
223+
pypi-release-wheels-publish:
224+
name: Publish to PyPI
225+
if: ${{ inputs.confirm_publish_to_pypi }}
226+
needs:
227+
- pypi-release-wheels-build
228+
- pypi-release-wheels-smoke-test
229+
runs-on: ubuntu-latest
230+
permissions:
231+
id-token: write
232+
packages: write
233+
contents: read
234+
steps:
235+
- name: Download built wheels
236+
uses: actions/download-artifact@v4
237+
with:
238+
pattern: release-wheels-*
239+
merge-multiple: true
240+
path: dist
241+
242+
- name: Publish to PyPI
243+
uses: pypa/gh-action-pypi-publish@release/v1
244+
with:
245+
verbose: true
246+
skip-existing: true
247+
user: __token__
248+
password: ${{ secrets.PYPI_API_TOKEN }}
249+
250+
concurrency:
251+
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-mlir-tensorrt-pypi
252+
cancel-in-progress: true

‎.github/workflows/mlir-tensorrt-release.yml‎

Lines changed: 3 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -23,8 +23,8 @@ jobs:
2323
# eg. 10.12 or 10.13
2424
MLIR_TRT_DOWNLOAD_TENSORRT_VERSION: ${{ matrix.trt }}
2525
ARCH: ${{ matrix.arch }}
26-
CMAKE_PRESET: distribution-wheels
27-
CCACHE_RESTORE_KEY: mlir-tensorrt-ccache-v1-${{ matrix.arch }}-distribution-wheels
26+
CMAKE_PRESET: ${{ matrix.cmake_preset }}
27+
CCACHE_RESTORE_KEY: mlir-tensorrt-ccache-v1-${{ matrix.arch }}-${{ matrix.cmake_preset }}
2828
CPM_RESTORE_KEY: mlir-tensorrt-cpm-v1
2929
runs-on: ${{ matrix.github_runner }}
3030
timeout-minutes: 120
@@ -144,17 +144,6 @@ jobs:
144144
cd mlir-tensorrt
145145
export JAX_PLATFORMS="mlir_tensorrt,cpu"
146146
147-
# Compute TensorRT library path from CPM cache by locating libnvinfer.so
148-
SEARCH_PATH=${CPM_SOURCE_CACHE}/tensorrt
149-
echo "SEARCH_PATH: ${SEARCH_PATH}"
150-
TensorRT_LIB_PATH="$(find ${SEARCH_PATH} -type f -name 'libnvinfer.so.*' 2>/dev/null | head -n1 | xargs -r dirname)"
151-
if [[ -z "${TensorRT_LIB_PATH}" ]]; then
152-
echo "Error: Could not locate libnvinfer.so under ${CPM_TRT_DIR}" >&2
153-
exit 1
154-
fi
155-
echo "TensorRT_LIB_PATH: ${TensorRT_LIB_PATH}"
156-
export LD_LIBRARY_PATH="${TensorRT_LIB_PATH}:$LD_LIBRARY_PATH"
157-
158147
# Choose uv extra based on CUDA_VERSION major (12.9 -> cu12, 13.0 -> cu13).
159148
matrix_cuda="${{ matrix.cuda }}"
160149
CUDA_MAJOR="${matrix_cuda%%.*}"
@@ -165,7 +154,7 @@ jobs:
165154
fi
166155
echo "UV_EXTRA: ${UV_EXTRA}"
167156
168-
uv sync --python 3.10 --extra ${UV_EXTRA}
157+
uv sync --python 3.10 --extra ${UV_EXTRA} --extra flashinfer-${UV_EXTRA}
169158
source .venv/bin/activate
170159
171160
uv pip install dist/mlir_tensorrt_jax-*-cp310-cp310*.whl

‎.github/workflows/mlir-tensorrt/generate-matrix.py‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -31,12 +31,19 @@
3131
"trt": "10.13",
3232
},
3333
],
34+
"pypi-release": [
35+
{
36+
"cuda": "13.0",
37+
"trt": "10.13",
38+
},
39+
],
3440
}
3541

3642
ARCH_LIST_DICT = {
3743
"test": ["x86_64"],
3844
"release": ["x86_64", "aarch64"],
3945
"nightly": ["x86_64", "aarch64"],
46+
"pypi-release": ["x86_64", "aarch64"],
4047
}
4148

4249
GH_RUNNER_DICT = {
@@ -47,9 +54,8 @@
4754
CMAKE_PRESET_DICT = {
4855
"nightly": "github-cicd",
4956
"test": "github-cicd",
50-
# release should use the release wheel build preset
51-
# TODO: add the release wheel build preset
52-
"release": "github-cicd",
57+
"release": "distribution-wheels",
58+
"pypi-release": "distribution-wheels",
5359
}
5460

5561
DOCKER_IMAGE_DICT = {
@@ -71,6 +77,9 @@
7177
"13.0": "ghcr.io/nvidia/tensorrt-incubator/mlir-tensorrt:cuda13.0-rockylinux9-0.1",
7278
},
7379
},
80+
"pypi-release": {
81+
"13.0": "ghcr.io/nvidia/tensorrt-incubator/mlir-tensorrt:cuda13.0-rockylinux8-0.1",
82+
},
7483
}
7584

7685

@@ -84,9 +93,9 @@ def main(args: list[str]) -> None:
8493
)
8594

8695
options = parser.parse_args(args)
87-
if options.channel not in ("nightly", "test", "release"):
96+
if options.channel not in ("nightly", "test", "release", "pypi-release"):
8897
raise Exception(
89-
"--channel is invalid, please choose from nightly, test or release"
98+
"--channel is invalid, please choose from nightly, test, release or pypi-release"
9099
)
91100

92101
channel = options.channel

0 commit comments

Comments
 (0)