1540 lines
56 KiB
Python
1540 lines
56 KiB
Python
import glob
|
|
import json
|
|
import os
|
|
import re
|
|
import sys
|
|
import time
|
|
import unicodedata
|
|
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"
|
|
PRECISION_BASELINE_TEST = "registered/debug_utils/test_nightly_precision_regression.py"
|
|
PRECISION_BASELINE_REFRESH_FLAG = "--refresh-precision-baseline"
|
|
|
|
|
|
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: <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: <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] <test_command> [<sha>]"
|
|
# Other workflows: "[stage-name] [<sha>]"
|
|
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 _strip_format_chars(s):
|
|
"""Remove Unicode format characters (category Cf: LRM/RLM U+200E/200F,
|
|
zero-width space/joiners U+200B-200D, word joiner U+2060, BOM U+FEFF).
|
|
|
|
GitHub's copy-path button and rich-text copy inject these invisibly;
|
|
a pasted `/rerun-test foo.py<U+200E>` then never matches any test file
|
|
(see PR #31059). They are display hints and never legitimate in a
|
|
command or path, so dropping them is always safe.
|
|
"""
|
|
return "".join(c for c in s if unicodedata.category(c) != "Cf")
|
|
|
|
|
|
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 AOT 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("python/sglang/kernels/aot/"):
|
|
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 _latest_run_per_workflow(runs):
|
|
"""
|
|
Collapse a head_sha's workflow runs to the newest run per workflow.
|
|
|
|
GitHub can have several runs of the *same* workflow at one commit — a
|
|
`synchronize` run superseded by a `labeled` run, or a run cancelled by
|
|
`cancel-in-progress` (pr-test.yml sets it) while its replacement is
|
|
already in flight. Only the newest one reflects current state; rerunning
|
|
the older ones spawns duplicates that fight the live run for runners.
|
|
"""
|
|
latest = {}
|
|
for run in runs:
|
|
current = latest.get(run.workflow_id)
|
|
if current is None or run.id > current.id:
|
|
latest[run.workflow_id] = run
|
|
return list(latest.values())
|
|
|
|
|
|
def handle_rerun_failed_ci(gh_repo, pr, comment, user_perms, react_on_success=True):
|
|
"""
|
|
Handles the /rerun-failed-ci command.
|
|
Reruns workflows that ended in 'failure', 'skipped', 'cancelled' or
|
|
'timed_out', restarting only the jobs that didn't succeed.
|
|
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 unsuccessful 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 (<python>, <cuda>)" (CUDA) and
|
|
# "Build Wheel Arm (<python>, <cuda>)" (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 that ended in failure, skipped, cancelled or timed_out.
|
|
#
|
|
# - 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.
|
|
# - cancelled / timed_out: rerun_failed_jobs() as well. GitHub restarts
|
|
# every job that didn't succeed — cancelled ones included — plus their
|
|
# dependents, and carries the passing jobs over untouched. A full
|
|
# rerun here would re-execute dozens of already-green GPU jobs to
|
|
# recover one cancelled partition.
|
|
# A run cancelled before any job failed has nothing for the endpoint to
|
|
# target, so fall back to run.rerun() when it rejects the request.
|
|
# - 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 = _latest_run_per_workflow(gh_repo.get_workflow_runs(head_sha=head_sha))
|
|
|
|
rerun_count = 0
|
|
for run in runs:
|
|
if run.status != "completed":
|
|
# A newer attempt is still in flight - nothing to recover.
|
|
continue
|
|
if run.conclusion not in ("failure", "skipped", "cancelled", "timed_out"):
|
|
# action_required (fork PR awaiting approval) is deliberately left
|
|
# out: it needs an approval, not a rerun.
|
|
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:
|
|
try:
|
|
print(" rerun_failed_jobs")
|
|
run.rerun_failed_jobs()
|
|
except Exception as e:
|
|
print(f" rerun_failed_jobs rejected ({e}) - full rerun")
|
|
run.rerun()
|
|
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, skipped, cancelled or timed-out 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",
|
|
"b200": _B200_DEFAULT_RUNNER,
|
|
}
|
|
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_suites(content, register_fn):
|
|
"""Pull every single-string `suite=` from `<register_fn>(...)` calls."""
|
|
out = []
|
|
for args in re.finditer(
|
|
rf"^[^#\n]*{register_fn}\s*\(([^)]*)\)", content, re.MULTILINE
|
|
):
|
|
m = re.search(r'suite\s*=\s*["\']([^"\']+)["\']', args.group(1))
|
|
if m:
|
|
out.append(m.group(1))
|
|
return out
|
|
|
|
|
|
def _extract_legacy_suites(content):
|
|
"""Pull every legacy single-string `suite=` from `register_cuda_ci(...)`
|
|
calls. Used only to report why such a file is not dispatchable."""
|
|
return _extract_suites(content, "register_cuda_ci")
|
|
|
|
|
|
# Backends with no job in rerun-test.yml (cuda / multimodal_gen / cpu only) and
|
|
# no runner_config in runner_configs.yml, so no dispatch can be built for them.
|
|
# Mirrors `REGISTER_MAPPING` in python/sglang/test/ci/ci_register.py.
|
|
_OTHER_BACKEND_REGISTERS = {
|
|
"register_amd_ci": "AMD",
|
|
"register_npu_ci": "NPU",
|
|
"register_xpu_ci": "XPU",
|
|
"register_musa_ci": "MUSA",
|
|
"register_mlx_ci": "MLX",
|
|
}
|
|
|
|
|
|
def _extract_other_backends(content):
|
|
"""Return (backend labels, suite names) for every non-CUDA/CPU registration."""
|
|
labels, suites = [], []
|
|
for register_fn, label in _OTHER_BACKEND_REGISTERS.items():
|
|
if re.search(rf"^[^#\n]*{register_fn}\s*\(", content, re.MULTILINE):
|
|
labels.append(label)
|
|
suites.extend(_extract_suites(content, register_fn))
|
|
return labels, sorted(set(suites))
|
|
|
|
|
|
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": "",
|
|
"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"]),
|
|
"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, 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. Every
|
|
dispatchable CUDA suite, per-commit and scheduled alike, goes through that
|
|
one path; the legacy single-string `suite=` carries no runner_config and is
|
|
reported as non-dispatchable.
|
|
|
|
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, 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_suites = _extract_legacy_suites(content)
|
|
|
|
if re.search(r"^[^#\n]*register_cpu_ci\s*\(", content, re.MULTILINE):
|
|
return [
|
|
{
|
|
"suite": "cpu",
|
|
"runner_label": "ubuntu-latest",
|
|
"install_script": "",
|
|
"install_timeout": "",
|
|
"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 registered with the legacy "
|
|
f"single-string `suite=`, which carries no runner_config and so "
|
|
f"is not dispatchable via /rerun-test. Re-register it with "
|
|
f"`stage=`/`runner_config=` (CUDA), or dispatch its own "
|
|
f"workflow (npu/amd).",
|
|
)
|
|
]
|
|
|
|
labels, suites = _extract_other_backends(content)
|
|
if labels:
|
|
backends = ", ".join(labels)
|
|
where = f" (suite `{suites[0]}`)" if suites else ""
|
|
return [
|
|
_dispatch_err(
|
|
suites[0] if suites else None,
|
|
f"`{full_path}` is registered for {backends}{where}, not for "
|
|
f"CUDA or CPU; rerun-test.yml has no {backends} job. Rerun it "
|
|
f"with /rerun-failed-ci, or dispatch the {backends} workflow "
|
|
f"manually.",
|
|
)
|
|
]
|
|
|
|
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": "",
|
|
"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']}, "
|
|
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"],
|
|
"rdma_devices": info["rdma_devices"],
|
|
"error": None,
|
|
}
|
|
)
|
|
return out
|
|
|
|
|
|
def _dispatch_batch(
|
|
gh_repo,
|
|
pr,
|
|
batch,
|
|
token,
|
|
reply_comment_id="",
|
|
reply_marker="",
|
|
refresh_precision_baseline=False,
|
|
):
|
|
"""
|
|
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 + 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"]
|
|
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",
|
|
"rdma_devices": rdma_devices,
|
|
"reply_comment_id": str(reply_comment_id) if reply_comment_id else "",
|
|
"reply_marker": reply_marker,
|
|
"refresh_precision_baseline": str(refresh_precision_baseline).lower(),
|
|
}
|
|
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.
|
|
"""
|
|
# A rerun dispatches rerun-test.yml, which never passes through pr-gate.yml,
|
|
# so it is unthrottled either way; gate on what pr-gate waives the limit for.
|
|
if user_perms.get("cooldown_interval_minutes") == 0:
|
|
return True
|
|
|
|
commenter = comment.user.login
|
|
perm = gh_repo.get_collaborator_permission(commenter)
|
|
if perm in ("admin", "write"):
|
|
print(f"Commenter {commenter} has write+ permission. Proceeding.")
|
|
return True
|
|
|
|
print(f"Permission denied: /{command_name} by {commenter} (permission: {perm}).")
|
|
comment.create_reaction("confused")
|
|
pr.create_issue_comment(
|
|
f"⛔ `/{command_name}` requires `cooldown_interval_minutes: 0` in "
|
|
"`.github/CI_PERMISSIONS.json`, or write permission on the repo.\n\n"
|
|
"Please ask a maintainer to run this command, or use the normal CI flow."
|
|
)
|
|
return False
|
|
|
|
|
|
def _check_precision_baseline_refresh_permissions(gh_repo, pr, comment):
|
|
commenter = comment.user.login
|
|
is_fork = pr.head.repo is None or pr.head.repo.full_name != gh_repo.full_name
|
|
if is_fork:
|
|
comment.create_reaction("confused")
|
|
pr.create_issue_comment(
|
|
"⛔ Precision baseline refresh is only available on PR branches in "
|
|
"this repository. Fork PR code cannot receive the baseline write token."
|
|
)
|
|
return False
|
|
|
|
perm = gh_repo.get_collaborator_permission(commenter)
|
|
if perm not in ("admin", "maintain", "write"):
|
|
comment.create_reaction("confused")
|
|
pr.create_issue_comment(
|
|
"⛔ Precision baseline refresh requires write permission on the repo."
|
|
)
|
|
return False
|
|
return True
|
|
|
|
|
|
def handle_rerun_test(
|
|
gh_repo,
|
|
pr,
|
|
comment,
|
|
user_perms,
|
|
test_specs,
|
|
token,
|
|
skip_permission_check=False,
|
|
command_label=None,
|
|
refresh_precision_baseline=False,
|
|
):
|
|
"""
|
|
Handles the /rerun-test command. Resolves all test specs, groups them by
|
|
dispatch shape (mode + runs_on + install_script + install_timeout +
|
|
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 (
|
|
refresh_precision_baseline
|
|
and not _check_precision_baseline_refresh_permissions(gh_repo, pr, comment)
|
|
):
|
|
return False
|
|
|
|
if not test_specs:
|
|
comment.create_reaction("confused")
|
|
pr.create_issue_comment(
|
|
"⛔ Please specify a test: `/rerun-test <file>::<TestClass.test_method>`\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)
|
|
|
|
if refresh_precision_baseline:
|
|
is_exact_precision_test = (
|
|
not resolve_failures
|
|
and len(resolved) == 1
|
|
and resolved[0]["test_command"] == PRECISION_BASELINE_TEST
|
|
)
|
|
if not is_exact_precision_test:
|
|
comment.create_reaction("confused")
|
|
pr.create_issue_comment(
|
|
"⛔ `--refresh-precision-baseline` must be used alone with "
|
|
f"`test/{PRECISION_BASELINE_TEST}`."
|
|
)
|
|
return False
|
|
|
|
# Phase 2: Group by dispatch shape.
|
|
groups = {}
|
|
for r in resolved:
|
|
key = (
|
|
r["mode"],
|
|
r["runs_on"],
|
|
r["install_script"],
|
|
r["install_timeout"],
|
|
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"<!--rrt:{idx}-->"
|
|
dispatch_results.append(
|
|
_dispatch_batch(
|
|
gh_repo,
|
|
pr,
|
|
batch,
|
|
token,
|
|
reply_comment_id=reply_comment.id,
|
|
reply_marker=marker,
|
|
refresh_precision_baseline=refresh_precision_baseline,
|
|
)
|
|
)
|
|
|
|
# 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 <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 = _strip_format_chars(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 on their own PRs, even if they are not
|
|
# listed in CI_PERMISSIONS.json.
|
|
# Note: /tag-run-ci-label still requires CI_PERMISSIONS.json.
|
|
# Authorship grants nothing for /rerun-test; that gate reads the commenter.
|
|
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
|
|
|
|
# No early exit on a missing entry: /rerun-test also gates on repo permission,
|
|
# so a write-holder absent from the file must still reach its handler.
|
|
if user_perms is None:
|
|
user_perms = {}
|
|
|
|
# 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-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"):
|
|
rerun_args = first_line.split()[1:]
|
|
refresh_precision_baseline = PRECISION_BASELINE_REFRESH_FLAG in rerun_args
|
|
test_specs = [
|
|
arg for arg in rerun_args if arg != PRECISION_BASELINE_REFRESH_FLAG
|
|
]
|
|
handle_rerun_test(
|
|
repo,
|
|
pr,
|
|
comment,
|
|
user_perms,
|
|
test_specs or None,
|
|
token,
|
|
command_label=first_line,
|
|
refresh_precision_baseline=refresh_precision_baseline,
|
|
)
|
|
|
|
else:
|
|
print(f"Unknown or ignored command: {first_line}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|