Add gating mechanism to ensure PyTorch wheels pass tests before releasing

.github/workflows/build_portable_linux_pytorch_wheels.yml

@@ -17,10 +17,18 @@ on:
         description: S3 subdirectory, not including the GPU-family
         required: true
         type: string
+      s3_staging_subdir:
+        description: S3 staging subdirectory, not including the GPU-family
+        required: true
+        type: string
       cloudfront_url:
         description: CloudFront URL pointing to Python index
         required: true
         type: string
+      cloudfront_staging_url:
+        description: CloudFront base URL pointing to staging Python index
+        required: true
+        type: string
       rocm_version:
         description: ROCm version to pip install
         type: string
@@ -55,10 +63,18 @@ on:
         description: S3 subdirectory, not including the GPU-family
         type: string
         default: "v2"
+      s3_staging_subdir:
+        description: S3 staging subdirectory, not including the GPU-family
+        type: string
+        default: "v2-staging"
       cloudfront_url:
         description: CloudFront base URL pointing to Python index
         type: string
         default: "https://d25kgig7rdsyks.cloudfront.net/v2"
+      cloudfront_staging_url:
+        description: CloudFront base URL pointing to staging Python index
+        type: string
+        default: "https://d25kgig7rdsyks.cloudfront.net/v2-staging"
       rocm_version:
         description: ROCm version to pip install
         type: string
@@ -89,7 +105,11 @@ jobs:
       S3_BUCKET_PY: "therock-${{ inputs.release_type }}-python"
       optional_build_prod_arguments: ""
     outputs:
+      cp_version: ${{ env.cp_version }}
       torch_version: ${{ steps.build-pytorch-wheels.outputs.torch_version }}
+      torchaudio_version: ${{ steps.build-pytorch-wheels.outputs.torchaudio_version }}
+      torchvision_version: ${{ steps.build-pytorch-wheels.outputs.torchvision_version }}
+      triton_version: ${{ steps.build-pytorch-wheels.outputs.triton_version }}
     steps:
       - name: Checkout
         uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
@@ -167,23 +187,24 @@ jobs:
         run: |
           python external-builds/pytorch/sanity_check_wheel.py ${{ env.PACKAGE_DIST_DIR }}/
 
-      - name: Upload wheels to S3
+      - name: Upload wheels to S3 staging
         if: ${{ github.repository_owner == 'ROCm' }}
         run: |
-          aws s3 cp ${{ env.PACKAGE_DIST_DIR }}/ s3://${{ env.S3_BUCKET_PY }}/${{ inputs.s3_subdir }}/${{ inputs.amdgpu_family }}/ \
+          aws s3 cp ${{ env.PACKAGE_DIST_DIR }}/ s3://${{ env.S3_BUCKET_PY }}/${{ inputs.s3_staging_subdir }}/${{ inputs.amdgpu_family }}/ \
             --recursive --exclude "*" --include "*.whl"
 
-      - name: (Re-)Generate Python package release index
+      - name: (Re-)Generate Python package release index for staging
         if: ${{ github.repository_owner == 'ROCm' }}
         run: |
           pip install boto3 packaging
-          python ./build_tools/third_party/s3_management/manage.py ${{ inputs.s3_subdir }}/${{ inputs.amdgpu_family }}
+          python ./build_tools/third_party/s3_management/manage.py ${{ inputs.s3_staging_subdir }}/${{ inputs.amdgpu_family }}
 
   generate_target_to_run:
     name: Generate target_to_run
     runs-on: ubuntu-24.04
     outputs:
       test_runs_on: ${{ steps.configure.outputs.test-runs-on }}
+      bypass_tests_for_releases: ${{ steps.configure.outputs.bypass_tests_for_releases }}
     steps:
       - name: Checking out repository
         uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
@@ -203,7 +224,79 @@ jobs:
     with:
       amdgpu_family: ${{ inputs.amdgpu_family }}
       test_runs_on: ${{ needs.generate_target_to_run.outputs.test_runs_on }}
-      cloudfront_url: ${{ inputs.cloudfront_url }}
+      cloudfront_url: ${{ inputs.cloudfront_staging_url }}
       python_version: ${{ inputs.python_version }}
       torch_version: ${{ needs.build_pytorch_wheels.outputs.torch_version }}
       pytorch_version: ${{ inputs.pytorch_version }}
+
+  upload_pytorch_wheels:
+    name: Release PyTorch Wheels to S3
+    needs: [build_pytorch_wheels, generate_target_to_run, test_pytorch_wheels]
+    if: always()
+    runs-on: ubuntu-24.04
+    env:
+      S3_BUCKET_PY: "therock-${{ inputs.release_type }}-python"
+      CP_VERSION: "${{ needs.build_pytorch_wheels.outputs.cp_version }}"
+      TORCH_VERSION: "${{ needs.build_pytorch_wheels.outputs.torch_version }}"
+      TORCHAUDIO_VERSION: "${{ needs.build_pytorch_wheels.outputs.torchaudio_version }}"
+      TORCHVISION_VERSION: "${{ needs.build_pytorch_wheels.outputs.torchvision_version }}"
+      TRITON_VERSION: "${{ needs.build_pytorch_wheels.outputs.triton_version }}"
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+
+      - name: Configure AWS Credentials
+        if: always()
+        uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4.3.1
+        with:
+          aws-region: us-east-2
+          role-to-assume: arn:aws:iam::692859939525:role/therock-${{ inputs.release_type }}-releases
+
+
+      - name: Determine upload flag
+        env:
+          BUILD_RESULT: ${{ needs.build_pytorch_wheels.result }}
+          TEST_RESULT: ${{ needs.test_pytorch_wheels.result }}
+          TEST_RUNS_ON: ${{ needs.generate_target_to_run.outputs.test_runs_on }}
+          BYPASS_TESTS_FOR_RELEASES: ${{ needs.generate_target_to_run.outputs.bypass_tests_for_releases }}
+        run: |
+          # 1) If the build failed → upload=false
+          if [[ "$BUILD_RESULT" != "success" ]]; then
+            echo "::warning::Build failed. Skipping upload."
+            echo "upload=false" >> "$GITHUB_ENV"
+
+          # 2) Else if there was a test runner AND tests failed or were skipped → upload=false
+          elif [[ -n "$TEST_RUNS_ON" && ( "$TEST_RESULT" == "failure" || "$TEST_RESULT" == "skipped" ) ]]; then
+            echo "::warning::Tests failed or were skipped (runner present). Skipping upload."
+            echo "upload=false" >> "$GITHUB_ENV"
+
+          # 3) Else if BYPASS_TESTS_FOR_RELEASES is not set and there was no test runner → upload=false
+          elif [[ -z "$BYPASS_TESTS_FOR_RELEASES" && -z "$TEST_RUNS_ON" ]]; then
+            echo "::warning::No test runner and BYPASS_TESTS_FOR_RELEASES not set. Skipping upload."
+            echo "upload=false" >> "$GITHUB_ENV"
+
+          # 4) Otherwise → upload=true
+          else
+            echo "upload=true" >> "$GITHUB_ENV"
+          fi
+
+      - name: Copy PyTorch wheels from staging to release S3
+        if: ${{ env.upload == 'true' }}
+        run: |
+          echo "Copying exact tested wheels to release S3 bucket..."
+          aws s3 cp \
+            s3://${S3_BUCKET_PY}/${{ inputs.s3_staging_subdir }}/${{ inputs.amdgpu_family }}/ \
+            s3://${S3_BUCKET_PY}/${{ inputs.s3_subdir }}/${{ inputs.amdgpu_family }}/ \
+            --recursive \
+            --exclude "*" \
+            --include "torch-${TORCH_VERSION}-${CP_VERSION}-linux_x86_64.whl" \
+            --include "torchaudio-${TORCHAUDIO_VERSION}-${CP_VERSION}-linux_x86_64.whl" \
+            --include "torchvision-${TORCHVISION_VERSION}-${CP_VERSION}-linux_x86_64.whl" \
+            --include "pytorch_triton_rocm-${TRITON_VERSION}-${CP_VERSION}-linux_x86_64.whl"
+
+      - name: (Re-)Generate Python package release index
+        if: ${{ env.upload == 'true' }}
+        run: |
+          pip install boto3 packaging
+          python ./build_tools/third_party/s3_management/manage.py ${{ inputs.s3_subdir }}/${{ inputs.amdgpu_family }}

.github/workflows/release_portable_linux_packages.yml

@@ -14,6 +14,10 @@ on:
         description: "Subdirectory to push the Python packages"
         type: string
         default: "v2"
+      s3_staging_subdir:
+        description: "Staging subdirectory to push the Python packages"
+        type: string
+        default: "v2-staging"
   # Trigger manually (typically to test the workflow or manually build a release [candidate])
   workflow_dispatch:
     inputs:
@@ -27,6 +31,10 @@ on:
         description: "Subdirectory to push the Python packages"
         type: string
         default: "v2"
+      s3_staging_subdir:
+        description: "Staging subdirectory to push the Python packages"
+        type: string
+        default: "v2-staging"
       families:
         description: "Comma separated list of AMD GPU families, e.g. `gfx94X,gfx103x`"
         type: string
@@ -44,6 +52,7 @@ jobs:
     runs-on: ubuntu-24.04
     env:
       S3_SUBDIR: ${{ inputs.s3_subdir || 'v2' }}
+      S3_STAGING_SUBDIR: ${{ inputs.s3_staging_subdir || 'v2-staging' }}
       release_type: ${{ inputs.release_type || 'nightly' }}
     outputs:
       version: ${{ steps.release_information.outputs.version }}
@@ -109,6 +118,7 @@ jobs:
       S3_BUCKET_TAR: "therock-${{ needs.setup_metadata.outputs.release_type }}-tarball"
       S3_BUCKET_PY: "therock-${{ needs.setup_metadata.outputs.release_type }}-python"
       S3_SUBDIR: ${{ inputs.s3_subdir || 'v2' }}
+      S3_STAGING_SUBDIR: ${{ inputs.s3_staging_subdir || 'v2-staging' }}
 
     steps:
       - name: "Checking out repository"
@@ -147,6 +157,11 @@ jobs:
           echo "Building ${{ env.DIST_ARCHIVE }}"
           tar cfz "${{ env.DIST_ARCHIVE }}" .
 
+      - name: Setup Python
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: 3.12
+
       - name: Build Python Packages
         run: |
           ./build_tools/linux_portable_build.py \
@@ -171,6 +186,22 @@ jobs:
           aws-region: us-east-2
           role-to-assume: arn:aws:iam::692859939525:role/therock-${{ env.RELEASE_TYPE }}-releases
 
+      - name: Upload Releases to staging S3
+        if: ${{ github.repository_owner == 'ROCm' }}
+        run: |
+          aws s3 cp ${{ env.OUTPUT_DIR }}/packages/dist/ s3://${{ env.S3_BUCKET_PY }}/${{ env.S3_STAGING_SUBDIR }}/${{ matrix.target_bundle.amdgpu_family }}/ \
+          --recursive --no-follow-symlinks \
+          --exclude "*" \
+          --include "*.whl" \
+          --include "*.tar.gz"
+
+      - name: (Re-)Generate Python package release index for staging
+        if: ${{ github.repository_owner == 'ROCm' }}
+        run: |
+          pip install boto3 packaging
+          python ./build_tools/third_party/s3_management/manage.py ${{ env.S3_STAGING_SUBDIR }}/${{ matrix.target_bundle.amdgpu_family }}
+
+      ## TODO: Restrict uploading to the non-staging S3 directory until ROCm sanity checks and all validation tests have successfully passed.
       - name: Upload Releases to S3
         if: ${{ github.repository_owner == 'ROCm' }}
         run: |
@@ -181,12 +212,6 @@ jobs:
           --include "*.whl" \
           --include "*.tar.gz"
 
-      - name: Setup Python
-        if: ${{ github.repository_owner == 'ROCm' }}
-        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
-        with:
-          python-version: 3.12
-
       - name: (Re-)Generate Python package release index
         if: ${{ github.repository_owner == 'ROCm' }}
         run: |

build_tools/github_actions/amdgpu_family_matrix.py

@@ -13,10 +13,12 @@
         "linux": {
             "test-runs-on": "",
             "family": "gfx110X-dgpu",
+            "bypass_tests_for_releases": True,
         },
         "windows": {
             "test-runs-on": "",
             "family": "gfx110X-dgpu",
+            "bypass_tests_for_releases": True,
         },
     },
 }
@@ -32,6 +34,7 @@
         "linux": {
             "test-runs-on": "",
             "family": "gfx1151",
+            "bypass_tests_for_releases": True,
         },
         "windows": {
             "test-runs-on": "windows-strix-halo-gpu-rocm",
@@ -42,10 +45,12 @@
         "linux": {
             "test-runs-on": "",  # removed due to machine issues, label is "linux-rx9070-gpu-rocm"
             "family": "gfx120X-all",
+            "bypass_tests_for_releases": True,
         },
         "windows": {
             "test-runs-on": "",
             "family": "gfx120X-all",
+            "bypass_tests_for_releases": True,
         },
     },
 }

build_tools/github_actions/configure_target_run.py

@@ -46,14 +46,50 @@ def get_runner_label(target: str, platform: str) -> str:
         if test_runs_on_machine:
             print(f"  Found runner: '{test_runs_on_machine}'")
             return test_runs_on_machine
+    return ""
+
+
+def get_upload_label(target: str, platform: str) -> str:
+    print(f"Searching for a runner for target '{target}' on platform '{platform}'")
+    amdgpu_family_info_matrix = (
+        amdgpu_family_info_matrix_presubmit | amdgpu_family_info_matrix_postsubmit
+    )
+    for key, info_for_key in amdgpu_family_info_matrix.items():
+        print(f"Cheecking key '{key}' with info:\n  {info_for_key}")
+        platform_for_key = info_for_key.get(platform)
+
+        if not platform_for_key:
+            # Some AMDGPU families are only supported on certain platforms.
+            print(f"  Skipping since this entry has no platform '{platform}'")
+            continue
+
+        # Check against both the inner "family" and the outer "key". If neither
+        # match then skip. Workflows are expected to use the inner "family"
+        # but manually triggered runs may use the outer "key" instead, so we'll
+        # be a bit lenient here.
+        # This needs a rework, see https://github.com/ROCm/TheRock/issues/1097.
+        family_for_platform = platform_for_key.get("family")
+        if target != family_for_platform and key not in target.lower():
+            print(
+                f"  Skipping since the target '{target}' does not match the family '{family_for_platform}'"
+            )
+            continue
 
+        # If there is no test machine available and bypass_tests_for_releases flag is True for GPU family and platform, output bypass_tests_for_releases as True
+        bypass_tests_for_releases = platform_for_key.get("bypass_tests_for_releases")
+        if bypass_tests_for_releases:
+            print(f"  bypass_tests_for_releases: True")
+            return bypass_tests_for_releases
     return ""
 
 
 def main(target: str, platform: str):
     runner_label = get_runner_label(target, platform)
     if runner_label:
         gha_set_output({"test-runs-on": runner_label})
+    upload_label = get_upload_label(target, platform)
+    if upload_label:
+        gha_set_output({"bypass_tests_for_releases": upload_label})
 
 
 if __name__ == "__main__":

build_tools/github_actions/write_torch_versions.py

@@ -11,7 +11,6 @@
 import os
 import glob
 import platform
-
 from github_actions_utils import *
 
 
@@ -100,10 +99,8 @@ def main(argv: list[str]):
         help="Path where wheels are located",
     )
     args = p.parse_args(argv)
-
     if not args.dist_dir.exists():
         raise FileNotFoundError(f"Dist dir '{args.dist_dir}' does not exist")
-
     all_versions = get_all_wheel_versions(args.dist_dir)
     _log("")
     gha_set_output(all_versions)

build_tools/third_party/s3_management/manage.py

@@ -40,6 +40,11 @@
     "v2/gfx120X-all",
     "v2/gfx94X-dcgpu",
     "v2/gfx950-dcgpu",
+    "v2-staging/gfx110X-dgpu",
+    "v2-staging/gfx1151",
+    "v2-staging/gfx120X-all",
+    "v2-staging/gfx94X-dcgpu",
+    "v2-staging/gfx950-dcgpu"
 ]
 
 CUSTOM_PREFIX = getenv('CUSTOM_PREFIX')