fix: match est_time updates by backend, not just suite (#22563)

This commit is contained in:
Cheng Wan
2026-04-10 17:54:50 -07:00
committed by GitHub
parent 75223c5404
commit b5e4ae7b1a
+17 -17
View File
@@ -113,8 +113,8 @@ def parse_job_logs(repo, job_id):
def collect_timings(repo): def collect_timings(repo):
"""Collect per-file elapsed times from recent scheduled CI runs. """Collect per-file elapsed times from recent scheduled CI runs.
Returns dict mapping (relative_path, suite) -> list of elapsed times Returns dict mapping (relative_path, suite, backend) -> list of elapsed
(newest first). times (newest first).
""" """
workflow_id = get_workflow_id(repo) workflow_id = get_workflow_id(repo)
print(f"Found workflow '{WORKFLOW_NAME}' (id={workflow_id})") print(f"Found workflow '{WORKFLOW_NAME}' (id={workflow_id})")
@@ -122,7 +122,7 @@ def collect_timings(repo):
runs = get_scheduled_runs(repo, workflow_id) runs = get_scheduled_runs(repo, workflow_id)
print(f"Found {len(runs)} completed scheduled runs on main") print(f"Found {len(runs)} completed scheduled runs on main")
# timings[(rel_path, suite)] = [elapsed1, elapsed2, ...] # timings[(rel_path, suite, backend)] = [elapsed1, elapsed2, ...]
timings = defaultdict(list) timings = defaultdict(list)
runs_processed = 0 runs_processed = 0
@@ -145,9 +145,10 @@ def collect_timings(repo):
for job in test_jobs: for job in test_jobs:
suite = job_name_to_suite(job["name"]) suite = job_name_to_suite(job["name"])
backend = determine_backend(job["name"])
entries = parse_job_logs(repo, job["id"]) entries = parse_job_logs(repo, job["id"])
for rel_path, elapsed in entries: for rel_path, elapsed in entries:
key = (rel_path, suite) key = (rel_path, suite, backend)
timings[key].append(elapsed) timings[key].append(elapsed)
if runs_processed >= MAX_RUNS: if runs_processed >= MAX_RUNS:
@@ -156,15 +157,15 @@ def collect_timings(repo):
print( print(
f"\nProcessed {runs_processed} runs, " f"\nProcessed {runs_processed} runs, "
f"collected timings for {len(timings)} (file, suite) pairs" f"collected timings for {len(timings)} (file, suite, backend) pairs"
) )
return timings return timings
def compute_medians(timings): def compute_medians(timings):
"""Compute median of last TARGET_DATA_POINTS timings for each (file, suite). """Compute median of last TARGET_DATA_POINTS timings for each entry.
Returns dict mapping (rel_path, suite) -> median (int). Returns dict mapping (rel_path, suite, backend) -> median (int).
Only includes entries with >= MIN_DATA_POINTS data points. Only includes entries with >= MIN_DATA_POINTS data points.
""" """
medians = {} medians = {}
@@ -188,12 +189,12 @@ def update_est_times(medians, dry_run=False):
updated = 0 updated = 0
skipped = 0 skipped = 0
# Group medians by file: {rel_path: {suite: median}} # Group medians by file: {rel_path: [(suite, backend, median), ...]}
by_file = defaultdict(dict) by_file = defaultdict(list)
for (rel_path, suite), median in medians.items(): for (rel_path, suite, backend), median in medians.items():
by_file[rel_path][suite] = median by_file[rel_path].append((suite, backend, median))
for rel_path, suite_medians in sorted(by_file.items()): for rel_path, entries in sorted(by_file.items()):
filepath = REPO_ROOT / rel_path filepath = REPO_ROOT / rel_path
if not filepath.exists(): if not filepath.exists():
print(f" SKIP {rel_path}: file not found") print(f" SKIP {rel_path}: file not found")
@@ -203,11 +204,11 @@ def update_est_times(medians, dry_run=False):
content = filepath.read_text() content = filepath.read_text()
new_content = content new_content = content
for suite, median in suite_medians.items(): for suite, backend, median in entries:
# Match registration calls with this specific suite. # Match registration calls with this specific backend and suite.
# Handles: register_cuda_ci(est_time=300, suite="stage-c-test-4-gpu-h100") # Handles: register_cuda_ci(est_time=300, suite="stage-c-test-4-gpu-h100")
pattern = re.compile( pattern = re.compile(
rf"(register_\w+_ci\(est_time=)(\d+)" rf"(register_{backend}_ci\(est_time=)(\d+)"
rf'(,\s*suite="{re.escape(suite)}")' rf'(,\s*suite="{re.escape(suite)}")'
) )
match = pattern.search(new_content) match = pattern.search(new_content)
@@ -219,7 +220,6 @@ def update_est_times(medians, dry_run=False):
continue continue
new_content = pattern.sub(rf"\g<1>{median}\3", new_content) new_content = pattern.sub(rf"\g<1>{median}\3", new_content)
backend = re.match(r"register_(\w+)_ci", match.group(0)).group(1)
print( print(
f" {rel_path}: register_{backend}_ci " f" {rel_path}: register_{backend}_ci "
f'suite="{suite}" est_time={old_val} -> {median}' f'suite="{suite}" est_time={old_val} -> {median}'
@@ -256,7 +256,7 @@ def main():
print("\nComputing medians...") print("\nComputing medians...")
medians = compute_medians(timings) medians = compute_medians(timings)
print(f"Computed medians for {len(medians)} (file, suite) pairs") print(f"Computed medians for {len(medians)} (file, suite, backend) entries")
print("\nUpdating est_time values...") print("\nUpdating est_time values...")
updated, skipped = update_est_times(medians, dry_run=args.dry_run) updated, skipped = update_est_times(medians, dry_run=args.dry_run)