Add AMD stage support to /rerun-stage command and fix related bugs (#14463)

.github/workflows/pr-test-amd.yml

@@ -20,6 +20,12 @@ on:
       - "sgl-kernel/**"
       - ".github/workflows/pr-test-amd.yml"
   workflow_dispatch:
+    inputs:
+      target_stage:
+        description: "Specific stage to run (optional, for quick testing)"
+        required: false
+        type: string
+        default: ""
 
 concurrency:
   group: pr-test-amd-${{ github.ref }}
@@ -54,7 +60,15 @@ jobs:
   # =============================================== sgl-kernel ====================================================
   sgl-kernel-unit-test-amd:
     needs: [check-changes]
-    if: needs.check-changes.outputs.sgl_kernel == 'true'
+    if: |
+      always() &&
+      (
+        (inputs.target_stage == 'sgl-kernel-unit-test-amd') ||
+        (
+          !inputs.target_stage &&
+          needs.check-changes.outputs.sgl_kernel == 'true'
+        )
+      )
     strategy:
       fail-fast: false
       matrix:
@@ -90,8 +104,16 @@ jobs:
 
   stage-a-test-1-amd:
     needs: [check-changes]
-    if: always() && !failure() && !cancelled() &&
+    if: |
+      always() &&
+      (
+        (inputs.target_stage == 'stage-a-test-1-amd') ||
+        (
+          !inputs.target_stage &&
+          (!failure() && !cancelled()) &&
           ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+        )
+      )
     strategy:
       fail-fast: false
       matrix:
@@ -120,8 +142,16 @@ jobs:
 
   unit-test-backend-1-gpu-amd:
     needs: [check-changes, stage-a-test-1-amd]
-    if: always() && !failure() && !cancelled() &&
+    if: |
+      always() &&
+      (
+        (inputs.target_stage == 'unit-test-backend-1-gpu-amd') ||
+        (
+          !inputs.target_stage &&
+          (!failure() && !cancelled()) &&
           ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+        )
+      )
     strategy:
       fail-fast: false
       matrix:
@@ -150,8 +180,16 @@ jobs:
 
   unit-test-backend-2-gpu-amd:
     needs: [check-changes, stage-a-test-1-amd]
-    if: always() && !failure() && !cancelled() &&
+    if: |
+      always() &&
+      (
+        (inputs.target_stage == 'unit-test-backend-2-gpu-amd') ||
+        (
+          !inputs.target_stage &&
+          (!failure() && !cancelled()) &&
           ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+        )
+      )
     strategy:
       fail-fast: false
       matrix:
@@ -181,8 +219,17 @@ jobs:
   unit-test-backend-8-gpu-amd:
     needs: [check-changes, unit-test-backend-2-gpu-amd]
     # Temporarily disabled - uncomment when ready to re-enable
-    if: false && always() && !failure() && !cancelled() &&
+    if: |
+      always() &&
+      (
+        (inputs.target_stage == 'unit-test-backend-8-gpu-amd') ||
+        (
+          false &&
+          !inputs.target_stage &&
+          (!failure() && !cancelled()) &&
           ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+        )
+      )
     env:
       RUNNER_LABELS: linux-mi300-gpu-8
     strategy:
@@ -219,8 +266,16 @@ jobs:
 
   performance-test-1-gpu-part-1-amd:
     needs: [check-changes, stage-a-test-1-amd]
-    if: always() && !failure() && !cancelled() &&
+    if: |
+      always() &&
+      (
+        (inputs.target_stage == 'performance-test-1-gpu-part-1-amd') ||
+        (
+          !inputs.target_stage &&
+          (!failure() && !cancelled()) &&
           ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+        )
+      )
     strategy:
       fail-fast: false
       matrix:
@@ -264,8 +319,16 @@ jobs:
 
   performance-test-1-gpu-part-2-amd:
     needs: [check-changes, stage-a-test-1-amd]
-    if: always() && !failure() && !cancelled() &&
+    if: |
+      always() &&
+      (
+        (inputs.target_stage == 'performance-test-1-gpu-part-2-amd') ||
+        (
+          !inputs.target_stage &&
+          (!failure() && !cancelled()) &&
           ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+        )
+      )
     strategy:
       fail-fast: false
       matrix:
@@ -303,8 +366,16 @@ jobs:
 
   performance-test-2-gpu-amd:
     needs: [check-changes, unit-test-backend-2-gpu-amd]
-    if: always() && !failure() && !cancelled() &&
+    if: |
+      always() &&
+      (
+        (inputs.target_stage == 'performance-test-2-gpu-amd') ||
+        (
+          !inputs.target_stage &&
+          (!failure() && !cancelled()) &&
           ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+        )
+      )
     strategy:
       fail-fast: false
       matrix:
@@ -352,8 +423,16 @@ jobs:
 
   accuracy-test-1-gpu-amd:
     needs: [check-changes, stage-a-test-1-amd]
-    if: always() && !failure() && !cancelled() &&
+    if: |
+      always() &&
+      (
+        (inputs.target_stage == 'accuracy-test-1-gpu-amd') ||
+        (
+          !inputs.target_stage &&
+          (!failure() && !cancelled()) &&
           ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+        )
+      )
     strategy:
       fail-fast: false
       matrix:
@@ -383,8 +462,16 @@ jobs:
 
   accuracy-test-2-gpu-amd:
     needs: [check-changes, accuracy-test-1-gpu-amd]
-    if: always() && !failure() && !cancelled() &&
+    if: |
+      always() &&
+      (
+        (inputs.target_stage == 'accuracy-test-2-gpu-amd') ||
+        (
+          !inputs.target_stage &&
+          (!failure() && !cancelled()) &&
           ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+        )
+      )
     strategy:
       fail-fast: false
       matrix:

.github/workflows/pr-test.yml

@@ -343,12 +343,15 @@ jobs:
   stage-a-test-1:
     needs: [check-changes, call-gate, sgl-kernel-build-wheels]
     if: |
+      always() &&
+      (
         (inputs.target_stage == 'stage-a-test-1') ||
         (
-        always() &&
+          !inputs.target_stage &&
           (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
@@ -380,12 +383,15 @@ jobs:
   multimodal-gen-test-1-gpu:
     needs: [check-changes, call-gate, sgl-kernel-build-wheels]
     if: |
+      always() &&
+      (
         (inputs.target_stage == 'multimodal-gen-test-1-gpu') ||
         (
-        always() &&
+          !inputs.target_stage &&
           (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
           needs.check-changes.outputs.multimodal_gen == 'true'
         )
+      )
     runs-on: 1-gpu-runner
     strategy:
       fail-fast: false
@@ -424,12 +430,15 @@ jobs:
   multimodal-gen-test-2-gpu:
     needs: [check-changes, call-gate, sgl-kernel-build-wheels]
     if: |
+      always() &&
+      (
         (inputs.target_stage == 'multimodal-gen-test-2-gpu') ||
         (
-        always() &&
+          !inputs.target_stage &&
           (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
           needs.check-changes.outputs.multimodal_gen == 'true'
         )
+      )
     runs-on: 2-gpu-runner
     strategy:
       fail-fast: false
@@ -468,12 +477,15 @@ jobs:
   quantization-test:
       needs: [check-changes, call-gate, stage-a-test-1]
       if: |
+        always() &&
+        (
           (inputs.target_stage == 'quantization-test') ||
           (
-          always() &&
+            !inputs.target_stage &&
             (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
@@ -500,12 +512,15 @@ jobs:
   unit-test-backend-1-gpu:
     needs: [check-changes, call-gate, stage-a-test-1]
     if: |
+      always() &&
+      (
         (inputs.target_stage == 'unit-test-backend-1-gpu') ||
         (
-        always() &&
+          !inputs.target_stage &&
           (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
@@ -539,12 +554,15 @@ jobs:
   unit-test-backend-2-gpu:
     needs: [check-changes, call-gate, unit-test-backend-1-gpu]
     if: |
+      always() &&
+      (
         (inputs.target_stage == 'unit-test-backend-2-gpu') ||
         (
-        always() &&
+          !inputs.target_stage &&
           (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
@@ -577,12 +595,15 @@ jobs:
   unit-test-backend-4-gpu:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu]
     if: |
+      always() &&
+      (
         (inputs.target_stage == 'unit-test-backend-4-gpu') ||
         (
-        always() &&
+          !inputs.target_stage &&
           (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
@@ -615,12 +636,15 @@ jobs:
   unit-test-backend-8-gpu-h200:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu]
     if: |
+      always() &&
+      (
         (inputs.target_stage == 'unit-test-backend-8-gpu-h200') ||
         (
-        always() &&
+          !inputs.target_stage &&
           (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
@@ -653,12 +677,15 @@ jobs:
   unit-test-backend-8-gpu-h20:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu]
     if: |
+      always() &&
+      (
         (inputs.target_stage == 'unit-test-backend-8-gpu-h20') ||
         (
-        always() &&
+          !inputs.target_stage &&
           (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"
@@ -692,12 +719,15 @@ jobs:
   performance-test-1-gpu-part-1:
     needs: [check-changes, call-gate, stage-a-test-1]
     if: |
+      always() &&
+      (
         (inputs.target_stage == 'performance-test-1-gpu-part-1') ||
         (
-        always() &&
+          !inputs.target_stage &&
           (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
@@ -758,12 +788,15 @@ jobs:
   performance-test-1-gpu-part-2:
     needs: [check-changes, call-gate, stage-a-test-1]
     if: |
+      always() &&
+      (
         (inputs.target_stage == 'performance-test-1-gpu-part-2') ||
         (
-        always() &&
+          !inputs.target_stage &&
           (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
@@ -816,12 +849,15 @@ jobs:
   performance-test-1-gpu-part-3:
     needs: [check-changes, call-gate, stage-a-test-1]
     if: |
+      always() &&
+      (
         (inputs.target_stage == 'performance-test-1-gpu-part-3') ||
         (
-        always() &&
+          !inputs.target_stage &&
           (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
@@ -868,12 +904,15 @@ jobs:
   performance-test-2-gpu:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu]
     if: |
+      always() &&
+      (
         (inputs.target_stage == 'performance-test-2-gpu') ||
         (
-        always() &&
+          !inputs.target_stage &&
           (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
@@ -932,12 +971,15 @@ jobs:
   accuracy-test-1-gpu:
     needs: [check-changes, call-gate, stage-a-test-1]
     if: |
+      always() &&
+      (
         (inputs.target_stage == 'accuracy-test-1-gpu') ||
         (
-        always() &&
+          !inputs.target_stage &&
           (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
@@ -969,12 +1011,15 @@ jobs:
   accuracy-test-2-gpu:
     needs: [check-changes, call-gate, accuracy-test-1-gpu]
     if: |
+      always() &&
+      (
         (inputs.target_stage == 'accuracy-test-2-gpu') ||
         (
-        always() &&
+          !inputs.target_stage &&
           (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
@@ -1006,12 +1051,15 @@ jobs:
   unit-test-deepep-4-gpu:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu]
     if: |
+      always() &&
+      (
         (inputs.target_stage == 'unit-test-deepep-4-gpu') ||
         (
-        always() &&
+          !inputs.target_stage &&
           (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
@@ -1040,12 +1088,15 @@ jobs:
   unit-test-deepep-8-gpu:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu]
     if: |
+      always() &&
+      (
         (inputs.target_stage == 'unit-test-deepep-8-gpu') ||
         (
-        always() &&
+          !inputs.target_stage &&
           (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
@@ -1074,12 +1125,15 @@ jobs:
   unit-test-backend-4-gpu-b200:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu]
     if: |
+      always() &&
+      (
         (inputs.target_stage == 'unit-test-backend-4-gpu-b200') ||
         (
-        always() &&
+          !inputs.target_stage &&
           (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
@@ -1113,12 +1167,15 @@ jobs:
   unit-test-backend-4-gpu-gb200:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu, sgl-kernel-build-wheels-arm]
     if: |
+      always() &&
+      (
         (inputs.target_stage == 'unit-test-backend-4-gpu-gb200') ||
         (
-        always() &&
+          !inputs.target_stage &&
           (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

scripts/ci/slash_command_handler.py

@@ -132,7 +132,7 @@ def handle_rerun_stage(
     if not stage_name:
         print("Error: No stage name provided")
         comment.create_reaction("confused")
-        comment.create_comment(
+        pr.create_issue_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`"
         )
@@ -140,8 +140,8 @@ def handle_rerun_stage(
 
     print(f"Permission granted. Triggering workflow_dispatch for stage '{stage_name}'.")
 
-    # Valid stage names that support target_stage
+    # Valid NVIDIA stage names that support target_stage
-    valid_stages = [
+    nvidia_stages = [
         "stage-a-test-1",
         "multimodal-gen-test-1-gpu",
         "multimodal-gen-test-2-gpu",
@@ -163,43 +163,69 @@ def handle_rerun_stage(
         "unit-test-backend-4-gpu-gb200",
     ]
 
+    # Valid AMD stage names that support target_stage
+    amd_stages = [
+        "sgl-kernel-unit-test-amd",
+        "stage-a-test-1-amd",
+        "unit-test-backend-1-gpu-amd",
+        "unit-test-backend-2-gpu-amd",
+        "unit-test-backend-8-gpu-amd",
+        "performance-test-1-gpu-part-1-amd",
+        "performance-test-1-gpu-part-2-amd",
+        "performance-test-2-gpu-amd",
+        "accuracy-test-1-gpu-amd",
+        "accuracy-test-2-gpu-amd",
+    ]
+
+    valid_stages = nvidia_stages + amd_stages
+    is_amd_stage = stage_name in amd_stages
+
     if stage_name not in valid_stages:
         comment.create_reaction("confused")
-        comment.create_comment(
+        pr.create_issue_comment(
             f"❌ Stage `{stage_name}` doesn't support isolated runs yet.\n\n"
-            f"Currently supported stages:\n"
+            f"**NVIDIA stages:**\n"
-            + "\n".join(f"- `{s}`" for s in valid_stages)
+            + "\n".join(f"- `{s}`" for s in nvidia_stages)
+            + "\n\n**AMD stages:**\n"
+            + "\n".join(f"- `{s}`" for s in amd_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
+        # Get the appropriate workflow based on stage type
+        workflow_name = "PR Test (AMD)" if is_amd_stage else "PR Test"
         workflows = gh_repo.get_workflows()
-        pr_test_workflow = None
+        target_workflow = None
         for wf in workflows:
-            if wf.name == "PR Test":
+            if wf.name == workflow_name:
-                pr_test_workflow = wf
+                target_workflow = wf
                 break
 
-        if not pr_test_workflow:
+        if not target_workflow:
-            print("Error: PR Test workflow not found")
+            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 on branch: {ref}")
+        print(f"Triggering {workflow_name} workflow on branch: {ref}")
 
-        success = pr_test_workflow.create_dispatch(
+        # AMD workflow doesn't have version input, only target_stage
+        if is_amd_stage:
+            inputs = {"target_stage": stage_name}
+        else:
+            inputs = {"version": "release", "target_stage": stage_name}
+
+        success = target_workflow.create_dispatch(
             ref=ref,
-            inputs={"version": "release", "target_stage": stage_name},
+            inputs=inputs,
         )
 
         if success:
             print(f"Successfully triggered workflow for stage '{stage_name}'")
             if react_on_success:
                 comment.create_reaction("+1")
-                comment.create_comment(
+                pr.create_issue_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."
                 )
@@ -211,7 +237,7 @@ def handle_rerun_stage(
     except Exception as e:
         print(f"Error triggering workflow_dispatch: {e}")
         comment.create_reaction("confused")
-        comment.create_comment(
+        pr.create_issue_comment(
             f"❌ Failed to trigger workflow: {str(e)}\n\n"
             f"Please check the logs or contact maintainers."
         )