[CI] Support CPU stage and auto-batch same-stage files in /rerun-test (#22081)

.github/workflows/rerun-test.yml

@@ -5,7 +5,7 @@ on:
   workflow_dispatch:
     inputs:
       test_command:
-        description: "Test command to run (e.g. 'registered/core/test_srt_endpoint.py TestSRTEndpoint.test_simple_decode')"
+        description: "Test command(s) to run, one per line (e.g. 'registered/core/test_srt_endpoint.py TestSRTEndpoint.test_simple_decode')"
         required: true
         type: string
       runner_label:
@@ -22,6 +22,7 @@ on:
           - 8-gpu-h200
           - 8-gpu-h20
           - 8-gpu-b200
+          - ubuntu-latest
       pr_head_sha:
         description: "PR head SHA to checkout (for /rerun-test on fork PRs)"
         required: false
@@ -32,6 +33,11 @@ on:
         required: false
         type: string
         default: "false"
+      is_cpu:
+        description: "Run as CPU-only test (uses ubuntu-latest with uv pip install)"
+        required: false
+        type: string
+        default: "false"
 
 env:
   SGLANG_IS_IN_CI: true
@@ -45,6 +51,7 @@ permissions:
 
 jobs:
   rerun-test-cuda:
+    if: inputs.is_cpu != 'true'
     runs-on: ${{ inputs.runner_label }}
     timeout-minutes: 120
     env:
@@ -77,7 +84,53 @@ jobs:
             source /etc/profile.d/sglang-ci.sh
           fi
           cd test/
-          python3 ${{ inputs.test_command }}
+          echo "${{ inputs.test_command }}" | while IFS= read -r cmd; do
+            [ -z "$cmd" ] && continue
+            echo ">>> Running: python3 $cmd"
+            python3 $cmd || exit 1
+          done
 
       - uses: ./.github/actions/upload-cuda-coredumps
         if: always()
+
+  rerun-test-cpu:
+    if: inputs.is_cpu == 'true'
+    runs-on: ubuntu-latest
+    timeout-minutes: 120
+    steps:
+      - name: Free disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc
+          df -h
+
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.pr_head_sha || github.sha }}
+
+      - uses: ./.github/actions/check-maintenance
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.10'
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+
+      - name: Install dependencies
+        timeout-minutes: 20
+        env:
+          UV_SYSTEM_PYTHON: "1"
+        run: |
+          uv pip install -e "python[dev]" --index-strategy unsafe-best-match --prerelease allow
+
+      - name: Run test
+        timeout-minutes: 60
+        run: |
+          cd test/
+          echo "${{ inputs.test_command }}" | while IFS= read -r cmd; do
+            [ -z "$cmd" ] && continue
+            echo ">>> Running: python3 $cmd"
+            python3 $cmd || exit 1
+          done

scripts/ci/utils/slash_command_handler.py

@@ -475,32 +475,23 @@ def resolve_test_file(file_part):
     return file_part, None
 
 
-def detect_cuda_suite(file_path_from_test):
+def detect_suite(file_path_from_test):
     """
-    Read a test file and extract the suite from register_cuda_ci(suite="...").
+    Read a test file and extract the suite from register_cuda_ci or register_cpu_ci.
 
-    Returns (suite_name, runner_label, use_deepep, error_message).
+    Returns (suite_name, runner_label, use_deepep, is_cpu, error_message).
     """
     full_path = f"test/{file_path_from_test}"
     with open(full_path, "r") as f:
         content = f.read()
 
+    # Try CUDA first
     match = re.search(
         r'^[^#\n]*register_cuda_ci\([^)]*suite\s*=\s*["\']([^"\']+)["\']',
         content,
         re.MULTILINE,
     )
-    if not match:
+    if match:
-        return (
-            None,
-            None,
-            False,
-            (
-                f"No `register_cuda_ci()` found in `{full_path}`.\n\n"
-                f"This file may not be a registered CUDA CI test."
-            ),
-        )
-
         suite = match.group(1)
         runner = CUDA_SUITE_TO_RUNNER.get(suite)
         if not runner:
@@ -509,22 +500,44 @@ def detect_cuda_suite(file_path_from_test):
                 suite,
                 None,
                 False,
+                False,
                 (
                     f"Unknown CUDA suite `{suite}` in `{full_path}`.\n\n"
                     f"Known suites: {known}"
                 ),
             )
         use_deepep = suite in DEEPEP_SUITES
-    return suite, runner, use_deepep, None
+        return suite, runner, use_deepep, False, None
+
+    # Try CPU
+    match = re.search(
+        r'^[^#\n]*register_cpu_ci\([^)]*suite\s*=\s*["\']([^"\']+)["\']',
+        content,
+        re.MULTILINE,
+    )
+    if match:
+        suite = match.group(1)
+        return suite, "ubuntu-latest", False, True, None
+
+    return (
+        None,
+        None,
+        False,
+        False,
+        (
+            f"No `register_cuda_ci()` or `register_cpu_ci()` found in `{full_path}`.\n\n"
+            f"This file may not be a registered CI test."
+        ),
+    )
 
 
-def _resolve_and_dispatch_ut(gh_repo, pr, test_spec, token):
+def _resolve_test_spec(test_spec):
     """
-    Resolve a single test spec and dispatch a workflow run.
+    Resolve a single test spec into its components without dispatching.
 
-    Returns a dict with keys: spec, success, test_command, runner_label, run_url, error.
+    Returns a dict with keys: spec, resolved_path, test_command, suite,
+    runner_label, use_deepep, is_cpu, error.
     """
-    # Parse spec: split on :: to get file path and optional test selector
     if "::" in test_spec:
         file_part, test_selector = test_spec.split("::", 1)
     else:
@@ -535,25 +548,48 @@ def _resolve_and_dispatch_ut(gh_repo, pr, test_spec, token):
     if test_selector:
         test_selector = test_selector.strip()
 
-    # Resolve file path
     resolved_path, err = resolve_test_file(file_part)
     if err:
-        return {"spec": test_spec, "success": False, "error": err}
+        return {"spec": test_spec, "error": err}
 
-    # Detect suite and runner
+    suite, runner_label, use_deepep, is_cpu, err = detect_suite(resolved_path)
-    suite, runner_label, use_deepep, err = detect_cuda_suite(resolved_path)
     if err:
-        return {"spec": test_spec, "success": False, "error": err}
+        return {"spec": test_spec, "error": err}
 
-    # Build test_command: file path (+ optional test selector as unittest arg)
     test_command = resolved_path
     if test_selector:
         test_command = f"{resolved_path} {test_selector}"
 
     print(
         f"Resolved: file={resolved_path}, selector={test_selector}, "
-        f"suite={suite}, runner={runner_label}, deepep={use_deepep}, command='{test_command}'"
+        f"suite={suite}, runner={runner_label}, deepep={use_deepep}, "
+        f"cpu={is_cpu}, command='{test_command}'"
     )
+    return {
+        "spec": test_spec,
+        "test_command": test_command,
+        "suite": suite,
+        "runner_label": runner_label,
+        "use_deepep": use_deepep,
+        "is_cpu": is_cpu,
+        "error": None,
+    }
+
+
+def _dispatch_batch(gh_repo, pr, batch, token):
+    """
+    Dispatch a single workflow run for a batch of resolved test specs
+    that share the same (runner_label, use_deepep, is_cpu).
+
+    Returns a dict with keys: specs, success, test_commands, runner_label, run_url, error.
+    """
+    test_commands = [r["test_command"] for r in batch]
+    runner_label = batch[0]["runner_label"]
+    use_deepep = batch[0]["use_deepep"]
+    is_cpu = batch[0]["is_cpu"]
+
+    # Join multiple commands with newlines for the workflow to iterate over
+    combined_command = "\n".join(test_commands)
 
     try:
         workflow_name = "Rerun Test"
@@ -566,7 +602,7 @@ def _resolve_and_dispatch_ut(gh_repo, pr, test_spec, token):
 
         if not target_workflow:
             return {
-                "spec": test_spec,
+                "specs": [r["spec"] for r in batch],
                 "success": False,
                 "error": f"{workflow_name} workflow not found",
             }
@@ -576,22 +612,18 @@ def _resolve_and_dispatch_ut(gh_repo, pr, test_spec, token):
         )
 
         pr_head_sha = None
+        inputs = {
+            "test_command": combined_command,
+            "runner_label": runner_label,
+            "use_deepep": str(use_deepep).lower(),
+            "is_cpu": str(is_cpu).lower(),
+        }
         if is_fork:
             ref = "main"
             pr_head_sha = pr.head.sha
-            inputs = {
+            inputs["pr_head_sha"] = pr_head_sha
-                "test_command": test_command,
-                "runner_label": runner_label,
-                "pr_head_sha": pr_head_sha,
-                "use_deepep": str(use_deepep).lower(),
-            }
         else:
             ref = pr.head.ref
-            inputs = {
-                "test_command": test_command,
-                "runner_label": runner_label,
-                "use_deepep": str(use_deepep).lower(),
-            }
 
         dispatch_time = time.time()
 
@@ -608,12 +640,12 @@ def _resolve_and_dispatch_ut(gh_repo, pr, test_spec, token):
         if not success:
             print(f"Dispatch failed: {dispatch_resp.status_code} {dispatch_resp.text}")
             return {
-                "spec": test_spec,
+                "specs": [r["spec"] for r in batch],
                 "success": False,
                 "error": f"Dispatch failed: {dispatch_resp.status_code}",
             }
 
-        print(f"Successfully triggered rerun-test: {test_command}")
+        print(f"Successfully triggered rerun-test: {combined_command}")
 
         run_url = find_workflow_run_url(
             gh_repo,
@@ -624,25 +656,29 @@ def _resolve_and_dispatch_ut(gh_repo, pr, test_spec, token):
             dispatch_time,
             pr_head_sha=pr_head_sha,
             max_wait=30,
-            test_command=test_command,
+            test_command=combined_command,
         )
         return {
-            "spec": test_spec,
+            "specs": [r["spec"] for r in batch],
             "success": True,
-            "test_command": test_command,
+            "test_commands": test_commands,
             "runner_label": runner_label,
             "run_url": run_url,
         }
 
     except Exception as e:
-        print(f"Error triggering rerun-test for {test_spec}: {e}")
+        print(f"Error triggering rerun-test for batch: {e}")
-        return {"spec": test_spec, "success": False, "error": str(e)}
+        return {
+            "specs": [r["spec"] for r in batch],
+            "success": False,
+            "error": str(e),
+        }
 
 
 def handle_rerun_test(gh_repo, pr, comment, user_perms, test_specs, token):
     """
-    Handles the /rerun-test command. Accepts a list of test specs and dispatches
+    Handles the /rerun-test command. Resolves all test specs, groups them by
-    a workflow run for each, posting a single consolidated comment.
+    (runner_label, use_deepep, is_cpu), and dispatches one workflow per group.
     """
     # SECURITY: For fork PRs, only allow /rerun-test if the commenter has write+ permission.
     # This command checks out and executes code from the PR branch on self-hosted GPU
@@ -681,36 +717,60 @@ def handle_rerun_test(gh_repo, pr, comment, user_perms, test_specs, token):
         )
         return False
 
-    results = []
+    # Phase 1: Resolve all specs
+    resolved = []
+    resolve_failures = []
     for spec in test_specs:
-        results.append(_resolve_and_dispatch_ut(gh_repo, pr, spec, token))
+        r = _resolve_test_spec(spec)
+        if r.get("error"):
+            resolve_failures.append(r)
+        else:
+            resolved.append(r)
+
+    # Phase 2: Group by (runner_label, use_deepep, is_cpu)
+    groups = {}
+    for r in resolved:
+        key = (r["runner_label"], r["use_deepep"], r["is_cpu"])
+        groups.setdefault(key, []).append(r)
+
+    # Phase 3: Dispatch one workflow per group
+    dispatch_results = []
+    for batch in groups.values():
+        dispatch_results.append(_dispatch_batch(gh_repo, pr, batch, token))
 
     # Build consolidated comment
-    successes = [r for r in results if r["success"]]
-    failures = [r for r in results if not r["success"]]
-
     lines = []
-    for r in successes:
+    for dr in dispatch_results:
-        if r.get("run_url"):
+        if dr["success"]:
+            cmds = "\n".join(
+                f"cd test/ && python3 {cmd}" for cmd in dr["test_commands"]
+            )
+            if dr.get("run_url"):
                 lines.append(
-                f"✅ `{r['runner_label']}`: [View workflow run]({r['run_url']})\n"
+                    f"✅ `{dr['runner_label']}` ({len(dr['test_commands'])} test{'s' if len(dr['test_commands']) > 1 else ''}): "
-                f"```\ncd test/ && python3 {r['test_command']}\n```"
+                    f"[View workflow run]({dr['run_url']})\n"
+                    f"```\n{cmds}\n```"
                 )
             else:
                 lines.append(
-                f"✅ `{r['runner_label']}`:\n"
+                    f"✅ `{dr['runner_label']}` ({len(dr['test_commands'])} test{'s' if len(dr['test_commands']) > 1 else ''}):\n"
-                f"```\ncd test/ && python3 {r['test_command']}\n```\n"
+                    f"```\n{cmds}\n```\n"
                     f"⚠️ Could not retrieve workflow run URL. "
                     f"Check the [Actions tab](https://github.com/{gh_repo.full_name}/actions) for progress."
                 )
-    for r in failures:
+        else:
+            specs_str = ", ".join(f"`{s}`" for s in dr["specs"])
+            lines.append(f"❌ {specs_str}: {dr['error']}")
+
+    for r in resolve_failures:
         lines.append(f"❌ `{r['spec']}`: {r['error']}")
 
     body = "\n\n".join(lines)
 
+    successes = [dr for dr in dispatch_results if dr["success"]]
     if successes:
         comment.create_reaction("+1")
-    if failures and not successes:
+    if not successes and (resolve_failures or dispatch_results):
         comment.create_reaction("confused")
 
     pr.create_issue_comment(body)