Improve flaky PR job reruns

Author: trask

.github/scripts/rerun-flaky-pr-jobs.py

@@ -1,8 +1,8 @@
 #!/usr/bin/env python3
 """Rerun up to two failed jobs for recent pull request CI runs.
 
-Scans recent failed runs of build-pull-request.yml, ignores the synthetic
-required-status-check job, and reruns eligible failed jobs once.
+Scans recent failed pull request workflow runs, ignores the synthetic
+required-status-check job, and reruns eligible failed jobs up to two times.
 """
 
 from __future__ import annotations
@@ -16,7 +16,8 @@
 from pathlib import Path
 
 LOOKBACK_HOURS = 2
-MAX_ELIGIBLE_FAILURES = 2
+MAX_FAILED_JOBS_PER_WORKFLOW_RUN = 2
+MAX_RERUN_ATTEMPTS = 2
 
 
 def main() -> None:
@@ -25,7 +26,7 @@ def main() -> None:
     lookback_cutoff = datetime.now(timezone.utc).timestamp() - LOOKBACK_HOURS * 60 * 60
 
     recent_runs = list_recent_pull_request_runs(owner, repo, lookback_cutoff)
-    latest_run_by_pull_request = latest_run_per_pull_request(recent_runs)
+    latest_run_by_pull_request = latest_run_per_pull_request_workflow(owner, repo, recent_runs)
 
     processed_runs: list[str] = []
     rerun_jobs: list[str] = []
@@ -34,8 +35,11 @@ def main() -> None:
         if run["status"] != "completed" or run.get("conclusion") != "failure":
             continue
 
-        if run["run_attempt"] > 1:
-            processed_runs.append(f"Skipped {format_run(run)}: already rerun once.")
+        rerun_attempts = run["run_attempt"] - 1
+        if rerun_attempts > MAX_RERUN_ATTEMPTS:
+            processed_runs.append(
+                f"Skipped {format_run(run)}: already rerun {rerun_attempts} times."
+            )
             continue
 
         jobs = list_jobs_for_run(owner, repo, run["id"])
@@ -49,21 +53,20 @@ def main() -> None:
             processed_runs.append(f"Skipped {format_run(run)}: only synthetic jobs failed.")
             continue
 
-        if len(failed_real_jobs) > MAX_ELIGIBLE_FAILURES:
+        if len(failed_real_jobs) > MAX_FAILED_JOBS_PER_WORKFLOW_RUN:
             processed_runs.append(
-                f"Skipped {format_run(run)}: {len(failed_real_jobs)} failed jobs exceeded limit {MAX_ELIGIBLE_FAILURES}."
+                f"Skipped {format_run(run)}: {len(failed_real_jobs)} failed jobs exceeded limit {MAX_FAILED_JOBS_PER_WORKFLOW_RUN}."
             )
             continue
 
-        for job in failed_real_jobs:
-            try:
-                github_request("POST", f"/repos/{owner}/{repo}/actions/jobs/{job['id']}/rerun")
-                rerun_jobs.append(f"{format_run(run)}: reran {job['name']} ({job['id']}).")
-            except subprocess.CalledProcessError as e:
-                message = read_process_error(e)
-                processed_runs.append(
-                    f"Failed rerun for {format_run(run)} job {job['name']} ({job['id']}): {message}"
-                )
+        try:
+            github_request("POST", f"/repos/{owner}/{repo}/actions/runs/{run['id']}/rerun-failed-jobs")
+            rerun_jobs.append(f"{format_run(run)}: reran failed jobs {format_jobs(failed_real_jobs)}.")
+        except subprocess.CalledProcessError as e:
+            message = read_process_error(e)
+            processed_runs.append(
+                f"Failed rerun for {format_run(run)} jobs {format_jobs(failed_real_jobs)}: {message}"
+            )
 
     if not processed_runs and not rerun_jobs:
         processed_runs.append("No recent failed PR runs matched the rerun policy.")
@@ -86,9 +89,9 @@ def list_recent_pull_request_runs(owner: str, repo: str, lookback_cutoff: float)
     page = 1
 
     while True:
-        response = github_request(
+        response = github_request_object(
             "GET",
-            f"/repos/{owner}/{repo}/actions/workflows/build-pull-request.yml/runs",
+            f"/repos/{owner}/{repo}/actions/runs",
             {
                 "event": "pull_request",
                 "per_page": str(per_page),
@@ -114,19 +117,21 @@ def list_recent_pull_request_runs(owner: str, repo: str, lookback_cutoff: float)
     return [run for run in runs if parse_github_time(run["created_at"]).timestamp() >= lookback_cutoff]
 
 
-def latest_run_per_pull_request(runs: list[dict]) -> dict[int, dict]:
-    latest_by_pr: dict[int, dict] = {}
+def latest_run_per_pull_request_workflow(owner: str, repo: str, runs: list[dict]) -> dict[tuple[int, int], dict]:
+    latest_by_pr_workflow: dict[tuple[int, int], dict] = {}
+    branch_cache: dict[tuple[str, str], int | None] = {}
 
     for run in runs:
-        pr_number = get_pr_number(run)
+        pr_number = resolve_pr_number(owner, repo, run, branch_cache)
         if pr_number is None:
             continue
 
-        existing = latest_by_pr.get(pr_number)
+        key = (pr_number, run["workflow_id"])
+        existing = latest_by_pr_workflow.get(key)
         if existing is None or parse_github_time(run["created_at"]) > parse_github_time(existing["created_at"]):
-            latest_by_pr[pr_number] = run
+            latest_by_pr_workflow[key] = run
 
-    return latest_by_pr
+    return latest_by_pr_workflow
 
 
 def list_jobs_for_run(owner: str, repo: str, run_id: int) -> list[dict]:
@@ -135,7 +140,7 @@ def list_jobs_for_run(owner: str, repo: str, run_id: int) -> list[dict]:
     page = 1
 
     while True:
-        response = github_request(
+        response = github_request_object(
             "GET",
             f"/repos/{owner}/{repo}/actions/runs/{run_id}/jobs",
             {
@@ -153,7 +158,7 @@ def list_jobs_for_run(owner: str, repo: str, run_id: int) -> list[dict]:
         page += 1
 
 
-def github_request(method: str, path: str, query: dict[str, str] | None = None) -> dict:
+def github_request(method: str, path: str, query: dict[str, str] | None = None) -> dict | list[dict]:
     url = path.removeprefix("/")
     if query:
         url += "?" + urllib.parse.urlencode(query)
@@ -180,6 +185,20 @@ def github_request(method: str, path: str, query: dict[str, str] | None = None)
     return json.loads(result.stdout)
 
 
+def github_request_object(method: str, path: str, query: dict[str, str] | None = None) -> dict:
+    response = github_request(method, path, query)
+    if isinstance(response, list):
+        raise TypeError(f"Expected object response for {path}")
+    return response
+
+
+def github_request_list(method: str, path: str, query: dict[str, str] | None = None) -> list[dict]:
+    response = github_request(method, path, query)
+    if not isinstance(response, list):
+        raise TypeError(f"Expected list response for {path}")
+    return response
+
+
 def read_process_error(error: subprocess.CalledProcessError) -> str:
     return error.stderr.strip() or error.stdout.strip() or f"exit code {error.returncode}"
 
@@ -194,19 +213,60 @@ def get_current_repository() -> str:
     return result.stdout.strip()
 
 
-def get_pr_number(run: dict) -> int | None:
+def resolve_pr_number(
+    owner: str, repo: str, run: dict, branch_cache: dict[tuple[str, str], int | None]
+) -> int | None:
+    cached_pr_number = run.get("resolved_pr_number")
+    if isinstance(cached_pr_number, int):
+        return cached_pr_number
+
     pull_requests = run.get("pull_requests") or []
-    if not pull_requests:
+    if pull_requests:
+        pr_number = pull_requests[0].get("number")
+        run["resolved_pr_number"] = pr_number
+        return pr_number
+
+    head_repository = run.get("head_repository") or {}
+    head_owner = (head_repository.get("owner") or {}).get("login")
+    head_branch = run.get("head_branch")
+    if not head_owner or not head_branch:
+        run["resolved_pr_number"] = None
         return None
-    return pull_requests[0].get("number")
+
+    cache_key = (head_owner, head_branch)
+    if cache_key not in branch_cache:
+        pull_requests = github_request_list(
+            "GET",
+            f"/repos/{owner}/{repo}/pulls",
+            {
+                "state": "open",
+                "head": f"{head_owner}:{head_branch}",
+                "per_page": "5",
+            },
+        )
+        pr_number = None
+        for pull_request in pull_requests:
+            if pull_request.get("head", {}).get("sha") == run.get("head_sha"):
+                pr_number = pull_request["number"]
+                break
+        if pr_number is None and pull_requests:
+            pr_number = pull_requests[0]["number"]
+        branch_cache[cache_key] = pr_number
+
+    run["resolved_pr_number"] = branch_cache[cache_key]
+    return branch_cache[cache_key]
 
 
 def format_run(run: dict) -> str:
-    pr_number = get_pr_number(run)
+    pr_number = run.get("resolved_pr_number")
     pr_label = f"PR #{pr_number}" if pr_number is not None else "PR unknown"
     return f"{pr_label}, run {run['id']}, attempt {run['run_attempt']}"
 
 
+def format_jobs(jobs: list[dict]) -> str:
+    return ", ".join(f"{job['name']} ({job['id']})" for job in jobs)
+
+
 def parse_github_time(value: str) -> datetime:
     return datetime.fromisoformat(value.replace("Z", "+00:00"))