[OMNIML-6183] [2/2] Run the GPU, example and regression tests in the merge queue only

.github/actions/pr-merge-base/action.yml

@@ -1,16 +1,23 @@
 name: PR merge base
 description: >
-  Resolve the commit to diff a copied PR branch against: the merge base of the PR's head and its
-  target branch. Outputs are empty for non-PR triggers, which have no diff to inspect.
+  Resolve the commit to diff a PR run against. For a copied PR branch it is the merge base of the
+  PR's head and its target branch; for a merge queue group it is the commit the group was built
+  on. Outputs are empty for non-PR triggers, which have no diff to inspect. Also says whether the
+  merge tests (the GPU, example and regression tests, which a PR run skips) run: always in a merge
+  queue group, and on a copied PR branch only when the PR carries the `run-merge-tests` label.
   Requires a prior actions/checkout with fetch-depth 0.
 
 outputs:
   merge_base:
     description: Commit to use as the base for changed-file comparisons.
-    value: ${{ steps.calculate-merge-base.outputs.merge-base }}
+    value: ${{ steps.calculate-merge-base.outputs.merge-base || steps.merge-group.outputs.base_sha }}
   head_sha:
-    description: The PR's head commit.
-    value: ${{ steps.pr-shas.outputs.head_sha }}
+    description: The PR's head commit, or the merge queue group's head commit.
+    value: ${{ steps.pr-shas.outputs.head_sha || steps.merge-group.outputs.head_sha }}
+  run_merge_tests:
+    description: >-
+      "true" if the merge tests run in this PR run, else "false". Empty for non-PR triggers.
+    value: ${{ steps.pr-shas.outputs.run_merge_tests || steps.merge-group.outputs.run_merge_tests }}
 
 runs:
   using: composite
@@ -27,6 +34,8 @@ runs:
       run: |
         echo "head_sha=$(echo "$PR_INFO" | jq -r '.head.sha')" >> $GITHUB_OUTPUT
         echo "base_sha=$(echo "$PR_INFO" | jq -r '.base.sha')" >> $GITHUB_OUTPUT
+        echo "run_merge_tests=$(echo "$PR_INFO" |
+          jq -r '[(.labels // [])[].name] | index("run-merge-tests") != null')" >> $GITHUB_OUTPUT
     # Get commit from the target branch that is present in the PR to use as base for changed files
     - if: startsWith(github.ref, 'refs/heads/pull-request/')
       id: calculate-merge-base
@@ -36,3 +45,15 @@ runs:
         # status and emit an empty base, silently changing which lanes run.
         merge_base=$(git merge-base "${{ steps.pr-shas.outputs.base_sha }}" "${{ steps.pr-shas.outputs.head_sha }}")
         echo "merge-base=$merge_base" | tee --append "${GITHUB_OUTPUT}"
+    # A merge queue group is the target branch plus the queued PRs on top of it, so its diff
+    # against the commit it was built on is what the queue is about to merge.
+    - if: github.event_name == 'merge_group'
+      id: merge-group
+      shell: bash
+      env:
+        BASE_SHA: ${{ github.event.merge_group.base_sha }}
+        HEAD_SHA: ${{ github.event.merge_group.head_sha }}
+      run: |
+        echo "base_sha=$BASE_SHA" | tee --append "${GITHUB_OUTPUT}"
+        echo "head_sha=$HEAD_SHA" | tee --append "${GITHUB_OUTPUT}"
+        echo "run_merge_tests=true" | tee --append "${GITHUB_OUTPUT}"

.github/workflows/_example_tests_runner.yml

@@ -107,7 +107,9 @@ jobs:
             # Keep legacy ONNX/TRT tests available manually during deprecation.
             pytest_args+=(--ignore=tests/examples/diffusers/test_diffusers_onnx_trt.py)
           fi
-          if [[ "${{ inputs.example }}" == "hf_ptq" && "${{ github.ref }}" == refs/heads/pull-request/* ]]; then
+          if [[ "${{ inputs.example }}" == "hf_ptq" &&
+                ( "${{ github.ref }}" == refs/heads/pull-request/* ||
+                  "${{ github.event_name }}" == "merge_group" ) ]]; then
             # One GPU on PR runners: small-model tests share it and the rest run alone on it
             # (tests/examples/hf_ptq/conftest.py). The 2-GPU nightly stays serial.
             pytest_args+=(-n 4)
@@ -118,6 +120,7 @@ jobs:
         run: |
           echo "::warning title=Allowed example failure::'${{ inputs.example }}' failed but is in the allow-failure list (vars.ALLOW_FAILURE_EXAMPLE_TESTS); not blocking. Remove it from the variable once fixed."
       - name: Upload coverage to Codecov
+        if: github.event_name != 'merge_group'
         uses: codecov/codecov-action@v7
         with:
           token: ${{ secrets.CODECOV_TOKEN }}

.github/workflows/_pr_gate.yml

@@ -20,25 +20,33 @@ on:
       any_changed:
         description: "Whether any relevant files changed"
         value: ${{ jobs.check-file-changes.outputs.any_changed }}
+      run_merge_tests:
+        description: >-
+          Whether the merge tests (the GPU, example and regression tests) run: always except on a PR
+          run without the run-merge-tests label
+        value: ${{ jobs.check-file-changes.outputs.run_merge_tests }}
 
 jobs:
   check-file-changes:
     runs-on: ubuntu-latest
     outputs:
       any_changed: ${{ steps.changed-tests.outputs.any_modified || steps.non-pr.outputs.any_changed }}
+      # Empty outside PR runs (nightly, on demand), which run everything
+      run_merge_tests: ${{ steps.base.outputs.run_merge_tests || 'true' }}
     steps:
-      # For non-PR triggers (schedule, workflow_dispatch), always run tests
+      # For non-PR triggers (schedule, workflow_dispatch), always run tests. A PR run is a push to a
+      # copy-pr-bot branch (pull-request/N) or a merge queue group, which both diff against the base.
       - id: non-pr
-        if: "!startsWith(github.ref, 'refs/heads/pull-request/')"
+        if: "!(startsWith(github.ref, 'refs/heads/pull-request/') || github.event_name == 'merge_group')"
         run: echo "any_changed=true" >> $GITHUB_OUTPUT
-      - if: startsWith(github.ref, 'refs/heads/pull-request/')
+      - if: startsWith(github.ref, 'refs/heads/pull-request/') || github.event_name == 'merge_group'
         uses: actions/checkout@v6
         with:
           fetch-depth: 0
-      - if: startsWith(github.ref, 'refs/heads/pull-request/')
+      - if: startsWith(github.ref, 'refs/heads/pull-request/') || github.event_name == 'merge_group'
         id: base
         uses: ./.github/actions/pr-merge-base
-      - if: startsWith(github.ref, 'refs/heads/pull-request/')
+      - if: startsWith(github.ref, 'refs/heads/pull-request/') || github.event_name == 'merge_group'
         name: Check for changes in test-relevant directories
         id: changed-tests
         uses: step-security/changed-files@v47.0.5
@@ -50,9 +58,11 @@ jobs:
           fail_on_initial_diff_error: true
   wait-checks:
     needs: [check-file-changes]
+    # Nothing to wait for when the merge tests are skipped
     if: >-
-      startsWith(github.ref, 'refs/heads/pull-request/') &&
-      needs.check-file-changes.outputs.any_changed == 'true'
+      (startsWith(github.ref, 'refs/heads/pull-request/') || github.event_name == 'merge_group') &&
+      needs.check-file-changes.outputs.any_changed == 'true' &&
+      needs.check-file-changes.outputs.run_merge_tests == 'true'
     runs-on: ubuntu-latest
     permissions:
       checks: read

.github/workflows/code_quality.yml

@@ -3,6 +3,10 @@ name: Code Quality
 on:
   pull_request:
     branches: [main, release/*, feature/*]
+  # The merge queue re-runs the required checks on the PR merged with the target branch
+  merge_group:
+    types: [checks_requested]
+    branches: [main, release/*, feature/*]
   schedule:
     - cron: "0 0 * * *" # Nightly
   workflow_dispatch:
@@ -21,12 +25,18 @@ jobs:
     steps:
       - uses: actions/checkout@v6
         with:
-          # PRs: TruffleHog needs the base commit, else it errors, scans nothing, and passes.
-          fetch-depth: ${{ github.event_name == 'pull_request' && '0' || '1' }}
+          # PRs and merge queue groups: TruffleHog needs the base commit, else it errors,
+          # scans nothing, and passes.
+          fetch-depth: >-
+            ${{ (github.event_name == 'pull_request' || github.event_name == 'merge_group') && '0' || '1' }}
       - uses: ./.github/actions/ubuntu-setup
       - name: Secret Scanning
         uses: trufflesecurity/trufflehog@v3.90.5
         with:
+          # Only set in a merge queue group, where the diff to scan is base..head (empty for the other events,
+          # which the action handles itself)
+          base: ${{ github.event.merge_group.base_sha }}
+          head: ${{ github.event.merge_group.head_sha }}
           # Exclude lob: its `(live|test)_[a-zA-Z0-9_]{35}` pattern matches any pytest function
           # whose name happens to be exactly 35 characters after `test_`, and it reports them as
           # verified. Only nightly sees them, since it scans all history while PRs scan their diff.

.github/workflows/example_tests.yml

@@ -4,6 +4,10 @@ on:
   push:
     branches: ["pull-request/[0-9]+"]
     # NOTE: paths cannot be used since push happens to copied PR and only latest commit to PR is used
+  # The merge queue re-runs the required checks on the PR merged with the target branch
+  merge_group:
+    types: [checks_requested]
+    branches: [main, release/*, feature/*]
   schedule:
     - cron: "0 0 * * *" # Nightly
   workflow_dispatch:
@@ -27,6 +31,7 @@ jobs:
       contents: read
     outputs:
       any_changed: ${{ steps.lanes.outputs.any_changed }}
+      run_merge_tests: ${{ steps.base.outputs.run_merge_tests || 'true' }}
       torch: ${{ steps.lanes.outputs.torch }}
       vllm: ${{ steps.lanes.outputs.vllm }}
       trtllm: ${{ steps.lanes.outputs.trtllm }}
@@ -37,10 +42,10 @@ jobs:
       - uses: actions/checkout@v6
         with:
           fetch-depth: 0
-      - if: startsWith(github.ref, 'refs/heads/pull-request/')
+      - if: startsWith(github.ref, 'refs/heads/pull-request/') || github.event_name == 'merge_group'
         id: base
         uses: ./.github/actions/pr-merge-base
-      - if: startsWith(github.ref, 'refs/heads/pull-request/')
+      - if: startsWith(github.ref, 'refs/heads/pull-request/') || github.event_name == 'merge_group'
         id: changed
         uses: step-security/changed-files@v47.0.5
         with:
@@ -87,7 +92,10 @@ jobs:
       - id: lanes
         env:
           # Nightly and on-demand runs have no diff to inspect, so they run everything.
-          NON_PR: ${{ !startsWith(github.ref, 'refs/heads/pull-request/') }}
+          NON_PR: >-
+            ${{ !(startsWith(github.ref, 'refs/heads/pull-request/') || github.event_name == 'merge_group') }}
+          # The example tests are merge tests: a PR run without the run-merge-tests label skips them
+          RUN_MERGE_TESTS: ${{ steps.base.outputs.run_merge_tests || 'true' }}
           COMMON: ${{ steps.changed.outputs.common_any_modified }}
           TORCH: ${{ steps.changed.outputs.torch_any_modified }}
           VLLM: ${{ steps.changed.outputs.vllm_any_modified }}
@@ -98,7 +106,8 @@ jobs:
         run: |
           any=false
           for lane in TORCH VLLM TRTLLM MEGATRON DIFFUSERS ONNX; do
-            if [[ "$NON_PR" == "true" || "$COMMON" == "true" || "${!lane}" == "true" ]]; then
+            if [[ "$RUN_MERGE_TESTS" == "true" &&
+                  ("$NON_PR" == "true" || "$COMMON" == "true" || "${!lane}" == "true") ]]; then
               run=true
               any=true
             else
@@ -108,7 +117,9 @@ jobs:
           done
           echo "any_changed=$any" | tee --append "$GITHUB_OUTPUT"
       - name: Wait for unit tests before spending GPU runners
-        if: ${{ steps.lanes.outputs.any_changed == 'true' && startsWith(github.ref, 'refs/heads/pull-request/') }}
+        if: >-
+          ${{ steps.lanes.outputs.any_changed == 'true' &&
+          (startsWith(github.ref, 'refs/heads/pull-request/') || github.event_name == 'merge_group') }}
         uses: poseidon/wait-for-status-checks@v0.7.0
         with:
           token: ${{ secrets.GITHUB_TOKEN }}
@@ -132,7 +143,9 @@ jobs:
       example: ${{ matrix.example }}
       timeout_minutes: 30
       pip_install_extras: "[hf,dev-test]"
-      runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
+      runner: >-
+        ${{ (startsWith(github.ref, 'refs/heads/pull-request/') || github.event_name == 'merge_group') &&
+        'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
       allow_failure: ${{ contains(format(',{0},', vars.ALLOW_FAILURE_EXAMPLE_TESTS), format(',{0},', matrix.example)) }}
 
   ##### vLLM Example Tests #####
@@ -153,7 +166,9 @@ jobs:
       example: ${{ matrix.example }}
       timeout_minutes: 30
       pip_install_extras: "[hf,dev-test]"
-      runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
+      runner: >-
+        ${{ (startsWith(github.ref, 'refs/heads/pull-request/') || github.event_name == 'merge_group') &&
+        'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
       allow_failure: ${{ contains(format(',{0},', vars.ALLOW_FAILURE_EXAMPLE_TESTS), format(',{0},', matrix.example)) }}
 
   ##### TensorRT-LLM Example Tests #####
@@ -172,7 +187,9 @@ jobs:
       docker_image: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20"
       example: ${{ matrix.example }}
       pip_install_extras: "[hf,dev-test]"
-      runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
+      runner: >-
+        ${{ (startsWith(github.ref, 'refs/heads/pull-request/') || github.event_name == 'merge_group') &&
+        'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
       allow_failure: ${{ contains(format(',{0},', vars.ALLOW_FAILURE_EXAMPLE_TESTS), format(',{0},', matrix.example)) }}
 
   ##### Megatron Example Tests #####
@@ -187,9 +204,13 @@ jobs:
       docker_image: "nvcr.io/nvidia/nemo:26.08"
       example: megatron_bridge
       # PR runs single-GPU (~11 min); the nightly runs every test multi-GPU (~58 min).
-      timeout_minutes: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 30 || 75 }}
+      timeout_minutes: >-
+        ${{ (startsWith(github.ref, 'refs/heads/pull-request/') ||
+        github.event_name == 'merge_group') && 30 || 75 }}
       pip_install_extras: "[hf,puzzletron,dev-test]"
-      runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
+      runner: >-
+        ${{ (startsWith(github.ref, 'refs/heads/pull-request/') || github.event_name == 'merge_group') &&
+        'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
       allow_failure: ${{ contains(format(',{0},', vars.ALLOW_FAILURE_EXAMPLE_TESTS), ',megatron_bridge,') }}
 
   ##### Diffusers Example Tests #####
@@ -206,7 +227,9 @@ jobs:
       example: diffusers
       timeout_minutes: 60
       pip_install_extras: "[onnx,hf,dev-test]"
-      runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
+      runner: >-
+        ${{ (startsWith(github.ref, 'refs/heads/pull-request/') || github.event_name == 'merge_group') &&
+        'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
       allow_failure: ${{ contains(format(',{0},', vars.ALLOW_FAILURE_EXAMPLE_TESTS), ',diffusers,') }}
 
   ##### ONNX/TensorRT Example Tests #####
@@ -229,16 +252,27 @@ jobs:
       example: ${{ matrix.example }}
       timeout_minutes: 45
       pip_install_extras: "[onnx,hf,dev-test]"
-      runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
+      runner: >-
+        ${{ (startsWith(github.ref, 'refs/heads/pull-request/') || github.event_name == 'merge_group') &&
+        'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
       allow_failure: ${{ contains(format(',{0},', vars.ALLOW_FAILURE_EXAMPLE_TESTS), format(',{0},', matrix.example)) }}
 
   ##### Required Check for PR #####
   example-pr-required-check:
     # Run even if example tests are skipped
-    if: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && always() }}
+    if: >-
+      ${{ (startsWith(github.ref, 'refs/heads/pull-request/') ||
+      github.event_name == 'merge_group') && always() }}
     needs: [pr-gate, torch, vllm, trtllm, megatron, diffusers, onnx]
     runs-on: ubuntu-latest
+    permissions:
+      pull-requests: write
     steps:
+      - name: Example tests are left to the merge queue
+        if: ${{ needs.pr-gate.outputs.run_merge_tests != 'true' }}
+        run: >-
+          echo "::notice::The example tests are merge tests and run in the merge queue. Add the
+          run-merge-tests label and re-run this workflow to run them on this PR."
       - name: Required example tests did not succeed
         if: |
           needs.pr-gate.result != 'success' ||
@@ -251,3 +285,15 @@ jobs:
             !contains(fromJSON('["success", "skipped"]'), needs.onnx.result)
           ))
         run: exit 1
+      - name: Remove the run-merge-tests label, it applies to one round of runs
+        if: >-
+          ${{ !cancelled() && startsWith(github.ref, 'refs/heads/pull-request/') &&
+          needs.pr-gate.outputs.any_changed == 'true' && needs.pr-gate.outputs.run_merge_tests == 'true' }}
+        env:
+          GH_TOKEN: ${{ github.token }}
+        run: |
+          pr="${GITHUB_REF_NAME#pull-request/}"
+          url="repos/${GITHUB_REPOSITORY}/issues/${pr}/labels/run-merge-tests"
+          if gh api -X DELETE "$url" --silent; then
+            echo "::notice::Removed the run-merge-tests label after this round; re-add it for the next push."
+          fi