Skip to content

Commit 09689ee

Browse files
lanluo-nvidialanluo-nvidia
committed
test
1 parent a0b1989 commit 09689ee

File tree

3 files changed

+383
-0
lines changed

3 files changed

+383
-0
lines changed

‎.github/scripts/generate-tensorrt-test-matrix.py‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
#!/usr/bin/env python3
2+
3+
import argparse
4+
import copy
5+
import json
6+
import sys
7+
8+
CUDA_VERSIONS_DICT = {
9+
"nightly": ["cu121"],
10+
"test": ["cu121", "cu124"],
11+
"release": ["cu121", "cu124"],
12+
}
13+
14+
PYTHON_VERSIONS_DICT = {
15+
"nightly": ["3.9"],
16+
"test": ["3.9", "3.12"],
17+
"release": ["3.9", "3.12"],
18+
}
19+
20+
TENSORRT_VERSIONS_DICT = {
21+
"windows": {
22+
"10.4": {
23+
"urls": "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.4.0/zip/TensorRT-10.4.0.26.Windows.win10.cuda-12.6.zip",
24+
"strip_prefix": "TensorRT-10.4.0.26",
25+
"sha256": "3a7de83778b9e9f812fd8901e07e0d7d6fc54ce633fcff2e340f994df2c6356c",
26+
},
27+
"10.5": {
28+
"urls": "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.5.0/zip/TensorRT-10.5.0.18.Windows.win10.cuda-12.6.zip",
29+
"strip_prefix": "TensorRT-10.5.0.18",
30+
"sha256": "e6436f4164db4e44d727354dccf7d93755efb70d6fbfd6fa95bdfeb2e7331b24",
31+
},
32+
},
33+
"linux": {
34+
"10.4": {
35+
"urls": "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.4.0/tars/TensorRT-10.4.0.26.Linux.x86_64-gnu.cuda-12.6.tar.gz",
36+
"strip_prefix": "TensorRT-10.4.0.26",
37+
"sha256": "cb0273ecb3ba4db8993a408eedd354712301a6c7f20704c52cdf9f78aa97bbdb",
38+
},
39+
"10.5": {
40+
"urls": "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.5.0/tars/TensorRT-10.5.0.18.Linux.x86_64-gnu.cuda-12.6.tar.gz",
41+
"strip_prefix": "TensorRT-10.5.0.18",
42+
"sha256": "f404d379d639552a3e026cd5267213bd6df18a4eb899d6e47815bbdb34854958",
43+
},
44+
},
45+
}
46+
47+
48+
def main(args: list[str]) -> None:
49+
parser = argparse.ArgumentParser()
50+
parser.add_argument(
51+
"--matrix",
52+
help="matrix",
53+
type=str,
54+
default="",
55+
)
56+
57+
options = parser.parse_args(args)
58+
if options.matrix == "":
59+
raise Exception(f"--matrix is empty, please provide the matrix json str")
60+
61+
matrix_dict = json.loads(options.matrix)
62+
includes = matrix_dict["include"]
63+
assert len(includes) > 0
64+
channel = includes[0].channel
65+
if "windows" in includes[0].validation_runner:
66+
arch = "windows"
67+
elif "linux" in includes[0].validation_runner:
68+
arch = "linux"
69+
else:
70+
raise Exception(
71+
f"{includes[0].validation_runner} is not the supported arch, currently only support windows and linux"
72+
)
73+
74+
cuda_versions = CUDA_VERSIONS_DICT[channel]
75+
python_versions = PYTHON_VERSIONS_DICT[channel]
76+
tensorrt_versions = TENSORRT_VERSIONS_DICT[arch]
77+
78+
filtered_includes = []
79+
for item in includes:
80+
if (
81+
item["desired_cuda"] in cuda_versions
82+
and item["python_version"] in python_versions
83+
):
84+
for _, tensorrt_json in tensorrt_versions.items():
85+
new_item = copy.deepcopy(item)
86+
new_item["tensorrt"] = tensorrt_json
87+
filtered_includes.append(new_item)
88+
filtered_matrix_dict = {}
89+
filtered_matrix_dict["include"] = filtered_includes
90+
print(json.dumps(filtered_matrix_dict))
91+
92+
93+
if __name__ == "__main__":
94+
main(sys.argv[1:])

‎.github/workflows/build-tensorrt-linux.yml‎

Lines changed: 220 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,220 @@
1+
name: Build Torch-TensorRT wheel on Linux
2+
3+
on:
4+
workflow_call:
5+
inputs:
6+
repository:
7+
description: 'Repository to checkout, defaults to ""'
8+
default: ""
9+
type: string
10+
ref:
11+
description: 'Reference to checkout, defaults to "nightly"'
12+
default: "nightly"
13+
type: string
14+
test-infra-repository:
15+
description: "Test infra repository to use"
16+
default: "pytorch/test-infra"
17+
type: string
18+
test-infra-ref:
19+
description: "Test infra reference to use"
20+
default: ""
21+
type: string
22+
build-matrix:
23+
description: "Build matrix to utilize"
24+
default: ""
25+
type: string
26+
pre-script:
27+
description: "Pre script to run prior to build"
28+
default: ""
29+
type: string
30+
post-script:
31+
description: "Post script to run prior to build"
32+
default: ""
33+
type: string
34+
smoke-test-script:
35+
description: "Script for Smoke Test for a specific domain"
36+
default: ""
37+
type: string
38+
env-var-script:
39+
description: "Script that sets Domain-Specific Environment Variables"
40+
default: ""
41+
type: string
42+
package-name:
43+
description: "Name of the actual python package that is imported"
44+
default: ""
45+
type: string
46+
trigger-event:
47+
description: "Trigger Event in caller that determines whether or not to upload"
48+
default: ""
49+
type: string
50+
cache-path:
51+
description: "The path(s) on the runner to cache or restore. The path is relative to repository."
52+
default: ""
53+
type: string
54+
cache-key:
55+
description: "The key created when saving a cache and the key used to search for a cache."
56+
default: ""
57+
type: string
58+
architecture:
59+
description: Architecture to build for x86_64 for default Linux, or aarch64 for Linux aarch64 builds
60+
required: false
61+
type: string
62+
default: x86_64
63+
submodules:
64+
description: Works as stated in actions/checkout, but the default value is recursive
65+
required: false
66+
type: string
67+
default: recursive
68+
setup-miniconda:
69+
description: Set to true if setup-miniconda is needed
70+
required: false
71+
type: boolean
72+
default: true
73+
74+
permissions:
75+
id-token: write
76+
contents: read
77+
78+
jobs:
79+
build:
80+
strategy:
81+
fail-fast: false
82+
matrix: ${{ fromJSON(inputs.build-matrix) }}
83+
env:
84+
PYTHON_VERSION: ${{ matrix.python_version }}
85+
PACKAGE_TYPE: wheel
86+
REPOSITORY: ${{ inputs.repository }}
87+
REF: ${{ inputs.ref }}
88+
CU_VERSION: ${{ matrix.desired_cuda }}
89+
UPLOAD_TO_BASE_BUCKET: ${{ matrix.upload_to_base_bucket }}
90+
ARCH: ${{ inputs.architecture }}
91+
TENSORRT_STRIP_PREFIX: ${{ matrix.tensorrt.strip_prefix }}
92+
TENSORRT_URLS: ${{ matrix.tensorrt.urls }}
93+
TENSORRT_SHA256: ${{ matrix.tensorrt.sha256 }}
94+
name: release_${{ matrix.build_name }}
95+
runs-on: ${{ matrix.validation_runner }}
96+
container:
97+
image: ${{ matrix.container_image }}
98+
options: ${{ matrix.gpu_arch_type == 'cuda' && '--gpus all' || ' ' }}
99+
# If a build is taking longer than 120 minutes on these runners we need
100+
# to have a conversation
101+
timeout-minutes: 120
102+
103+
steps:
104+
- name: Clean workspace
105+
shell: bash -l {0}
106+
run: |
107+
set -x
108+
echo "::group::Cleanup debug output"
109+
rm -rf "${GITHUB_WORKSPACE}"
110+
mkdir -p "${GITHUB_WORKSPACE}"
111+
if [[ "${{ inputs.architecture }}" = "aarch64" ]]; then
112+
rm -rf "${RUNNER_TEMP}/*"
113+
fi
114+
echo "::endgroup::"
115+
- uses: actions/checkout@v3
116+
with:
117+
# Support the use case where we need to checkout someone's fork
118+
repository: ${{ inputs.test-infra-repository }}
119+
ref: ${{ inputs.test-infra-ref }}
120+
path: test-infra
121+
- uses: actions/checkout@v3
122+
if: ${{ env.ARCH == 'aarch64' }}
123+
with:
124+
# Support the use case where we need to checkout someone's fork
125+
repository: "pytorch/builder"
126+
ref: "main"
127+
path: builder
128+
- name: Set linux aarch64 CI
129+
if: ${{ inputs.architecture == 'aarch64' }}
130+
shell: bash -l {0}
131+
env:
132+
DESIRED_PYTHON: ${{ matrix.python_version }}
133+
run: |
134+
set +e
135+
# TODO: This is temporary aarch64 setup script, this should be integrated into aarch64 docker.
136+
${GITHUB_WORKSPACE}/builder/aarch64_linux/aarch64_ci_setup.sh
137+
echo "/opt/conda/bin" >> $GITHUB_PATH
138+
set -e
139+
- uses: ./test-infra/.github/actions/set-channel
140+
- name: Set PYTORCH_VERSION
141+
if: ${{ env.CHANNEL == 'test' }}
142+
run: |
143+
# When building RC, set the version to be the current candidate version,
144+
# otherwise, leave it alone so nightly will pick up the latest
145+
echo "PYTORCH_VERSION=${{ matrix.stable_version }}" >> "${GITHUB_ENV}"
146+
- uses: ./test-infra/.github/actions/setup-binary-builds
147+
env:
148+
PLATFORM: ${{ inputs.architecture == 'aarch64' && 'linux-aarch64' || ''}}
149+
with:
150+
repository: ${{ inputs.repository }}
151+
ref: ${{ inputs.ref }}
152+
submodules: ${{ inputs.submodules }}
153+
setup-miniconda: ${{ inputs.setup-miniconda }}
154+
python-version: ${{ env.PYTHON_VERSION }}
155+
cuda-version: ${{ env.CU_VERSION }}
156+
arch: ${{ env.ARCH }}
157+
- name: Combine Env Var and Build Env Files
158+
if: ${{ inputs.env-var-script != '' }}
159+
working-directory: ${{ inputs.repository }}
160+
shell: bash -l {0}
161+
run: |
162+
cat "${{ inputs.env-var-script }}" >> "${BUILD_ENV_FILE}"
163+
- name: Install torch dependency
164+
shell: bash -l {0}
165+
run: |
166+
set -x
167+
# shellcheck disable=SC1090
168+
source "${BUILD_ENV_FILE}"
169+
# shellcheck disable=SC2086
170+
${CONDA_RUN} ${PIP_INSTALL_TORCH}
171+
- name: Run Pre-Script with Caching
172+
if: ${{ inputs.pre-script != '' }}
173+
uses: ./test-infra/.github/actions/run-script-with-cache
174+
with:
175+
cache-path: ${{ inputs.cache-path }}
176+
cache-key: ${{ inputs.cache-key }}
177+
repository: ${{ inputs.repository }}
178+
script: ${{ inputs.pre-script }}
179+
- name: Build clean
180+
working-directory: ${{ inputs.repository }}
181+
shell: bash -l {0}
182+
run: |
183+
set -x
184+
source "${BUILD_ENV_FILE}"
185+
${CONDA_RUN} python setup.py clean
186+
- name: Build the wheel (bdist_wheel)
187+
working-directory: ${{ inputs.repository }}
188+
shell: bash -l {0}
189+
run: |
190+
set -x
191+
source "${BUILD_ENV_FILE}"
192+
${CONDA_RUN} python setup.py bdist_wheel
193+
194+
- name: Run Post-Script
195+
if: ${{ inputs.post-script != '' }}
196+
uses: ./test-infra/.github/actions/run-script-with-cache
197+
with:
198+
repository: ${{ inputs.repository }}
199+
script: ${{ inputs.post-script }}
200+
- name: Smoke Test
201+
shell: bash -l {0}
202+
env:
203+
PACKAGE_NAME: ${{ inputs.package-name }}
204+
SMOKE_TEST_SCRIPT: ${{ inputs.smoke-test-script }}
205+
run: |
206+
set -x
207+
source "${BUILD_ENV_FILE}"
208+
# TODO: add smoke test for the auditwheel tarball built
209+
210+
# NB: Only upload to GitHub after passing smoke tests
211+
- name: Upload wheel to GitHub
212+
continue-on-error: true
213+
uses: actions/upload-artifact@v3
214+
with:
215+
name: ${{ env.ARTIFACT_NAME }}
216+
path: ${{ inputs.repository }}/dist
217+
218+
concurrency:
219+
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ inputs.repository }}-${{ github.event_name == 'workflow_dispatch' }}-${{ inputs.job-name }}-${{ inputs.cxx11-tarball-release }}
220+
cancel-in-progress: true

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

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
name: Build and test Torch-TensorRT on Linux
2+
3+
on:
4+
pull_request:
5+
workflow_dispatch:
6+
7+
permissions:
8+
id-token: write
9+
contents: read
10+
packages: write
11+
12+
jobs:
13+
generate-matrix:
14+
uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
15+
with:
16+
package-type: wheel
17+
os: linux
18+
test-infra-repository: pytorch/test-infra
19+
test-infra-ref: main
20+
with-rocm: false
21+
with-cpu: false
22+
23+
generate-tensorrt-matrix:
24+
needs: [generate-matrix]
25+
outputs:
26+
matrix: ${{ steps.generate.outputs.matrix }}
27+
runs-on: ubuntu-latest
28+
steps:
29+
- uses: actions/setup-python@v4
30+
with:
31+
python-version: '3.10'
32+
- uses: actions/checkout@v3
33+
with:
34+
repository: pytorch/tensorrt
35+
- name: Generate tensorrt matrix
36+
id: generate
37+
run: |
38+
set -eou pipefail
39+
MATRIX_BLOB=${{ toJSON(needs.generate-matrix.outputs.matrix) }}
40+
MATRIX_BLOB="$(python3 .github/scripts/generate-tensorrt-test-matrix.py --matrix "${MATRIX_BLOB}")"
41+
echo "${MATRIX_BLOB}"
42+
echo "matrix=${MATRIX_BLOB}" >> "${GITHUB_OUTPUT}"
43+
44+
build:
45+
needs: [generate-tensorrt-matrix]
46+
name: Build torch-tensorrt whl package
47+
strategy:
48+
fail-fast: false
49+
matrix:
50+
include:
51+
- repository: pytorch/tensorrt
52+
pre-script: packaging/pre_build_script.sh
53+
env-var-script: packaging/env_vars.txt
54+
post-script: packaging/post_build_script.sh
55+
smoke-test-script: packaging/smoke_test_script.sh
56+
package-name: torch_tensorrt
57+
uses: ./.github/workflows/build-torch-tensorrt-linux.yml
58+
with:
59+
repository: ${{ matrix.repository }}
60+
ref: ""
61+
test-infra-repository: pytorch/test-infra
62+
test-infra-ref: main
63+
build-matrix: ${{ needs.generate-tensorrt-matrix.outputs.matrix }}
64+
pre-script: ${{ matrix.pre-script }}
65+
env-var-script: ${{ matrix.env-var-script }}
66+
post-script: ${{ matrix.post-script }}
67+
package-name: ${{ matrix.package-name }}
68+
smoke-test-script: ${{ matrix.smoke-test-script }}
69+
trigger-event: ${{ github.event_name }}

0 commit comments

Comments
 (0)