Enable /rerun-stage slash command to work on fork PRs (#16128)

This commit is contained in:
Alison Shao
2026-01-02 10:36:12 -08:00
committed by GitHub
parent f07e76b229
commit 17041f4673
3 changed files with 94 additions and 56 deletions
+21 -14
View File
@@ -24,6 +24,11 @@ on:
required: false
type: string
default: ""
pr_head_sha:
description: "PR head SHA to checkout (for /rerun-stage on fork PRs)"
required: false
type: string
default: ""
workflow_call:
inputs:
ref:
@@ -36,8 +41,10 @@ on:
required: false
type: boolean
default: false
concurrency:
group: pr-test-amd-${{ inputs.ref || github.ref }}
# Include pr_head_sha in group for /rerun-stage dispatches to avoid collisions with main branch runs
group: pr-test-amd-${{ inputs.pr_head_sha || inputs.ref || github.ref }}
cancel-in-progress: ${{ github.event_name != 'workflow_call' }}
jobs:
@@ -55,7 +62,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Determine run mode
id: run-mode
@@ -110,7 +117,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Ensure VRAM is clear
run: bash scripts/ensure_vram_clear.sh rocm
@@ -158,7 +165,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Ensure VRAM is clear
run: bash scripts/ensure_vram_clear.sh rocm
@@ -191,7 +198,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Ensure VRAM is clear
run: bash scripts/ensure_vram_clear.sh rocm
@@ -312,7 +319,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Ensure VRAM is clear
run: bash scripts/ensure_vram_clear.sh rocm
@@ -439,7 +446,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Ensure VRAM is clear
run: bash scripts/ensure_vram_clear.sh rocm
@@ -479,7 +486,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Ensure VRAM is clear
run: bash scripts/ensure_vram_clear.sh rocm
@@ -521,7 +528,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Ensure VRAM is clear
run: bash scripts/ensure_vram_clear.sh rocm
@@ -566,7 +573,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Ensure VRAM is clear
run: bash scripts/ensure_vram_clear.sh rocm
@@ -621,7 +628,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Ensure VRAM is clear
run: bash scripts/ensure_vram_clear.sh rocm
@@ -670,7 +677,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Ensure VRAM is clear
run: bash scripts/ensure_vram_clear.sh rocm
@@ -729,7 +736,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Ensure VRAM is clear
run: bash scripts/ensure_vram_clear.sh rocm
@@ -771,7 +778,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Ensure VRAM is clear
run: bash scripts/ensure_vram_clear.sh rocm
+43 -35
View File
@@ -25,6 +25,11 @@ on:
required: false
type: boolean
default: false
pr_head_sha:
description: "PR head SHA to checkout (for /rerun-stage on fork PRs)"
required: false
type: string
default: ""
workflow_call:
inputs:
ref:
@@ -39,7 +44,8 @@ on:
default: false
concurrency:
group: pr-test-${{ inputs.ref || github.ref }}
# Include pr_head_sha in group for /rerun-stage dispatches to avoid collisions with main branch runs
group: pr-test-${{ inputs.pr_head_sha || inputs.ref || github.ref }}
cancel-in-progress: ${{ github.event_name != 'workflow_call' }}
env:
@@ -62,7 +68,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Determine run mode
id: run-mode
@@ -190,7 +196,7 @@ jobs:
- uses: actions/checkout@v4
with:
submodules: "recursive"
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v5
@@ -232,7 +238,7 @@ jobs:
- uses: actions/checkout@v4
with:
submodules: "recursive"
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v5
@@ -261,7 +267,7 @@ jobs:
steps:
- uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Cleanup
run: |
@@ -294,7 +300,7 @@ jobs:
steps:
- uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Cleanup
run: |
@@ -328,7 +334,7 @@ jobs:
steps:
- uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Cleanup
run: |
@@ -373,7 +379,7 @@ jobs:
steps:
- uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Cleanup
run: |
@@ -432,14 +438,16 @@ jobs:
jit-kernel-unit-test:
needs: [check-changes, call-gate]
if: needs.check-changes.outputs.jit_kernel == 'true'
if: |
!inputs.target_stage &&
needs.check-changes.outputs.jit_kernel == 'true'
runs-on: 1-gpu-runner
env:
RUNNER_LABELS: 1-gpu-runner
steps:
- uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Install dependencies
run: |
@@ -472,7 +480,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -521,7 +529,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -564,7 +572,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -607,7 +615,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -650,7 +658,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -693,7 +701,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -734,7 +742,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -783,7 +791,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -828,7 +836,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -880,7 +888,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -931,7 +939,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -979,7 +987,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -1030,7 +1038,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -1081,7 +1089,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -1139,7 +1147,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -1186,7 +1194,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -1257,7 +1265,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -1320,7 +1328,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -1377,7 +1385,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -1446,7 +1454,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -1488,7 +1496,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -1530,7 +1538,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -1577,7 +1585,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -1629,7 +1637,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
@@ -1678,7 +1686,7 @@ jobs:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || github.ref }}
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
+30 -7
View File
@@ -215,15 +215,38 @@ def handle_rerun_stage(
print(f"Error: {workflow_name} workflow not found")
return False
# Trigger workflow_dispatch on the PR's head branch
ref = pr.head.ref
print(f"Triggering {workflow_name} workflow on branch: {ref}")
# Check if PR is from a fork by comparing repo owners
# Handle case where fork repo may have been deleted (pr.head.repo is None)
is_fork = (
pr.head.repo is None or pr.head.repo.owner.login != gh_repo.owner.login
)
print(f"PR is from fork: {is_fork}")
# AMD workflow doesn't have version input, only target_stage
if is_amd_stage:
inputs = {"target_stage": stage_name}
if is_fork:
# For fork PRs: dispatch on main and pass SHA as input
# This is needed because fork branch names don't exist in the main repo
ref = "main"
pr_head_sha = pr.head.sha
print(
f"Triggering {workflow_name} workflow on ref: {ref}, PR head SHA: {pr_head_sha}"
)
if is_amd_stage:
inputs = {"target_stage": stage_name, "pr_head_sha": pr_head_sha}
else:
inputs = {
"version": "release",
"target_stage": stage_name,
"pr_head_sha": pr_head_sha,
}
else:
inputs = {"version": "release", "target_stage": stage_name}
# For non-fork PRs: dispatch on the PR branch directly
# This allows testing workflow changes before merge
ref = pr.head.ref
print(f"Triggering {workflow_name} workflow on branch: {ref}")
if is_amd_stage:
inputs = {"target_stage": stage_name}
else:
inputs = {"version": "release", "target_stage": stage_name}
# Use requests directly as PyGithub's create_dispatch only accepts HTTP 204
dispatch_url = f"https://api.github.com/repos/{gh_repo.full_name}/actions/workflows/{target_workflow.id}/dispatches"