Simplify Release workflow and Add windows zip in the release artifacts

.github/scripts/generate-release-matrix.py

@@ -5,16 +5,17 @@
 import sys
 
 RELEASE_CUDA_VERSION = {
-    "wheel": ["cu128"],
-    "tarball": ["cu128"],
+    "wheel": ["cu129"],
+    "tarball": ["cu129"],
 }
 RELEASE_PYTHON_VERSION = {
-    "wheel": ["3.9", "3.10", "3.11", "3.12", "3.13"],
+    "wheel": ["3.10", "3.11", "3.12", "3.13"],
     "tarball": ["3.11"],
 }
+sbsa_container_image: str = "quay.io/pypa/manylinux_2_34_aarch64"
 
 CXX11_TARBALL_CONTAINER_IMAGE = {
-    "cu128": "pytorch/libtorch-cxx11-builder:cuda12.8-main",
+    "cu129": "pytorch/libtorch-cxx11-builder:cuda12.9-main",
 }
 
 
@@ -56,10 +57,12 @@ def main(args: list[str]) -> None:
             item["desired_cuda"] in cuda_versions
             and item["python_version"] in python_versions
         ):
-            if options.tarball_matrix != "":
-                item["container_image"] = CXX11_TARBALL_CONTAINER_IMAGE[
-                    item["desired_cuda"]
-                ]
+            if item["gpu_arch_type"] == "cuda-aarch64":
+                item["container_image"] = sbsa_container_image
+            # if options.tarball_matrix != "":
+            #     item["container_image"] = CXX11_TARBALL_CONTAINER_IMAGE[
+            #         item["desired_cuda"]
+            #     ]
             filtered_includes.append(item)
     filtered_matrix_dict = {}
     filtered_matrix_dict["include"] = filtered_includes

.github/workflows/build_wheels_linux.yml

@@ -117,6 +117,16 @@ on:
         required: false
         default: false
         type: boolean
+      is-release-wheel:
+        description: Set to true if the build is for release wheel
+        required: false
+        default: false
+        type: boolean
+      is-release-tarball:
+        description: Set to true if the build is for release tarball
+        required: false
+        default: false
+        type: boolean
     secrets:
       PYPI_API_TOKEN:
         description: An optional token to upload to pypi
@@ -265,22 +275,42 @@ jobs:
           export PYTORCH_VERSION="$(${CONDA_RUN} pip show torch | grep ^Version: | sed 's/Version: *//' | sed 's/+.\+//')"
           ${CONDA_RUN} python setup.py clean
           echo "Successfully ran `python setup.py clean`"
-          if [[ "$BUILD_VERSION" != *"+"${CU_VERSION} ]]; then
-            BUILD_VERSION="${BUILD_VERSION}+${CU_VERSION}"
+          if [[ ${{ inputs.is-release-wheel }} == true || ${{ inputs.is-release-tarball }} == true ]]; then
+              # release version for upload to pypi
+              # BUILD_VERSION example: 2.4.0+cu121, we don't want the +cu121 part, so remove +cu121
+              BUILD_VERSION=${BUILD_VERSION%+*}
+              if [[ ${{ inputs.is-release-tarball }} == true ]]; then
+                  mkdir -p release/tarball
+                  TRT_VERSION=$(cat dev_dep_versions.yml | grep __tensorrt_version__ | sed 's/__tensorrt_version__: //g' | sed 's/"//g')
+                  bazel build //:libtorchtrt --compilation_mode opt --config=linux
+                  cp bazel-bin/libtorchtrt.tar.gz \
+                  release/tarball/libtorchtrt-${BUILD_VERSION}-tensorrt${TRT_VERSION}-cuda${CU_VERSION:2}-libtorch${PYTORCH_VERSION}-x86_64-linux.tar.gz
+                  return 0
+              else
+                  if [[ ${{ inputs.is-release-wheel }} == true ]]; then
+                    param="--release"
+                  fi
+              fi
+          else
+              # release version for upload to pytorch index
+              if [[ "$BUILD_VERSION" != *"+"${CU_VERSION} ]]; then
+                BUILD_VERSION="${BUILD_VERSION}+${CU_VERSION}"
+              fi
           fi
+
           echo "BUILD_VERSION=$BUILD_VERSION"
           echo "USE_TRT_RTX=$USE_TRT_RTX"
           echo "LD_LIBRARY_PATH=$LD_LIBRARY_PATH"
           if [[ ${{ inputs.use-rtx }} == true ]]; then
             echo "Building tensorrt-rtx wheel"
-            ${CONDA_RUN} python setup.py bdist_wheel --use-rtx
+            ${CONDA_RUN} python setup.py bdist_wheel ${param} --use-rtx
           else
             if [[ ${{ inputs.is-jetpack }} == true ]]; then
               echo "Building tensorrt wheel for jetpack"
-              ${CONDA_RUN} python setup.py bdist_wheel --jetpack
+              ${CONDA_RUN} python setup.py bdist_wheel ${param} --jetpack
             else
               echo "Building standard tensorrt wheel"
-              ${CONDA_RUN} python setup.py bdist_wheel
+              ${CONDA_RUN} python setup.py bdist_wheel ${param}
             fi
           fi
       - name: Repair Manylinux_2_28 Wheel
@@ -289,7 +319,7 @@ jobs:
           PACKAGE_NAME: ${{ inputs.package-name }}
           SMOKE_TEST_SCRIPT: ${{ inputs.smoke-test-script }}
         # TODO: lan to verify whether manylinux repair is needed for jetpack
-        #if: ${{ inputs.is-jetpack == true }}
+        if: ${{ inputs.is-release-tarball == false }}
         run: |
           set -euxo pipefail
           source "${BUILD_ENV_FILE}"
@@ -312,6 +342,7 @@ jobs:
         env:
           PACKAGE_NAME: ${{ inputs.package-name }}
           SMOKE_TEST_SCRIPT: ${{ inputs.smoke-test-script }}
+        if: ${{ inputs.is-release-tarball == false }}
         run: |
           set -euxo pipefail
           source "${BUILD_ENV_FILE}"
@@ -348,18 +379,27 @@ jobs:
       # NB: Only upload to GitHub after passing smoke tests
 
       - name: Upload wheel to GitHub
+        if: ${{ inputs.is-release-tarball == false }}
         continue-on-error: true
         uses: actions/upload-artifact@v4
         with:
           name: ${{ env.ARTIFACT_NAME }}
           path: ${{ inputs.repository }}/dist/
+      - name: Upload cxx11 tarball to GitHub
+        if: ${{ inputs.is-release-tarball == true }}
+        continue-on-error: true
+        uses: actions/upload-artifact@v4
+        with:
+          name: cxx11-tarball-${{ env.PYTHON_VERSION }}-${{ env.CU_VERSION }}
+          path: ${{ inputs.repository }}/release/tarball/
 
   upload:
     needs: build
     name: upload-wheel-${{ matrix.python_version }}-${{ matrix.desired_cuda }}-${{ matrix.gpu_arch_type }}-${{ inputs.is-jetpack }}
     uses: pytorch/test-infra/.github/workflows/_binary_upload.yml@main
     #  if it is not jetpack nor rtx, then upload to pytorch index
-    if: ${{ inputs.is-jetpack == false && inputs.use-rtx == false }}
+    # if it is the release wheel or tarball, then skip upload to pytorch index
+    if: ${{ inputs.is-jetpack == false && inputs.use-rtx == false && inputs.is-release-wheel == false && inputs.is-release-tarball == false }}
     with:
       repository: ${{ inputs.repository }}
       ref: ${{ inputs.ref }}
@@ -373,5 +413,5 @@ jobs:
       PYPI_API_TOKEN: ${{ secrets.PYPI_API_TOKEN }}
 
 concurrency:
-  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{inputs.use-rtx}}-${{inputs.architecture}}-${{inputs.is-jetpack}}-${{ inputs.repository }}-${{ github.event_name == 'workflow_dispatch' }}
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{inputs.is-release-wheel}}-${{inputs.is-release-tarball}}-${{inputs.use-rtx}}-${{inputs.architecture}}-${{inputs.is-jetpack}}-${{ inputs.repository }}-${{ github.event_name == 'workflow_dispatch' }}
   cancel-in-progress: true

.github/workflows/build_wheels_windows.yml

@@ -89,6 +89,16 @@ on:
         description: 'CPU architecture to build for'
         default: "x64"
         type: string
+      is-release-tarball:
+        description: Set to true if the build is for release tarball
+        required: false
+        default: false
+        type: boolean
+      is-release-wheel:
+        description: Set to true if the build is for release wheel
+        required: false
+        default: false
+        type: boolean
 
 permissions:
   id-token: write
@@ -277,6 +287,20 @@ jobs:
           BUILD_PARAMS: ${{ inputs.wheel-build-params }}
         run: |
           source "${BUILD_ENV_FILE}"
+          if [[ ${{ inputs.is-release-wheel }} == true || ${{ inputs.is-release-tarball }} == true ]]; then
+            # release version for upload to pypi
+            # BUILD_VERSION example: 2.4.0+cu121, we don't want the +cu121 part, so remove +cu121
+            BUILD_VERSION=${BUILD_VERSION%+*}
+            if [[ ${{ inputs.is-release-tarball }} == true ]]; then
+              mkdir -p release/tarball
+              TRT_VERSION=$(cat dev_dep_versions.yml | grep __tensorrt_version__ | sed 's/__tensorrt_version__: //g' | sed 's/"//g')
+              bazel build //:libtorchtrt --compilation_mode opt --config=windows
+              ls -lart bazel-bin/
+              cp bazel-bin/libtorchtrt.zip\
+              release/tarball/libtorchtrt-${BUILD_VERSION}-tensorrt${TRT_VERSION}-cuda${CU_VERSION:2}-libtorch${PYTORCH_VERSION}-x86_64-windows.zip
+              return 0
+            fi
+          fi
           if [[ -z "${ENV_SCRIPT}" ]]; then
             ${CONDA_RUN} python setup.py bdist_wheel
           else
@@ -311,7 +335,7 @@ jobs:
           source "${BUILD_ENV_FILE}"
           ${CONDA_RUN} ${ENV_SCRIPT} ${POST_SCRIPT}
       - name: Smoke Test X64
-        if: inputs.architecture == 'x64'
+        if: ${{ inputs.architecture == 'x64' && inputs.is-release-tarball == false }}
         env:
           ENV_SCRIPT: ${{ inputs.env-script }}
           PACKAGE_NAME: ${{ inputs.package-name }}
@@ -365,11 +389,19 @@ jobs:
         run: |
           echo "ARTIFACT_NAME=${REPOSITORY//\//_}_${REF//\//_}_${PYTHON_VERSION}_${CU_VERSION}_${ARCH}" >> "${GITHUB_ENV}"
       - name: Upload wheel to GitHub
+        if: ${{ inputs.is-release-tarball == false }}
         continue-on-error: true
         uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
         with:
           name: ${{ env.ARTIFACT_NAME }}
           path: ${{ inputs.repository }}/dist/
+      - name: Upload zip artifact to GitHub
+        if: ${{ inputs.is-release-tarball == true }}
+        continue-on-error: true
+        uses: actions/upload-artifact@v4
+        with:
+          name: cxx11-zip-${{ env.PYTHON_VERSION }}-${{ env.CU_VERSION }}
+          path: ${{ inputs.repository }}/release/tarball/
       - uses: ./test-infra/.github/actions/teardown-windows
         if: inputs.architecture == 'x64'
         name: Teardown Windows
@@ -378,7 +410,8 @@ jobs:
     needs: build
     uses: pytorch/test-infra/.github/workflows/_binary_upload.yml@main
     # for tensorrt-rtx build, do not upload to pytorch indexat at all
-    if: ${{ inputs.use-rtx == false }}
+    # if it is the release wheel or tarball, then skip upload to pytorch index(release wheel are uploaded to pypi)
+    if: ${{ inputs.use-rtx == false && inputs.is-release-wheel == false && inputs.is-release-tarball == false }}
     with:
       repository: ${{ inputs.repository }}
       ref: ${{ inputs.ref }}
@@ -389,5 +422,5 @@ jobs:
       architecture: ${{ inputs.architecture }}
 
 concurrency:
-  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ inputs.repository }}-${{ github.event_name == 'workflow_dispatch' }}
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ inputs.repository }}-${{ inputs.is-release-wheel }}-${{ inputs.is-release-tarball }}-${{ github.event_name == 'workflow_dispatch' }}
   cancel-in-progress: true

.github/workflows/release-linux-aarch64.yml

@@ -0,0 +1,132 @@
+name: Release aarch64 Linux wheels and tarball artifacts
+
+on:
+  push:
+    tags:
+      # NOTE: Binary build pipelines should only get triggered on release candidate builds
+      # Release candidate tags look like: v1.11.0-rc1
+      - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
+  workflow_dispatch:
+
+permissions:
+  id-token: write
+  contents: read
+  packages: write
+
+jobs:
+  generate-matrix:
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
+    if: ${{ contains(github.event.pull_request.labels.*.name, 'build-release-artifacts') || startsWith(github.event.ref, 'refs/tags/v') }}
+    with:
+      package-type: wheel
+      os: linux-aarch64
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      with-rocm: false
+      with-cpu: false
+
+  generate-release-tarball-matrix:
+    needs: [generate-matrix]
+    outputs:
+      matrix: ${{ steps.generate.outputs.matrix }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+      - uses: actions/checkout@v4
+        with:
+          repository: pytorch/tensorrt
+      - name: Generate release matrix
+        id: generate
+        run: |
+          set -eou pipefail
+          MATRIX_BLOB=${{ toJSON(needs.generate-matrix.outputs.matrix) }}
+          MATRIX_BLOB="$(python3 .github/scripts/generate-release-matrix.py --tarball_matrix "${MATRIX_BLOB}")"
+          echo "${MATRIX_BLOB}"
+          echo "matrix=${MATRIX_BLOB}" >> "${GITHUB_OUTPUT}"
+
+  release-cxx11-tarball-artifacts:
+    needs: [generate-release-tarball-matrix]
+    name: Release aarch64 torch-tensorrt cxx11 tarball artifacts
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - repository: pytorch/tensorrt
+            package-name: torch_tensorrt
+            pre-script: packaging/pre_build_script.sh
+            env-var-script: packaging/env_vars.txt
+            post-script: packaging/post_build_script.sh
+            smoke-test-script: packaging/smoke_test_script.sh
+            is-release-tarball: "true"
+    uses: ./.github/workflows/build_wheels_linux.yml
+    with:
+      repository: ${{ matrix.repository }}
+      ref: ""
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      build-matrix: ${{ needs.generate-release-tarball-matrix.outputs.matrix }}
+      pre-script: ${{ matrix.pre-script }}
+      env-var-script: ${{ matrix.env-var-script }}
+      post-script: ${{ matrix.post-script }}
+      package-name: ${{ matrix.package-name }}
+      smoke-test-script: ${{ matrix.smoke-test-script }}
+      trigger-event: ${{ github.event_name }}
+      is-release-tarball: true
+      architecture: "aarch64"
+
+  generate-release-wheel-matrix:
+    needs: [generate-matrix]
+    outputs:
+      matrix: ${{ steps.generate.outputs.matrix }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+      - uses: actions/checkout@v4
+        with:
+          repository: pytorch/tensorrt
+      - name: Generate release matrix
+        id: generate
+        run: |
+          set -eou pipefail
+          MATRIX_BLOB=${{ toJSON(needs.generate-matrix.outputs.matrix) }}
+          MATRIX_BLOB="$(python3 .github/scripts/generate-release-matrix.py --wheel_matrix "${MATRIX_BLOB}")"
+          echo "${MATRIX_BLOB}"
+          echo "matrix=${MATRIX_BLOB}" >> "${GITHUB_OUTPUT}"
+
+  release-wheel-artifacts:
+    name: Release aarch64 torch-tensorrt wheel artifacts
+    needs: [generate-release-wheel-matrix]
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - repository: pytorch/tensorrt
+            package-name: torch_tensorrt
+            pre-script: packaging/pre_build_script.sh
+            env-var-script: packaging/env_vars.txt
+            post-script: packaging/post_build_script.sh
+            smoke-test-script: packaging/smoke_test_script.sh
+            is-release-wheel:  true
+    uses: ./.github/workflows/build_wheels_linux.yml
+    with:
+      repository: ${{ matrix.repository }}
+      ref: ""
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      build-matrix: ${{ needs.generate-release-wheel-matrix.outputs.matrix }}
+      pre-script: ${{ matrix.pre-script }}
+      env-var-script: ${{ matrix.env-var-script }}
+      post-script: ${{ matrix.post-script }}
+      package-name: ${{ matrix.package-name }}
+      smoke-test-script: ${{ matrix.smoke-test-script }}
+      trigger-event: ${{ github.event_name }}
+      is-release-wheel: true
+      architecture: "aarch64"
+
+concurrency:
+  group: ${{ github.workflow }}-aarch64-${{ github.event.pull_request.number || github.ref_name }}-${{ inputs.repository }}-${{ github.event_name == 'workflow_dispatch' }}-${{ inputs.job-name }}
+  cancel-in-progress: true