Add /rerun-stage slash command to rerun specific PR test stages (#14262)

.github/workflows/pr-test.yml

@@ -15,6 +15,11 @@ on:
         options:
           - "release"
           - "nightly"
+      target_stage:
+        description: "Specific stage to run (optional, for quick testing)"
+        required: false
+        type: string
+        default: ""
 
 concurrency:
   group: pr-test-${{ github.ref }}
@@ -338,9 +343,12 @@ jobs:
   stage-a-test-1:
     needs: [check-changes, call-gate, sgl-kernel-build-wheels]
     if: |
+      (inputs.target_stage == 'stage-a-test-1') ||
+      (
         always() &&
         (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
         ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+      )
     runs-on: 1-gpu-runner
     env:
       RUNNER_LABELS: 1-gpu-runner
@@ -372,9 +380,12 @@ jobs:
   multimodal-gen-test-1-gpu:
     needs: [check-changes, call-gate, sgl-kernel-build-wheels]
     if: |
+      (inputs.target_stage == 'multimodal-gen-test-1-gpu') ||
+      (
         always() &&
         (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
         needs.check-changes.outputs.multimodal_gen == 'true'
+      )
     runs-on: 1-gpu-runner
     strategy:
       fail-fast: false
@@ -413,9 +424,12 @@ jobs:
   multimodal-gen-test-2-gpu:
     needs: [check-changes, call-gate, sgl-kernel-build-wheels]
     if: |
+      (inputs.target_stage == 'multimodal-gen-test-2-gpu') ||
+      (
         always() &&
         (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
         needs.check-changes.outputs.multimodal_gen == 'true'
+      )
     runs-on: 2-gpu-runner
     strategy:
       fail-fast: false
@@ -454,9 +468,12 @@ jobs:
   quantization-test:
       needs: [check-changes, call-gate, stage-a-test-1]
       if: |
+        (inputs.target_stage == 'quantization-test') ||
+        (
           always() &&
           (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
           ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+        )
       runs-on: 1-gpu-runner
       steps:
         - name: Checkout code
@@ -483,9 +500,12 @@ jobs:
   unit-test-backend-1-gpu:
     needs: [check-changes, call-gate, stage-a-test-1]
     if: |
+      (inputs.target_stage == 'unit-test-backend-1-gpu') ||
+      (
         always() &&
         (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
         ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+      )
     runs-on: 1-gpu-runner
     env:
       RUNNER_LABELS: 1-gpu-runner
@@ -519,9 +539,12 @@ jobs:
   unit-test-backend-2-gpu:
     needs: [check-changes, call-gate, unit-test-backend-1-gpu]
     if: |
+      (inputs.target_stage == 'unit-test-backend-2-gpu') ||
+      (
         always() &&
         (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
         ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+      )
     runs-on: 2-gpu-runner
     env:
       RUNNER_LABELS: 2-gpu-runner
@@ -554,9 +577,12 @@ jobs:
   unit-test-backend-4-gpu:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu]
     if: |
+      (inputs.target_stage == 'unit-test-backend-4-gpu') ||
+      (
         always() &&
         (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
         ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+      )
     runs-on: 4-gpu-h100
     env:
       RUNNER_LABELS: 4-gpu-h100
@@ -589,9 +615,12 @@ jobs:
   unit-test-backend-8-gpu-h200:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu]
     if: |
+      (inputs.target_stage == 'unit-test-backend-8-gpu-h200') ||
+      (
         always() &&
         (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
         ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+      )
     runs-on: 8-gpu-h200
     env:
       RUNNER_LABELS: 8-gpu-h200
@@ -624,9 +653,12 @@ jobs:
   unit-test-backend-8-gpu-h20:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu]
     if: |
+      (inputs.target_stage == 'unit-test-backend-8-gpu-h20') ||
+      (
         always() &&
         (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
         ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+      )
     runs-on: 8-gpu-h20
     env:
       SGLANG_CI_RDMA_ALL_DEVICES: "mlx5_1,mlx5_2,mlx5_3,mlx5_4"
@@ -660,9 +692,12 @@ jobs:
   performance-test-1-gpu-part-1:
     needs: [check-changes, call-gate, stage-a-test-1]
     if: |
+      (inputs.target_stage == 'performance-test-1-gpu-part-1') ||
+      (
         always() &&
         (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
         ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+      )
     runs-on: 1-gpu-runner
     env:
       RUNNER_LABELS: 1-gpu-runner
@@ -723,9 +758,12 @@ jobs:
   performance-test-1-gpu-part-2:
     needs: [check-changes, call-gate, stage-a-test-1]
     if: |
+      (inputs.target_stage == 'performance-test-1-gpu-part-2') ||
+      (
         always() &&
         (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
         ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+      )
     runs-on: 1-gpu-runner
     env:
       RUNNER_LABELS: 1-gpu-runner
@@ -778,9 +816,12 @@ jobs:
   performance-test-1-gpu-part-3:
     needs: [check-changes, call-gate, stage-a-test-1]
     if: |
+      (inputs.target_stage == 'performance-test-1-gpu-part-3') ||
+      (
         always() &&
         (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
         ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+      )
     runs-on: 1-gpu-runner
     env:
       RUNNER_LABELS: 1-gpu-runner
@@ -827,9 +868,12 @@ jobs:
   performance-test-2-gpu:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu]
     if: |
+      (inputs.target_stage == 'performance-test-2-gpu') ||
+      (
         always() &&
         (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
         ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+      )
     runs-on: 2-gpu-runner
     env:
       RUNNER_LABELS: 2-gpu-runner
@@ -888,9 +932,12 @@ jobs:
   accuracy-test-1-gpu:
     needs: [check-changes, call-gate, stage-a-test-1]
     if: |
+      (inputs.target_stage == 'accuracy-test-1-gpu') ||
+      (
         always() &&
         (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
         ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+      )
     runs-on: 1-gpu-runner
     env:
       RUNNER_LABELS: 1-gpu-runner
@@ -922,9 +969,12 @@ jobs:
   accuracy-test-2-gpu:
     needs: [check-changes, call-gate, accuracy-test-1-gpu]
     if: |
+      (inputs.target_stage == 'accuracy-test-2-gpu') ||
+      (
         always() &&
         (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
         ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+      )
     runs-on: 2-gpu-runner
     env:
       RUNNER_LABELS: 2-gpu-runner
@@ -956,9 +1006,12 @@ jobs:
   unit-test-deepep-4-gpu:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu]
     if: |
+      (inputs.target_stage == 'unit-test-deepep-4-gpu') ||
+      (
         always() &&
         (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
         ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+      )
     runs-on: 4-gpu-h100
     env:
       RUNNER_LABELS: 4-gpu-h100
@@ -987,9 +1040,12 @@ jobs:
   unit-test-deepep-8-gpu:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu]
     if: |
+      (inputs.target_stage == 'unit-test-deepep-8-gpu') ||
+      (
         always() &&
         (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
         ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+      )
     runs-on: 8-gpu-h200
     env:
       RUNNER_LABELS: 8-gpu-h200
@@ -1018,9 +1074,12 @@ jobs:
   unit-test-backend-4-gpu-b200:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu]
     if: |
+      (inputs.target_stage == 'unit-test-backend-4-gpu-b200') ||
+      (
         always() &&
         (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
         ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+      )
     runs-on: 4-gpu-b200
     env:
       RUNNER_LABELS: 4-gpu-b200
@@ -1054,9 +1113,12 @@ jobs:
   unit-test-backend-4-gpu-gb200:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu, sgl-kernel-build-wheels-arm]
     if: |
+      (inputs.target_stage == 'unit-test-backend-4-gpu-gb200') ||
+      (
         always() &&
         (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
         ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+      )
     runs-on: 4-gpu-gb200
     env:
       RUNNER_LABELS: 4-gpu-gb200

.github/workflows/slash-command-handler.yml

@@ -17,7 +17,8 @@ jobs:
       github.event.issue.pull_request &&
       (startsWith(github.event.comment.body, '/tag-run-ci-label') ||
        startsWith(github.event.comment.body, '/rerun-failed-ci') ||
-       startsWith(github.event.comment.body, '/tag-and-rerun-ci'))
+       startsWith(github.event.comment.body, '/tag-and-rerun-ci') ||
+       startsWith(github.event.comment.body, '/rerun-stage'))
     runs-on: ubuntu-latest
 
     steps:

scripts/ci/slash_command_handler.py

@@ -117,6 +117,107 @@ def handle_rerun_failed_ci(gh_repo, pr, comment, user_perms, react_on_success=Tr
         return False
 
 
+def handle_rerun_stage(
+    gh_repo, pr, comment, user_perms, stage_name, react_on_success=True
+):
+    """
+    Handles the /rerun-stage <stage-name> command.
+    Triggers a workflow_dispatch to run only the specified stage, skipping dependencies.
+    Returns True if action was taken, False otherwise.
+    """
+    if not user_perms.get("can_rerun_stage", False):
+        print("Permission denied: can_rerun_stage is false.")
+        return False
+
+    if not stage_name:
+        print("Error: No stage name provided")
+        comment.create_reaction("confused")
+        comment.create_comment(
+            f"❌ Please specify a stage name: `/rerun-stage <stage-name>`\n\n"
+            f"Examples: `/rerun-stage unit-test-backend-4-gpu`, `/rerun-stage accuracy-test-1-gpu`"
+        )
+        return False
+
+    print(f"Permission granted. Triggering workflow_dispatch for stage '{stage_name}'.")
+
+    # Valid stage names that support target_stage
+    valid_stages = [
+        "stage-a-test-1",
+        "multimodal-gen-test-1-gpu",
+        "multimodal-gen-test-2-gpu",
+        "quantization-test",
+        "unit-test-backend-1-gpu",
+        "unit-test-backend-2-gpu",
+        "unit-test-backend-4-gpu",
+        "unit-test-backend-8-gpu-h200",
+        "unit-test-backend-8-gpu-h20",
+        "performance-test-1-gpu-part-1",
+        "performance-test-1-gpu-part-2",
+        "performance-test-1-gpu-part-3",
+        "performance-test-2-gpu",
+        "accuracy-test-1-gpu",
+        "accuracy-test-2-gpu",
+        "unit-test-deepep-4-gpu",
+        "unit-test-deepep-8-gpu",
+        "unit-test-backend-4-gpu-b200",
+        "unit-test-backend-4-gpu-gb200",
+    ]
+
+    if stage_name not in valid_stages:
+        comment.create_reaction("confused")
+        comment.create_comment(
+            f"❌ Stage `{stage_name}` doesn't support isolated runs yet.\n\n"
+            f"Currently supported stages:\n"
+            + "\n".join(f"- `{s}`" for s in valid_stages)
+            + "\n\nOther stages will be added soon. For now, use `/rerun-failed-ci` for those stages."
+        )
+        return False
+
+    try:
+        # Get the PR Test workflow
+        workflows = gh_repo.get_workflows()
+        pr_test_workflow = None
+        for wf in workflows:
+            if wf.name == "PR Test":
+                pr_test_workflow = wf
+                break
+
+        if not pr_test_workflow:
+            print("Error: PR Test workflow not found")
+            return False
+
+        # Trigger workflow_dispatch on the PR's head branch
+        ref = pr.head.ref
+        print(f"Triggering workflow on branch: {ref}")
+
+        success = pr_test_workflow.create_dispatch(
+            ref=ref,
+            inputs={"version": "release", "target_stage": stage_name},
+        )
+
+        if success:
+            print(f"Successfully triggered workflow for stage '{stage_name}'")
+            if react_on_success:
+                comment.create_reaction("+1")
+                comment.create_comment(
+                    f"✅ Triggered `{stage_name}` to run independently (skipping dependencies).\n\n"
+                    f"Check the [Actions tab](https://github.com/{gh_repo.full_name}/actions) for progress."
+                )
+            return True
+        else:
+            print("Failed to trigger workflow_dispatch")
+            return False
+
+    except Exception as e:
+        print(f"Error triggering workflow_dispatch: {e}")
+        comment.create_reaction("confused")
+        comment.create_comment(
+            f"❌ Failed to trigger workflow: {str(e)}\n\n"
+            f"Please check the logs or contact maintainers."
+        )
+        return False
+
+
 def main():
     # 1. Load Environment Variables
     token = get_env_var("GITHUB_TOKEN")
@@ -168,6 +269,12 @@ def main():
         else:
             print("Combined command finished, but no actions were taken.")
 
+    elif first_line.startswith("/rerun-stage"):
+        # Extract stage name from command
+        parts = first_line.split(maxsplit=1)
+        stage_name = parts[1].strip() if len(parts) > 1 else None
+        handle_rerun_stage(repo, pr, comment, user_perms, stage_name)
+
     else:
         print(f"Unknown or ignored command: {first_line}")