import glob import json import os import re import sys import time from datetime import datetime, timezone import requests from github import Auth, Github # Import scripts/ci/runner_configs.py (sibling-up dir) for runner_config -> runs_on lookup. sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")) import runner_configs as _runner_configs # noqa: E402 # rerun-test workflow doesn't build sgl-kernel, so b200 stages always use the # non-kernel pool when resolving the `$b200_runner` sentinel from runner_configs.yml. _B200_DEFAULT_RUNNER = "4-gpu-b200" # install_script values from runner_configs.yml are passed verbatim into a # `bash ${{ inputs.install_script }}` step in rerun-test.yml. GHA expression # substitution happens before bash parses, so shell metacharacters in the # string would inject. Restrict the allowed shape to `scripts/ci/cuda/*.sh` # (single path component under that dir, no whitespace/operators). _ALLOWED_INSTALL_SCRIPT = re.compile(r"^scripts/ci/cuda/[\w.-]+\.sh$") # Configuration PERMISSIONS_FILE_PATH = ".github/CI_PERMISSIONS.json" MAINTENANCE_ISSUE_NUMBER = 21065 def _check_rebase_gate(gh_repo, pr, token): """ Pre-dispatch gate mirroring `.github/actions/check-maintenance/action.yml`. Without this, /rerun-test would dispatch a workflow_run on a PR that's behind a required base, the action would catch it, and every job in the run would fail at the gate — wasting runner time and producing N error annotations instead of one comment. Pre-checking here short-circuits the dispatch and posts a single explanatory comment. Mirrors the action's two independent modes driven by issue #21065: (1) Full-pause: maintenance issue is OPEN (2) Rebase-required: issue body contains `MIN_BASE_SHA: ` Both bypassed by the `bypass-maintenance` PR label. Returns (allowed: bool, message: Optional[str]). When allowed=False, caller MUST post `message` to the PR and skip dispatch. Fail-open on API errors (matches the action's behavior). """ headers = { "Authorization": f"Bearer {token}", "Accept": "application/vnd.github+json", "X-GitHub-Api-Version": "2022-11-28", } repo_full_name = gh_repo.full_name try: issue_resp = requests.get( f"https://api.github.com/repos/{repo_full_name}/issues/{MAINTENANCE_ISSUE_NUMBER}", headers=headers, timeout=15, ) if issue_resp.status_code != 200: print( f"check_rebase_gate: issue fetch returned {issue_resp.status_code}; fail-open" ) return True, None issue_data = issue_resp.json() except Exception as e: print(f"check_rebase_gate: issue fetch failed ({e}); fail-open") return True, None issue_state = (issue_data.get("state") or "").lower() issue_body = issue_data.get("body") or "" min_base_sha = None # First MIN_BASE_SHA: line wins. Match the action's parser: # tolerate optional backticks and either ':' or '=' separator. for line in issue_body.replace("\r", "").split("\n"): m = re.match( r"^\s*`?MIN_BASE_SHA`?\s*[:=]\s*`?([A-Fa-f0-9]+)`?", line, ) if m: candidate = m.group(1) if 7 <= len(candidate) <= 40: min_base_sha = candidate break gate_active = (issue_state == "open") or bool(min_base_sha) if not gate_active: return True, None bypass = any( (lbl.name if hasattr(lbl, "name") else lbl.get("name")) == "bypass-maintenance" for lbl in pr.get_labels() ) if bypass: print("check_rebase_gate: PR has bypass-maintenance label; allowing dispatch") return True, None if issue_state == "open": msg = ( "## ⚠️ CI Maintenance Mode is Active\n" "The CI infrastructure is currently under maintenance. " "All PR CI runs are paused until maintenance is complete. " "**Merging non-CI-fix PRs is prohibited during maintenance mode.**\n\n" f"Follow [issue #{MAINTENANCE_ISSUE_NUMBER}]" f"(https://github.com/{repo_full_name}/issues/{MAINTENANCE_ISSUE_NUMBER}) " "for status updates. Re-run was not dispatched." ) return False, msg # MIN_BASE_SHA set, issue not OPEN — check rebase status. pr_head_sha = pr.head.sha try: compare_resp = requests.get( f"https://api.github.com/repos/{repo_full_name}/compare/{min_base_sha}...{pr_head_sha}", headers=headers, timeout=15, ) if compare_resp.status_code != 200: print( f"check_rebase_gate: compare API returned {compare_resp.status_code}; fail-open" ) return True, None status = compare_resp.json().get("status", "unknown") except Exception as e: print(f"check_rebase_gate: compare API failed ({e}); fail-open") return True, None if status in ("ahead", "identical"): return True, None msg = ( "## ⚠️ Rebase Required Before Re-run\n" f"A major update has landed on `main`. Your PR is `{status}` relative " f"to required base commit `{min_base_sha[:12]}`.\n\n" "**Re-run was not dispatched.** What to do:\n" "- Rebase your branch onto the latest `main` and push again\n" f"- Follow [issue #{MAINTENANCE_ISSUE_NUMBER}]" f"(https://github.com/{repo_full_name}/issues/{MAINTENANCE_ISSUE_NUMBER}) for context\n" "- CI-fix PRs may request the `bypass-maintenance` label to skip this check" ) return False, msg def find_workflow_run_url( gh_repo, workflow_id, ref, target_stage, token, dispatch_time, pr_head_sha=None, max_wait=30, test_command=None, ): """ Poll for the workflow run URL after dispatch. Uses the dynamic run-name feature to identify runs: - Fork PRs: display_title = "[stage-name] sha" - Non-fork PRs: display_title = "[stage-name]" Args: gh_repo: PyGithub repository object workflow_id: ID of the workflow that was dispatched ref: Branch/ref the workflow was dispatched on target_stage: The stage name we're looking for token: GitHub API token dispatch_time: Unix timestamp when dispatch was triggered pr_head_sha: PR head SHA (for fork PRs, used to match display_title) max_wait: Maximum seconds to wait for the run to appear Returns: The workflow run URL if found, None otherwise. """ # Build expected display_title based on workflow's run-name. # rerun-test includes test_command: "[rerun-test] []" # Other workflows: "[stage-name] []" suffix = f" {test_command}" if test_command else "" if pr_head_sha: expected_title = f"[{target_stage}]{suffix} {pr_head_sha}" else: expected_title = f"[{target_stage}]{suffix}" print(f"Looking for workflow run with display_title: {expected_title}") for attempt in range(max_wait // 5): time.sleep(5) # Get recent workflow_dispatch runs for this workflow runs_url = f"https://api.github.com/repos/{gh_repo.full_name}/actions/workflows/{workflow_id}/runs" runs_resp = requests.get( runs_url, params={"event": "workflow_dispatch", "branch": ref, "per_page": 10}, headers={ "Authorization": f"Bearer {token}", "Accept": "application/vnd.github+json", }, ) if runs_resp.status_code != 200: print(f"Failed to fetch workflow runs: {runs_resp.status_code}") continue for run in runs_resp.json().get("workflow_runs", []): # Skip runs created before our dispatch (with 10s tolerance) run_created = datetime.fromisoformat( run["created_at"].replace("Z", "+00:00") ).timestamp() if run_created < dispatch_time - 10: continue # Match by display_title (set by workflow's run-name directive) # This is immediately available, unlike job names which require waiting. # # GitHub caps display_title at 512 chars and appends "..." when the # workflow's run-name is longer (e.g. /rerun-test with many files # newline-joins them into one long command). Accept the un-truncated # prefix when the rest matches our submission verbatim — concurrent # dispatches use different commands (different files per group), so # the prefix is enough to disambiguate. display_title = run.get("display_title", "") if display_title == expected_title or ( display_title.endswith("...") and expected_title.startswith(display_title[:-3]) ): print( f"Found matching workflow run: {run['id']} with title '{display_title}'" ) return run["html_url"] print(f"Could not find workflow run after {max_wait} seconds") return None def get_env_var(name): val = os.getenv(name) if not val: print(f"Error: Environment variable {name} not set.") sys.exit(1) return val def load_permissions(user_login): """ Reads the permissions JSON from the local file system and returns the permissions dict for the specific user. """ try: print(f"Loading permissions from {PERMISSIONS_FILE_PATH}...") if not os.path.exists(PERMISSIONS_FILE_PATH): print(f"Error: Permissions file not found at {PERMISSIONS_FILE_PATH}") return None with open(PERMISSIONS_FILE_PATH, "r") as f: data = json.load(f) user_perms = data.get(user_login) if not user_perms: print(f"User '{user_login}' not found in permissions file.") return None return user_perms except Exception as e: print(f"Failed to load or parse permissions file: {e}") sys.exit(1) def has_sgl_kernel_changes(pr): """ Check if the PR has changes to the sgl-kernel directory. This is used to determine if we need a full workflow rerun (to rebuild the kernel) vs just rerunning failed jobs. """ try: files = pr.get_files() for f in files: if f.filename.startswith("sgl-kernel/"): return True return False except Exception as e: print(f"Warning: Could not check PR files for sgl-kernel changes: {e}") # Default to False to avoid unnecessary full reruns return False def handle_tag_run_ci( gh_repo, pr, comment, user_perms, react_on_success=True, tag_extra=False ): """ Handles the /tag-run-ci-label command. When tag_extra is True (triggered by the `extra` argument), also adds the `run-ci-extra` label. pr-test-extra.yml gates on BOTH `run-ci` and `run-ci-extra`, so both must be present for the extra workflow to run — we always add `run-ci` alongside `run-ci-extra`. Reuses the same `can_tag_run_ci_label` permission. How fresh runs get dispatched: pr-test.yml and pr-test-extra.yml both include `labeled` in `on.pull_request.types`, so adding a label fires a new `pull_request.labeled` event with the up-to-date label set in its payload, which spawns a fresh workflow run that satisfies the `check-changes.if` gate. Note that this is the ONLY way to "un-skip" a label-gated run — `run.rerun()` on a previously-skipped pull_request run reuses the original event payload (frozen labels), so it would skip again. handle_rerun_failed_ci can't recover label-skipped runs; the labeled event is the recovery mechanism. Returns True if action was taken, False otherwise. """ if not user_perms.get("can_tag_run_ci_label", False): print("Permission denied: can_tag_run_ci_label is false.") return False labels = ["run-ci"] if tag_extra: labels.append("run-ci-extra") print(f"Permission granted. Adding labels: {labels}.") for label in labels: pr.add_to_labels(label) if react_on_success: comment.create_reaction("+1") print("Labels added and comment reacted.") else: print("Labels added (reaction suppressed).") return True def handle_rerun_failed_ci(gh_repo, pr, comment, user_perms, react_on_success=True): """ Handles the /rerun-failed-ci command. Reruns workflows with 'failure' or 'skipped' conclusions. Returns True if action was taken, False otherwise. """ if not user_perms.get("can_rerun_failed_ci", False): print("Permission denied: can_rerun_failed_ci is false.") return False print("Permission granted. Triggering rerun of failed or skipped workflows.") # Check if PR has sgl-kernel changes - if so, we may need full reruns # to ensure sgl-kernel-build-wheels runs and produces fresh artifacts. # However, if the wheel already built successfully for this commit, # we can just rerun failed jobs — the artifact is already there. sgl_kernel_changes = has_sgl_kernel_changes(pr) if sgl_kernel_changes: print("PR has sgl-kernel changes - checking if kernel wheel already built") # Get the SHA of the latest commit in the PR head_sha = pr.head.sha print(f"Checking workflows for commit: {head_sha}") # If PR has sgl-kernel changes, check whether ALL wheel builds already # succeeded for this commit (CUDA + ARM). If so, we can use # rerun_failed_jobs and avoid retriggering all tests. If any wheel # build is pending/failed, a dependent job could fail for missing # artifacts, so fall back to full rerun. # Check-runs display names: "Build Wheel (, )" (CUDA) and # "Build Wheel Arm (, )" (ARM). The YAML job ids # sgl-kernel-build-wheels{,-arm} are NOT what the check-runs API # returns — it returns the job's `name:` field. kernel_wheel_built = False if sgl_kernel_changes: try: wheel_builds = [ cr for cr in gh_repo.get_commit(head_sha).get_check_runs() if cr.name.startswith("Build Wheel") ] kernel_wheel_built = bool(wheel_builds) and all( cr.conclusion == "success" for cr in wheel_builds ) print( f"All {len(wheel_builds)} kernel wheel build(s) passed - using rerun_failed_jobs" if kernel_wheel_built else f"Kernel wheel not fully built " f"({sum(1 for c in wheel_builds if c.conclusion == 'success')}" f"/{len(wheel_builds)} success) - will use full rerun" ) except Exception as e: print( f"Failed to check kernel wheel status: {e} - falling back to full rerun" ) # Rerun workflows with conclusion=failure or conclusion=skipped. # # - failure: use rerun_failed_jobs() which reruns failed jobs *and their # dependent jobs* (GitHub API). Fast-fail cascades call # core.setFailed(...) so their conclusion is "failure" and are covered. # - skipped: the entire run was skipped (no jobs ran), so there are no # failed jobs for rerun_failed_jobs() to target. Use run.rerun(). # # Caveat: GitHub's `run.rerun()` reuses the original event payload, so # reruns of `pull_request`-event runs that were skipped because their # `if` evaluated to false (e.g. missing label) will skip again — the # label set in the frozen payload doesn't update. To un-skip a # label-gated workflow, add the missing label (the `labeled` event # dispatches a fresh run with the current label set); this function # cannot recover those by rerun alone. # - kernel wheel escape: if the PR touches sgl-kernel and not all wheel # builds are success yet, full-rerun failure runs too — Build Wheel # lives in pr-test-sgl-kernel.yml, consumers in pr-test.yml, and # rerun_failed_jobs() is scoped to a single workflow run. runs = gh_repo.get_workflow_runs(head_sha=head_sha) rerun_count = 0 for run in runs: if run.status != "completed": continue if run.conclusion not in ("failure", "skipped"): continue print(f"Processing {run.conclusion} workflow: {run.name} (ID: {run.id})") try: if run.conclusion == "skipped" or ( sgl_kernel_changes and not kernel_wheel_built ): print(" Full rerun") run.rerun() else: print(" rerun_failed_jobs") run.rerun_failed_jobs() rerun_count += 1 except Exception as e: print(f"Failed to rerun workflow {run.id}: {e}") if rerun_count > 0: print(f"Triggered rerun for {rerun_count} workflows.") if react_on_success: comment.create_reaction("+1") return True else: print("No failed or skipped workflows found to rerun.") return False MULTIMODAL_TEST_DIR = "python/sglang/multimodal_gen/test" MULTIMODAL_PATH_TO_RUNNER = { "2_gpu": "2-gpu-h100", "2-gpu": "2-gpu-h100", } MULTIMODAL_DEFAULT_RUNNER = "1-gpu-h100" def _known_test_groups(): groups = [] for group_dir in glob.glob("test/registered/*"): if os.path.isdir(group_dir): groups.append(os.path.basename(group_dir)) return sorted(groups) def resolve_test_group_specs(group_name): """ Resolve a test group name into /rerun-test specs. A group maps to a directory under test/registered/. For example, "hicache" maps to all test_*.py files under test/registered/hicache/. Returns (test_specs, error_message). On success error_message is None. """ group_name = group_name.strip().strip("/") if ( not group_name or group_name.startswith(".") or "/." in group_name or ".." in group_name.split("/") ): return [], f"Invalid test group `{group_name}`." group_dir = os.path.join("test", "registered", group_name) if not os.path.isdir(group_dir): known = ", ".join(f"`{g}`" for g in _known_test_groups()) return ( [], f"Unknown test group `{group_name}`.\n\nKnown groups: {known}", ) test_files = sorted( glob.glob(os.path.join(group_dir, "**", "test_*.py"), recursive=True) ) if not test_files: return [], f"No registered test files found in `{group_dir}`." return [os.path.relpath(path, "test") for path in test_files], None # A spec is treated as a wildcard pattern (expanding to many files) when its # file part contains a glob metacharacter. Plain specs keep the existing # single-file resolution, which requires a unique match. _GLOB_METACHARS = ("*", "?", "[") def _is_glob_pattern(file_part): return any(ch in file_part for ch in _GLOB_METACHARS) def expand_glob_spec(file_part): """ Expand a wildcard file_part into matching test files (repo-relative paths). Globs are matched against the same locations resolve_test_file() searches — test/registered/ and the multimodal_gen test dir — so e.g. `test_*backend*.py` reruns every backend test without hand-enumerating each file. Two constraints keep a broad pattern from pulling in non-tests: a match must live under a known test root and be named `test_*.py`. glob's `*` matches path separators only via `**`, so a bare pattern is searched recursively under each root; a path-ful pattern is anchored. Returns (sorted_repo_relative_paths, error). On success error is None. """ pat = file_part if pat.startswith("test/"): pat = pat[len("test/") :] matches = set() if "/" in pat: # Path-ful pattern. Glob from the repo root (handles fully qualified # multimodal paths like python/sglang/multimodal_gen/test/**/test_*.py) # and under test/ (handles test/-relative patterns like # registered/attention/test_*.py). for base in (".", "test"): matches.update(glob.glob(os.path.join(base, pat), recursive=True)) else: # Bare pattern: search recursively under each known test root. for root in ("test/registered", MULTIMODAL_TEST_DIR): matches.update(glob.glob(os.path.join(root, "**", pat), recursive=True)) # If the literal glob matched any directories — e.g. `unittest/*` matching # subdirs `dense/`, `dsa/`, … — descend recursively into them for # test_*.py. glob's `*` doesn't cross `/`, so without this a tree of # subdirs containing tests would return only the top-level entries (and # then get filtered out as non-test files). expanded = set() for p in matches: if os.path.isdir(p): expanded.update( glob.glob(os.path.join(p, "**", "test_*.py"), recursive=True) ) else: expanded.add(p) matches = expanded def _under_test_root(path): return path.startswith("test/registered/") or path.startswith( MULTIMODAL_TEST_DIR + "/" ) files = sorted( { os.path.normpath(p) for p in matches if os.path.isfile(p) and os.path.basename(p).startswith("test_") and p.endswith(".py") and _under_test_root(os.path.normpath(p)) } ) if not files: return [], ( f"No test files matched wildcard `{file_part}` under " f"`test/registered/` or `{MULTIMODAL_TEST_DIR}/` " f"(patterns only match files named `test_*.py`)." ) return files, None def resolve_test_file(file_part): """ Resolve a user-provided file path to a path relative to test/ or full path for multimodal. Supports: - Full path: test/registered/core/test_srt_endpoint.py - Relative to test/: registered/core/test_srt_endpoint.py - Bare filename: test_srt_endpoint.py (glob-matched, must be unique) - Multimodal paths: python/sglang/multimodal_gen/test/server/test_server_a.py Returns (resolved_path, is_multimodal, error_message). On success error_message is None. """ # Check if it's explicitly a multimodal path multimodal_prefixes = [ "python/sglang/multimodal_gen/test/", "sglang/multimodal_gen/test/", "multimodal_gen/test/", ] for prefix in multimodal_prefixes: if file_part.startswith(prefix): full_path = ( file_part if file_part.startswith("python/") else f"python/sglang/multimodal_gen/test/{file_part[len(prefix):]}" ) if not os.path.isfile(full_path): return None, False, f"File not found: `{full_path}`" return full_path, True, None # Existing logic for test/registered/ paths if file_part.startswith("test/"): file_part = file_part[len("test/") :] if "/" not in file_part: # Try test/registered/ first matches = glob.glob(f"test/registered/**/{file_part}", recursive=True) # Try multimodal test directory mm_matches = glob.glob(f"{MULTIMODAL_TEST_DIR}/**/{file_part}", recursive=True) # Filter to only test files mm_matches = [m for m in mm_matches if os.path.basename(m).startswith("test_")] if len(matches) == 1 and len(mm_matches) == 0: return matches[0][len("test/") :], False, None if len(matches) == 0 and len(mm_matches) == 1: return mm_matches[0], True, None all_matches = matches + mm_matches if len(all_matches) == 0: return ( None, False, f"No test file found matching `{file_part}` under `test/registered/` or `{MULTIMODAL_TEST_DIR}/`.", ) if len(all_matches) > 1: match_list = "\n".join(f"- `{m}`" for m in sorted(all_matches)) return ( None, False, ( f"Ambiguous filename `{file_part}` — matched {len(all_matches)} files:\n\n" f"{match_list}\n\n" f"Please provide the full path, e.g. `/rerun-test {all_matches[0]}`" ), ) # Shouldn't reach here, but handle gracefully if mm_matches: return mm_matches[0], True, None return matches[0][len("test/") :], False, None # Path with directory - check test/ location full_path = f"test/{file_part}" if os.path.isfile(full_path): return file_part, False, None return None, False, f"File not found: `{full_path}`" def detect_multimodal_suite(file_path): """ Determine runner for a multimodal gen test file based on its path. Returns (runner_label, error_message). """ # Check path components and basename for GPU count hints for pattern, runner in MULTIMODAL_PATH_TO_RUNNER.items(): if pattern in file_path: return runner, None return MULTIMODAL_DEFAULT_RUNNER, None def _extract_runner_configs(content): """Pull `(runner_config, args_str)` from EVERY `register_cuda_ci(...)` call. A test file can register itself on multiple pools (e.g. both `4-gpu-b200` and `1-gpu-large`). The earlier `re.search` variant returned only the first match, so /rerun-test silently dropped every registration after the first — multi-pool files only ever ran on the pool listed first. `re.finditer` is what makes the fan-out happen. """ out = [] for args in re.finditer( r"^[^#\n]*register_cuda_ci\s*\(([^)]*)\)", content, re.MULTILINE ): m = re.search(r'runner_config\s*=\s*["\']([^"\']+)["\']', args.group(1)) if m: out.append((m.group(1), args.group(1))) return out def _extract_legacy_suites(content): """Pull every legacy single-string `suite=` from `register_cuda_ci(...)` calls. Mirrors _extract_runner_configs for the legacy nightly/weekly shape: a file may register on multiple pools, so collect all of them rather than the first. """ out = [] for args in re.finditer( r"^[^#\n]*register_cuda_ci\s*\(([^)]*)\)", content, re.MULTILINE ): m = re.search(r'suite\s*=\s*["\']([^"\']+)["\']', args.group(1)) if m: out.append(m.group(1)) return out # Legacy nightly/weekly CUDA suites register with a single-string `suite=` # instead of `runner_config=`, so they carry no runner metadata of their own. # Map each to the runner_config in scripts/ci/runner_configs.yml whose hardware # matches the runner the nightly/weekly pipeline actually uses (see # .github/workflows/{nightly,weekly}-test-nvidia.yml), so /rerun-test can still # dispatch a single nightly/weekly test. The runner label, install script, # timeout, grace_blackwell, and rdma_devices are then resolved from # runner_configs.yml as usual, keeping that file the single source of truth for # runner details. # # Suites on hardware with no matching runner_config (e.g. nightly-4-gpu-gb300) # and non-CUDA suites (npu/amd) are intentionally absent and stay # non-dispatchable until a matching runner_config exists. _LEGACY_SUITE_TO_RUNNER_CONFIG = { "nightly-1-gpu": "1-gpu-large", "nightly-kernel-1-gpu": "1-gpu-large", "nightly-eval-text-2-gpu": "2-gpu-large", "nightly-perf-text-2-gpu": "2-gpu-large", "nightly-eval-vlm-2-gpu": "2-gpu-large", "nightly-perf-vlm-2-gpu": "2-gpu-large", "nightly-4-gpu": "4-gpu-h100", "nightly-4-gpu-b200": "4-gpu-b200", "nightly-8-gpu-common": ["8-gpu-h200", "8-gpu-b200"], "nightly-8-gpu-h200": "8-gpu-h200", "nightly-kernel-8-gpu-h200": "8-gpu-h200", "nightly-precision-8-gpu-h200": "8-gpu-h200", "nightly-8-gpu-h20": "8-gpu-h20", "nightly-8-gpu-b200": "8-gpu-b200", "weekly-8-gpu-h200": "8-gpu-h200", } def _dispatch_err(suite, msg): """Build a detect_suite error result for the given suite.""" return { "suite": suite, "runner_label": None, "install_script": "", "install_timeout": "", "grace_blackwell": "0", "rdma_devices": "", "is_cpu": False, "error": msg, } def _resolve_runner_config(rc, full_path, suite): """Resolve a runner_config key into a detect_suite dispatch dict. Returns the dispatch dict on success, or an error dict on failure. """ configs = _runner_configs.load() cfg = configs.get(rc) if cfg is None: known = ", ".join(f"`{k}`" for k in sorted(configs)) return _dispatch_err( suite, f"Unknown runner_config `{rc}` in `{full_path}` " f"— not in scripts/ci/runner_configs.yml.\n\n" f"Known runner_configs: {known}", ) install_script = cfg["install"] if not _ALLOWED_INSTALL_SCRIPT.match(install_script): return _dispatch_err( suite, f"Disallowed `install` value `{install_script}` for runner_config " f"`{rc}` in scripts/ci/runner_configs.yml. The slash handler " f"passes this string verbatim into a shell step, so it must " f"match `scripts/ci/cuda/*.sh`.", ) runs_on = cfg.get("runs_on") # Resolve $b200_runner sentinel: rerun-test never builds sgl-kernel, # so always pick the non-kernel b200 pool. if runs_on == "$b200_runner": runs_on = _B200_DEFAULT_RUNNER return { "suite": suite, "runner_label": runs_on, "install_script": install_script, "install_timeout": str(cfg["install_timeout"]), "grace_blackwell": str(cfg.get("grace_blackwell", "0")), "rdma_devices": cfg.get("rdma_devices", ""), "is_cpu": False, "error": None, } def detect_suite(file_path_from_test): """ Read a test file and extract dispatch info from register_cuda_ci or register_cpu_ci. A CUDA file can carry multiple `register_cuda_ci(...)` calls — one per pool it should run on — so this returns a *list* of dispatch dicts, one per registration. Runner label, install script, timeout, grace_blackwell, and rdma_devices are all resolved from scripts/ci/runner_configs.yml — the same single source of truth that drives the main PR test pipeline. Legacy nightly/weekly CUDA suites (single-string `suite=`) are dispatchable too: each suite name is mapped to the matching runner_config via _LEGACY_SUITE_TO_RUNNER_CONFIG, then resolved the same way. CPU files yield a single-element list. A file with no recognised (or no dispatchable) registration yields a one-element list whose dict has an `error` set. Each dict has keys: suite, runner_label, install_script, install_timeout, grace_blackwell, rdma_devices, is_cpu, error. """ full_path = f"test/{file_path_from_test}" with open(full_path, "r") as f: content = f.read() cuda_calls = _extract_runner_configs(content) if cuda_calls: results = [] for rc, args_str in cuda_calls: stage_m = re.search(r'stage\s*=\s*["\']([^"\']+)["\']', args_str) suite = f"{stage_m.group(1)}-test-{rc}" if stage_m else rc results.append(_resolve_runner_config(rc, full_path, suite)) return results # Legacy nightly/weekly CUDA suites: single-string `suite=`, no # runner_config. Map each mappable suite to its runner_config and resolve. legacy_suites = _extract_legacy_suites(content) mappable = [s for s in legacy_suites if s in _LEGACY_SUITE_TO_RUNNER_CONFIG] if mappable: results = [] for s in mappable: rcs = _LEGACY_SUITE_TO_RUNNER_CONFIG[s] if isinstance(rcs, str): rcs = [rcs] for rc in rcs: results.append(_resolve_runner_config(rc, full_path, s)) return results if re.search(r"^[^#\n]*register_cpu_ci\s*\(", content, re.MULTILINE): return [ { "suite": "cpu", "runner_label": "ubuntu-latest", "install_script": "", "install_timeout": "", "grace_blackwell": "0", "rdma_devices": "", "is_cpu": True, "error": None, } ] if legacy_suites: suite = legacy_suites[0] return [ _dispatch_err( suite, f"Suite `{suite}` in `{full_path}` is not dispatchable via " f"/rerun-test. It has no entry in _LEGACY_SUITE_TO_RUNNER_CONFIG " f"— either it is a non-CUDA suite (npu/amd) or it runs on " f"hardware with no matching runner_config in " f"scripts/ci/runner_configs.yml.", ) ] return [ _dispatch_err( None, f"No `register_cuda_ci(runner_config=...)` or `register_cpu_ci()` " f"found in `{full_path}`. This file may not be a registered CI test.", ) ] def _resolve_test_spec(test_spec): """ Resolve a single test spec into one or more dispatch entries. A file registered on N pools (multiple `register_cuda_ci(...)` calls) yields N entries — handle_rerun_test's grouping then sends one workflow per pool. Multimodal and CPU files yield a single entry. Resolution errors yield a one-element list with {"spec","error"}. """ if "::" in test_spec: file_part, test_selector = test_spec.split("::", 1) else: file_part = test_spec test_selector = None file_part = file_part.strip() if test_selector: test_selector = test_selector.strip() resolved_path, is_multimodal, err = resolve_test_file(file_part) if err: return [{"spec": test_spec, "error": err}] if is_multimodal: runner_label, err = detect_multimodal_suite(resolved_path) if err: return [{"spec": test_spec, "error": err}] # For multimodal pytest tests, use :: separator for test selection test_command = resolved_path if test_selector: test_command = f"{resolved_path}::{test_selector}" print( f"Resolved (multimodal_gen): file={resolved_path}, " f"selector={test_selector}, runner={runner_label}, " f"command='{test_command}'" ) return [ { "spec": test_spec, "test_command": test_command, "mode": "multimodal_gen", "runs_on": runner_label, "install_script": "", "install_timeout": "", "grace_blackwell": "0", "rdma_devices": "", "error": None, } ] test_command = resolved_path if test_selector: test_command = f"{resolved_path} {test_selector}" out = [] for info in detect_suite(resolved_path): if info["error"]: out.append({"spec": test_spec, "error": info["error"]}) continue mode = "cpu" if info["is_cpu"] else "cuda" print( f"Resolved: file={resolved_path}, selector={test_selector}, " f"suite={info['suite']}, mode={mode}, runs_on={info['runner_label']}, " f"install={info['install_script']}, grace_blackwell={info['grace_blackwell']}, " f"rdma={info['rdma_devices']}, " f"command='{test_command}'" ) out.append( { "spec": test_spec, "test_command": test_command, "mode": mode, "runs_on": info["runner_label"], "install_script": info["install_script"], "install_timeout": info["install_timeout"], "grace_blackwell": info["grace_blackwell"], "rdma_devices": info["rdma_devices"], "error": None, } ) return out def _dispatch_batch(gh_repo, pr, batch, token, reply_comment_id="", reply_marker=""): """ Dispatch a single workflow run for a batch of resolved test specs that share the same dispatch shape (mode + runs_on + install_script + install_timeout + grace_blackwell + rdma_devices). Returns a dict with keys: specs, success, test_commands, runner_label, run_url, error. """ test_commands = [r["test_command"] for r in batch] mode = batch[0]["mode"] runs_on = batch[0]["runs_on"] install_script = batch[0]["install_script"] install_timeout = batch[0]["install_timeout"] grace_blackwell = batch[0]["grace_blackwell"] rdma_devices = batch[0]["rdma_devices"] # Join multiple commands with newlines for the workflow to iterate over combined_command = "\n".join(test_commands) try: workflow_name = "Rerun Test" workflows = gh_repo.get_workflows() target_workflow = None for wf in workflows: if wf.name == workflow_name: target_workflow = wf break if not target_workflow: return { "specs": [r["spec"] for r in batch], "success": False, "error": f"{workflow_name} workflow not found", } is_fork = ( pr.head.repo is None or pr.head.repo.owner.login != gh_repo.owner.login ) pr_head_sha = None inputs = { "mode": mode, "test_command": combined_command, "runs_on": runs_on or "", "install_script": install_script, "install_timeout": install_timeout or "20", "grace_blackwell": grace_blackwell or "0", "rdma_devices": rdma_devices, "reply_comment_id": str(reply_comment_id) if reply_comment_id else "", "reply_marker": reply_marker, } if is_fork: ref = "main" pr_head_sha = pr.head.sha inputs["pr_head_sha"] = pr_head_sha else: ref = pr.head.ref dispatch_time = time.time() dispatch_url = f"https://api.github.com/repos/{gh_repo.full_name}/actions/workflows/{target_workflow.id}/dispatches" dispatch_resp = requests.post( dispatch_url, json={"ref": ref, "inputs": inputs}, headers={ "Authorization": f"Bearer {token}", "Accept": "application/vnd.github+json", }, ) success = dispatch_resp.status_code in (200, 204) if not success: print(f"Dispatch failed: {dispatch_resp.status_code} {dispatch_resp.text}") return { "specs": [r["spec"] for r in batch], "success": False, "error": f"Dispatch failed: {dispatch_resp.status_code}", } print(f"Successfully triggered rerun-test: {combined_command}") run_url = find_workflow_run_url( gh_repo, target_workflow.id, ref, "rerun-test", token, dispatch_time, pr_head_sha=pr_head_sha, max_wait=30, test_command=combined_command, ) return { "specs": [r["spec"] for r in batch], "success": True, "test_commands": test_commands, "mode": mode, "runs_on": runs_on, "run_url": run_url, "reply_marker": reply_marker, } except Exception as e: print(f"Error triggering rerun-test for batch: {e}") return { "specs": [r["spec"] for r in batch], "success": False, "error": str(e), } def _check_rerun_test_permissions(gh_repo, pr, comment, user_perms, command_name): """ Check permissions shared by /rerun-test and /rerun-group. """ # SECURITY: These commands check out and execute code from the PR branch on # self-hosted GPU runners, so fork PRs require a trusted collaborator. is_fork = pr.head.repo is None or pr.head.repo.owner.login != gh_repo.owner.login if is_fork: commenter = comment.user.login perm = gh_repo.get_collaborator_permission(commenter) if perm not in ("admin", "write"): print(f"Permission denied: /{command_name} on fork PR by {commenter}.") comment.create_reaction("confused") pr.create_issue_comment( f"⛔ `/{command_name}` is not available for fork PRs unless the commenter " "has write permission on the repo.\n\n" "Please ask a maintainer to run this command, or use the normal CI flow." ) return False print(f"Fork PR, but commenter {commenter} has write+ permission. Proceeding.") if not ( user_perms.get("can_rerun_test", False) or user_perms.get("can_rerun_stage", False) ): print("Permission denied: neither can_rerun_test nor can_rerun_stage is true.") return False return True def handle_rerun_test( gh_repo, pr, comment, user_perms, test_specs, token, skip_permission_check=False, command_label=None, ): """ Handles the /rerun-test command. Resolves all test specs, groups them by dispatch shape (mode + runs_on + install_script + install_timeout + grace_blackwell + rdma_devices), and dispatches one workflow per group. """ if not skip_permission_check and not _check_rerun_test_permissions( gh_repo, pr, comment, user_perms, "rerun-test" ): return False if not test_specs: comment.create_reaction("confused") pr.create_issue_comment( "⛔ Please specify a test: `/rerun-test ::`\n\n" "Examples:\n" "- `/rerun-test test/registered/core/test_srt_endpoint.py::TestSRTEndpoint.test_simple_decode`\n" "- `/rerun-test registered/core/test_srt_endpoint.py::TestSRTEndpoint`\n" "- `/rerun-test test_srt_endpoint.py`\n" "- `/rerun-test test_a.py test_b.py test_c.py` (multiple tests)\n" "- `/rerun-test test_*backend*.py` (wildcard — reruns every matching " "file; wrap the pattern in backticks so GitHub keeps the `*` literal)" ) return False allowed, gate_msg = _check_rebase_gate(gh_repo, pr, token) if not allowed: comment.create_reaction("confused") pr.create_issue_comment(gate_msg) return False # Phase 0: Expand wildcard specs into concrete test files. A spec whose # file part contains a glob metacharacter (* ? [) expands to every # matching file; plain specs pass through to single-file resolution. resolve_failures = [] seen_failures = set() def _record_failure(spec, error): # De-dupe failures by canonical path so the same un-dispatchable file # reported two ways (explicit + glob, or two globs) yields one line. key = spec.strip().strip("\"'`") if key.startswith("test/"): key = key[len("test/") :] if key not in seen_failures: seen_failures.add(key) resolve_failures.append({"spec": spec, "error": error}) expanded_specs = [] for spec in test_specs: # Quotes/backticks are never meaningful here (the command isn't # shell-parsed), so strip them. Backtick-wrapping is in fact the # recommended way to write a glob: plain `*backend*` renders as italics # in a GitHub comment, but `` `test_*backend*.py` `` stays literal — and # either way the handler reads the raw body, so the `*` survives. file_part = spec.split("::", 1)[0].strip().strip("\"'`") if not _is_glob_pattern(file_part): expanded_specs.append(spec) continue if "::" in spec: _record_failure( spec, "Wildcard patterns can't be combined with a `::test` " "selector — drop the `::...` to rerun whole files.", ) continue matched, err = expand_glob_spec(file_part) if err: _record_failure(spec, err) continue expanded_specs.extend(matched) # Phase 1: Resolve all specs, de-duping by the *resolved* identity. A glob # expands to `test/`-prefixed paths while an explicit spec keeps the form # the user typed, so the same file requested both ways resolves to two # different raw strings — keying de-dup on the resolved identity is what # collapses them into a single dispatch instead of running twice. # `runs_on` is included in the key so a file registered on N pools fans # out into one dispatch per pool, instead of all-but-the-first collapsing. resolved = [] seen_commands = set() for spec in expanded_specs: for r in _resolve_test_spec(spec): if r.get("error"): _record_failure(r["spec"], r["error"]) continue key = (r["mode"], r["runs_on"], r["test_command"]) if key in seen_commands: continue seen_commands.add(key) resolved.append(r) # Phase 2: Group by dispatch shape. groups = {} for r in resolved: key = ( r["mode"], r["runs_on"], r["install_script"], r["install_timeout"], r["grace_blackwell"], r["rdma_devices"], ) groups.setdefault(key, []).append(r) # Phase 3a: Create placeholder reply comment so we have its ID before # dispatching workflows. This lets each dispatched run write its # success/failure result back to the right line in this comment. dispatching = f"`{command_label}`" if command_label else "rerun-test workflow(s)" reply_comment = pr.create_issue_comment(f"🚀 Dispatching {dispatching}...") # Phase 3b: Dispatch one workflow per group, with a unique per-batch # marker each. The marker is an HTML comment that the writeback step # uses to locate the line and replace 🚀 with ✅/❌. dispatch_results = [] for idx, batch in enumerate(groups.values()): marker = f"" dispatch_results.append( _dispatch_batch( gh_repo, pr, batch, token, reply_comment_id=reply_comment.id, reply_marker=marker, ) ) # Build consolidated comment body (markers placed at line ends so the # writeback step can locate and update each line). lines = [] for dr in dispatch_results: if dr["success"]: if dr["mode"] == "multimodal_gen": cmds = "\n".join( f"python3 -m pytest {cmd} -x" for cmd in dr["test_commands"] ) else: cmds = "\n".join( f"cd test/ && python3 {cmd}" for cmd in dr["test_commands"] ) marker = dr.get("reply_marker", "") label = dr["runs_on"] or dr["mode"] if dr.get("run_url"): lines.append( f"🚀 `{label}` ({len(dr['test_commands'])} test{'s' if len(dr['test_commands']) > 1 else ''}): " f"⏳ [View workflow run]({dr['run_url']}) {marker}\n" f"```\n{cmds}\n```" ) else: lines.append( f"🚀 `{label}` ({len(dr['test_commands'])} test{'s' if len(dr['test_commands']) > 1 else ''}): ⏳ {marker}\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." ) 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) # Echo the originating command so each reply is self-identifying when # several /rerun-test commands are in flight at once. Backtick-wrapping # also keeps any `*` in the pattern from rendering as italics. if command_label: body = f"Results for `{command_label}`:\n\n{body}" successes = [dr for dr in dispatch_results if dr["success"]] if successes: comment.create_reaction("+1") if not successes and (resolve_failures or dispatch_results): comment.create_reaction("confused") reply_comment.edit(body) return len(successes) > 0 def handle_rerun_group( gh_repo, pr, comment, user_perms, group_names, token, command_label=None ): """ Handles the /rerun-group command. Expands one or more registered test groups into test file specs, then reuses /rerun-test dispatch behavior. """ if not _check_rerun_test_permissions( gh_repo, pr, comment, user_perms, "rerun-group" ): return False if not group_names: comment.create_reaction("confused") pr.create_issue_comment( "⛔ Please specify a test group: `/rerun-group `\n\n" "Example:\n" "- `/rerun-group hicache`" ) return False test_specs = [] failures = [] seen = set() for group_name in group_names: specs, err = resolve_test_group_specs(group_name) if err: failures.append((group_name, err)) continue for spec in specs: if spec not in seen: test_specs.append(spec) seen.add(spec) if failures: comment.create_reaction("confused") lines = [f"⛔ `{group}`: {err}" for group, err in failures] pr.create_issue_comment("\n\n".join(lines)) return False return handle_rerun_test( gh_repo, pr, comment, user_perms, test_specs, token, skip_permission_check=True, command_label=command_label, ) def main(): # 1. Load Environment Variables token = get_env_var("GITHUB_TOKEN") repo_name = get_env_var("REPO_FULL_NAME") pr_number = int(get_env_var("PR_NUMBER")) comment_id = int(get_env_var("COMMENT_ID")) comment_body = get_env_var("COMMENT_BODY").strip() user_login = get_env_var("USER_LOGIN") # 2. Load Permissions (local file check first to avoid unnecessary API calls) user_perms = load_permissions(user_login) # 3. Initialize GitHub API with Auth auth = Auth.Token(token) g = Github(auth=auth) repo = g.get_repo(repo_name) pr = repo.get_pull(pr_number) comment = repo.get_issue(pr_number).get_comment(comment_id) # PR authors can always rerun failed CI and rerun individual UTs on their own PRs, # even if they are not listed in CI_PERMISSIONS.json. # Note: /tag-run-ci-label still requires CI_PERMISSIONS.json. # Note: /rerun-test is blocked entirely for fork PRs in handle_rerun_test() itself. if pr.user.login == user_login: if user_perms is None: print( f"User {user_login} is the PR author (not in CI_PERMISSIONS.json). " "Granting CI rerun permissions." ) user_perms = {} else: print( f"User {user_login} is the PR author and has existing CI permissions." ) user_perms["can_rerun_failed_ci"] = True user_perms["can_rerun_test"] = True if not user_perms: print(f"User {user_login} does not have any configured permissions. Exiting.") return # 4. Parse Command and Execute first_line = comment_body.split("\n")[0].strip() # `extra` argument opts in to also tagging `run-ci-extra`. Both # `/tag-run-ci-label extra` and `/tag-and-rerun-ci extra` share this # parser so the surface is symmetric. tokens = first_line.split() tag_extra = len(tokens) > 1 and "extra" in tokens[1:] if first_line.startswith("/tag-run-ci-label"): handle_tag_run_ci(repo, pr, comment, user_perms, tag_extra=tag_extra) elif first_line.startswith("/rerun-failed-ci"): handle_rerun_failed_ci(repo, pr, comment, user_perms) elif first_line.startswith("/tag-and-rerun-ci"): # Perform both actions, but suppress individual reactions print(f"Processing combined command: /tag-and-rerun-ci (tag_extra={tag_extra})") tagged = handle_tag_run_ci( repo, pr, comment, user_perms, react_on_success=False, tag_extra=tag_extra ) # Wait for the label to propagate before triggering rerun if tagged: print("Waiting 5 seconds for label to propagate...") time.sleep(5) rerun = handle_rerun_failed_ci( repo, pr, comment, user_perms, react_on_success=False ) # If at least one action was successful, add the reaction here if tagged or rerun: comment.create_reaction("+1") print("Combined command processed successfully; reaction added.") else: print("Combined command finished, but no actions were taken.") elif first_line.startswith("/rerun-stage"): print("/rerun-stage is deprecated; posting deprecation notice.") comment.create_reaction("-1") pr.create_issue_comment( "⚠️ **`/rerun-stage` has been deprecated.**\n\n" "Stage granularity is too coarse — a stage usually doesn't map to one " "feature, so rerunning a stage re-pays the cost of unrelated tests. " "If you don't know which exact test files to rerun, you shouldn't be " "using `/rerun-stage` or `/rerun-test` in the first place.\n\n" "**Use one of these instead:**\n" "- **Selective tests** (you know exactly which files to rerun):\n" " ```\n" " /rerun-test test_foo.py test_bar.py\n" " ```\n" "- **Rerun only failed jobs**:\n" " ```\n" " /rerun-failed-ci\n" " ```\n" "- **Full CI rerun** (with extra coverage): add the `run-ci` or " "`run-ci-extra` label and push a new commit (or use `/tag-and-rerun-ci`).\n\n" "**AMD CI**: stage-level dispatch is still available via " "Actions UI → *PR Test (AMD)* / *PR Test ROCm 7.2 (AMD)* → " "*Run workflow* → pick a stage from the dropdown." ) elif first_line.startswith("/rerun-group"): group_names = first_line.split()[1:] handle_rerun_group( repo, pr, comment, user_perms, group_names or None, token, command_label=first_line, ) elif first_line.startswith("/rerun-test"): test_specs = first_line.split()[1:] handle_rerun_test( repo, pr, comment, user_perms, test_specs or None, token, command_label=first_line, ) else: print(f"Unknown or ignored command: {first_line}") if __name__ == "__main__": main()