diff --git a/scripts/ci/update_est_time.py b/scripts/ci/update_est_time.py index 0f19ce2c1..5468f59e9 100755 --- a/scripts/ci/update_est_time.py +++ b/scripts/ci/update_est_time.py @@ -22,7 +22,7 @@ REPO_ROOT = Path(__file__).resolve().parent.parent.parent # Regex to extract per-file elapsed time from CI logs. # Matches lines like: # filename='/actions-runner/_work/sglang/sglang/test/registered/core/test_x.py', elapsed=120, ... -# filename='/actions-runner/_work/sglang/sglang/python/sglang/jit_kernel/tests/test_x.py', elapsed=120, ... +# filename='/actions-runner/_work/sglang/sglang/python/sglang/jit_kernel/tests/test_x.py', ... LOG_PATTERN = re.compile( r"filename='[^']*?/sglang/((?:test|python)/[^']+\.py)', elapsed=(\d+)," ) @@ -73,6 +73,15 @@ def get_successful_jobs(repo, run_id): return [j for j in data["jobs"] if j["conclusion"] == "success"] +def job_name_to_suite(job_name): + """Extract the suite name from a job name. + + Job names look like "stage-c-test-4-gpu-h100 (2)" or "stage-a-test-cpu". + Strip the partition suffix " (N)" to get the suite name. + """ + return re.sub(r"\s*\(\d+\)$", "", job_name) + + def determine_backend(job_name): """Determine backend from job name.""" name = job_name.lower() @@ -104,7 +113,7 @@ 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, backend) -> list of elapsed times + Returns dict mapping (relative_path, suite) -> list of elapsed times (newest first). """ workflow_id = get_workflow_id(repo) @@ -113,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, backend)] = [elapsed1, elapsed2, ...] + # timings[(rel_path, suite)] = [elapsed1, elapsed2, ...] timings = defaultdict(list) runs_processed = 0 @@ -135,10 +144,10 @@ def collect_timings(repo): ) for job in test_jobs: - backend = determine_backend(job["name"]) + suite = job_name_to_suite(job["name"]) entries = parse_job_logs(repo, job["id"]) for rel_path, elapsed in entries: - key = (rel_path, backend) + key = (rel_path, suite) timings[key].append(elapsed) if runs_processed >= MAX_RUNS: @@ -147,15 +156,15 @@ def collect_timings(repo): print( f"\nProcessed {runs_processed} runs, " - f"collected timings for {len(timings)} (file, backend) pairs" + f"collected timings for {len(timings)} (file, suite) pairs" ) return timings def compute_medians(timings): - """Compute median of last TARGET_DATA_POINTS timings for each (file, backend). + """Compute median of last TARGET_DATA_POINTS timings for each (file, suite). - Returns dict mapping (rel_path, backend) -> median (int). + Returns dict mapping (rel_path, suite) -> median (int). Only includes entries with >= MIN_DATA_POINTS data points. """ medians = {} @@ -170,17 +179,21 @@ def compute_medians(timings): def update_est_times(medians, dry_run=False): """Update est_time values in source files. + Each registration call is matched by both the function name and suite, + so files with multiple registrations for different suites get the correct + per-suite median. + Returns (updated_count, skipped_count). """ updated = 0 skipped = 0 - # Group medians by file + # Group medians by file: {rel_path: {suite: median}} by_file = defaultdict(dict) - for (rel_path, backend), median in medians.items(): - by_file[rel_path][backend] = median + for (rel_path, suite), median in medians.items(): + by_file[rel_path][suite] = median - for rel_path, backend_medians in sorted(by_file.items()): + for rel_path, suite_medians in sorted(by_file.items()): filepath = REPO_ROOT / rel_path if not filepath.exists(): print(f" SKIP {rel_path}: file not found") @@ -190,8 +203,13 @@ def update_est_times(medians, dry_run=False): content = filepath.read_text() new_content = content - for backend, median in backend_medians.items(): - pattern = re.compile(rf"(register_{backend}_ci\(est_time=)(\d+)") + for suite, median in suite_medians.items(): + # Match registration calls with this specific 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'(,\s*suite="{re.escape(suite)}")' + ) match = pattern.search(new_content) if not match: continue @@ -200,10 +218,11 @@ def update_est_times(medians, dry_run=False): if old_val == median: continue - new_content = pattern.sub(rf"\g<1>{median}", 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( f" {rel_path}: register_{backend}_ci " - f"est_time={old_val} -> {median}" + f'suite="{suite}" est_time={old_val} -> {median}' ) if new_content != content: @@ -237,7 +256,7 @@ def main(): print("\nComputing medians...") medians = compute_medians(timings) - print(f"Computed medians for {len(medians)} (file, backend) pairs") + print(f"Computed medians for {len(medians)} (file, suite) pairs") print("\nUpdating est_time values...") updated, skipped = update_est_times(medians, dry_run=args.dry_run)