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):
"""Collect per-file elapsed times from recent scheduled CI runs.
Returns dict mapping (relative_path, suite) -> list of elapsed times
(newest first).
Returns dict mapping (relative_path, suite, backend) -> list of elapsed
times (newest first).
"""
workflow_id = get_workflow_id(repo)
print(f"Found workflow '{WORKFLOW_NAME}' (id={workflow_id})")
@@ -122,7 +122,7 @@ def collect_timings(repo):
runs = get_scheduled_runs(repo, workflow_id)
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)
runs_processed = 0
@@ -145,9 +145,10 @@ def collect_timings(repo):
for job in test_jobs:
suite = job_name_to_suite(job["name"])
backend = determine_backend(job["name"])
entries = parse_job_logs(repo, job["id"])
for rel_path, elapsed in entries:
key = (rel_path, suite)
key = (rel_path, suite, backend)
timings[key].append(elapsed)
if runs_processed >= MAX_RUNS:
@@ -156,15 +157,15 @@ def collect_timings(repo):
print(
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
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.
"""
medians = {}
@@ -188,12 +189,12 @@ def update_est_times(medians, dry_run=False):
updated = 0
skipped = 0
# Group medians by file: {rel_path: {suite: median}}
by_file = defaultdict(dict)
for (rel_path, suite), median in medians.items():
by_file[rel_path][suite] = median
# Group medians by file: {rel_path: [(suite, backend, median), ...]}
by_file = defaultdict(list)
for (rel_path, suite, backend), median in medians.items():
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
if not filepath.exists():
print(f" SKIP {rel_path}: file not found")
@@ -203,11 +204,11 @@ def update_est_times(medians, dry_run=False):
content = filepath.read_text()
new_content = content
for suite, median in suite_medians.items():
# Match registration calls with this specific suite.
for suite, backend, median in entries:
# Match registration calls with this specific backend and suite.
# Handles: register_cuda_ci(est_time=300, suite="stage-c-test-4-gpu-h100")
pattern = re.compile(
rf"(register_\w+_ci\(est_time=)(\d+)"
rf"(register_{backend}_ci\(est_time=)(\d+)"
rf'(,\s*suite="{re.escape(suite)}")'
)
match = pattern.search(new_content)
@@ -219,7 +220,6 @@ def update_est_times(medians, dry_run=False):
continue
new_content = pattern.sub(rf"\g<1>{median}\3", new_content)
backend = re.match(r"register_(\w+)_ci", match.group(0)).group(1)
print(
f" {rel_path}: register_{backend}_ci "
f'suite="{suite}" est_time={old_val} -> {median}'
@@ -256,7 +256,7 @@ def main():
print("\nComputing medians...")
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...")
updated, skipped = update_est_times(medians, dry_run=args.dry_run)