From b5e4ae7b1a37a4b37109b751151437d48c1090a1 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Fri, 10 Apr 2026 17:54:50 -0700 Subject: [PATCH] fix: match est_time updates by backend, not just suite (#22563) --- scripts/ci/update_est_time.py | 34 +++++++++++++++++----------------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/scripts/ci/update_est_time.py b/scripts/ci/update_est_time.py index 5468f59e9..59a718729 100755 --- a/scripts/ci/update_est_time.py +++ b/scripts/ci/update_est_time.py @@ -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)