reuse the generate binary builds from test-infra

.github/scripts/filter-matrix.py

@@ -2,9 +2,11 @@
 
 import argparse
 import json
+import os
 import sys
 from typing import List
 
+# this was introduced to avoid build for py3.13, currently we support py3.13, but keep this for future use with other python versions
 disabled_python_versions: List[str] = []
 
 
@@ -17,15 +19,34 @@ def main(args: list[str]) -> None:
         default="",
     )
 
+    parser.add_argument(
+        "--limit-pr-builds",
+        help="Limit PR builds to single python/cuda config(py3.11/cu12.8): true or false",
+        type=str,
+        choices=["true", "false"],
+        default=os.getenv("LIMIT_PR_BUILDS", "false"),
+    )
+
     options = parser.parse_args(args)
 
     if options.matrix == "":
         raise Exception("--matrix needs to be provided")
+
     matrix_dict = json.loads(options.matrix)
     includes = matrix_dict["include"]
     filtered_includes = []
     for item in includes:
-        if item["python_version"] not in disabled_python_versions:
+        if item["python_version"] in disabled_python_versions:
+            continue
+        if options.limit_pr_builds == "true":
+            # currently if it is the pr build, it build using py3.9 with all cuda versions, we want to change to py3.11 with singlecu12.8
+            if item["desired_cuda"] == "cu128":
+                item["python_version"] = "3.11"
+                build_names = item["build_name"].split("-")
+                build_names[1] = "py3_11"
+                item["build_name"] = "-".join(build_names)
+                filtered_includes.append(item)
+        else:
             filtered_includes.append(item)
     filtered_matrix_dict = {}
     filtered_matrix_dict["include"] = filtered_includes

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

@@ -15,7 +15,7 @@ on:
 
 jobs:
   generate-matrix:
-    uses: ./.github/workflows/generate_binary_build_matrix.yml
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
     with:
       package-type: wheel
       os: linux
@@ -26,6 +26,8 @@ jobs:
 
   filter-matrix:
     needs: [generate-matrix]
+    env:
+      LIMIT_PR_BUILDS: ${{ github.event_name == 'pull_request' && !contains( github.event.pull_request.labels.*.name, 'ciflow/binaries/all') }}
     outputs:
       matrix: ${{ steps.generate.outputs.matrix }}
     runs-on: ubuntu-latest
@@ -36,7 +38,7 @@ jobs:
       - uses: actions/checkout@v4
         with:
           repository: pytorch/tensorrt
-      - name: Generate release matrix
+      - name: Generate matrix
         id: generate
         run: |
           set -eou pipefail

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

@@ -12,7 +12,7 @@ permissions:
 
 jobs:
   generate-matrix:
-    uses: ./.github/workflows/generate_binary_build_matrix.yml
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
     with:
       package-type: wheel
       os: linux

.github/workflows/build-test-tensorrt-windows.yml

@@ -12,7 +12,7 @@ permissions:
 
 jobs:
   generate-matrix:
-    uses: ./.github/workflows/generate_binary_build_matrix.yml
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
     with:
       package-type: wheel
       os: windows

.github/workflows/build-test-windows.yml

@@ -15,7 +15,7 @@ on:
 
 jobs:
   generate-matrix:
-    uses: ./.github/workflows/generate_binary_build_matrix.yml
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
     with:
       package-type: wheel
       os: windows
@@ -24,16 +24,39 @@ jobs:
       with-rocm: false
       with-cpu: false
 
+  filter-matrix:
+    needs: [generate-matrix]
+    env:
+      LIMIT_PR_BUILDS: ${{ github.event_name == 'pull_request' && !contains( github.event.pull_request.labels.*.name, 'ciflow/binaries/all') }}
+    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 matrix
+        id: generate
+        run: |
+          set -eou pipefail
+          MATRIX_BLOB=${{ toJSON(needs.generate-matrix.outputs.matrix) }}
+          MATRIX_BLOB="$(python3 .github/scripts/filter-matrix.py --matrix "${MATRIX_BLOB}")"
+          echo "${MATRIX_BLOB}"
+          echo "matrix=${MATRIX_BLOB}" >> "${GITHUB_OUTPUT}"
+
   substitute-runner:
-    needs: generate-matrix
+    needs: filter-matrix
     outputs:
       matrix: ${{ steps.substitute.outputs.matrix }}
     runs-on: ubuntu-latest
     steps:
       - name: Substitute runner
         id: substitute
         run: |
-          echo matrix="$(echo '${{ needs.generate-matrix.outputs.matrix }}' | sed -e 's/windows.g4dn.xlarge/windows.g5.4xlarge.nvidia.gpu/g')" >> ${GITHUB_OUTPUT}
+          echo matrix="$(echo '${{ needs.filter-matrix.outputs.matrix }}' | sed -e 's/windows.g4dn.xlarge/windows.g5.4xlarge.nvidia.gpu/g')" >> ${GITHUB_OUTPUT}
 
   build:
     needs: substitute-runner

.github/workflows/release-linux.yml

@@ -15,7 +15,7 @@ permissions:
 
 jobs:
   generate-matrix:
-    uses: ./.github/workflows/generate_binary_build_matrix.yml
+    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

.github/workflows/release-windows.yml

@@ -15,7 +15,7 @@ permissions:
 
 jobs:
   generate-matrix:
-    uses: ./.github/workflows/generate_binary_build_matrix.yml
+    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