[CI] Fix CI authorization notification fallback (#51095)

Signed-off-by: khluu <khluu000@gmail.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>

.github/workflows/add_label_automerge.yml

@@ -7,7 +7,7 @@ on:
             - auto_merge_enabled
 jobs:
     add-label-on-auto-merge:
-        runs-on: ubuntu-latest
+        runs-on: [self-hosted, linux, x64, vllm-runners]
         steps:
             -   name: Add label
                 uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0

.github/workflows/notify-ci-authorized.yml

@@ -11,7 +11,8 @@ concurrency:
   group: >-
     notify-ci-authorized-${{
       github.event.pull_request.number ||
-      github.event.workflow_run.pull_requests[0].number
+      github.event.workflow_run.pull_requests[0].number ||
+      github.event.workflow_run.id
     }}
   cancel-in-progress: false
 
@@ -29,9 +30,8 @@ jobs:
       github.event.label.name == 'ready-run-all-tests')) ||
       (github.event_name == 'workflow_run' &&
       github.event.workflow_run.event == 'pull_request_review' &&
-      github.event.workflow_run.conclusion == 'success' &&
-      github.event.workflow_run.pull_requests[0])
-    runs-on: ubuntu-latest
+      github.event.workflow_run.conclusion == 'success')
+    runs-on: [self-hosted, linux, x64, vllm-runners]
     steps:
       - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
         with:

.github/workflows/record-ci-approval.yml

@@ -9,6 +9,6 @@ permissions: {}
 jobs:
   approved:
     if: github.event.review.state == 'approved'
-    runs-on: ubuntu-latest
+    runs-on: [self-hosted, linux, x64, vllm-runners]
     steps:
       - run: echo "Approval recorded"

.github/workflows/scripts/run_ci_command.py

@@ -203,6 +203,13 @@ class GitHubClient:
     def get_pr(self, number: int) -> dict[str, Any]:
         return self._request(self._repo_path(f"/pulls/{number}"))
 
+    def list_pulls_for_commit(self, commit: str) -> list[dict[str, Any]]:
+        commit = urllib.parse.quote(commit, safe="")
+        response = self._request(self._repo_path(f"/commits/{commit}/pulls"))
+        if not isinstance(response, list):
+            raise ApiError(None, "GitHub API returned an invalid pull request list.")
+        return response
+
     def get_permission(self, actor: str) -> str:
         username = urllib.parse.quote(actor, safe="")
         try:
@@ -622,7 +629,6 @@ def notify_authorized(
         if (
             event.get("action") != "labeled"
             or event["label"]["name"] not in READY_LABELS
-            or has_trusted_approval(github, pr["number"], trusted_users)
         ):
             return
     elif "review" in event:
@@ -650,6 +656,40 @@ def notify_authorized(
     )
 
 
+def resolve_workflow_run_pr(
+    workflow_run: Mapping[str, Any],
+    github: GitHubClient,
+) -> dict[str, Any] | None:
+    head_sha = str(workflow_run.get("head_sha", ""))
+    if not head_sha:
+        return None
+
+    seen: set[int] = set()
+
+    def find_matching_pr(
+        candidates: Sequence[Mapping[str, Any]],
+    ) -> dict[str, Any] | None:
+        for candidate in candidates:
+            candidate_number = candidate.get("number")
+            if not isinstance(candidate_number, int) or candidate_number in seen:
+                continue
+            seen.add(candidate_number)
+            try:
+                pr = github.get_pr(candidate_number)
+            except ApiError as error:
+                if error.status == 404:
+                    continue
+                raise
+            if pr["state"] == "open" and pr["head"]["sha"] == head_sha:
+                return pr
+        return None
+
+    associated_pr = find_matching_pr(workflow_run.get("pull_requests") or [])
+    if associated_pr is not None:
+        return associated_pr
+    return find_matching_pr(github.list_pulls_for_commit(head_sha))
+
+
 def handle_run_ci(
     *,
     actor: str,
@@ -897,10 +937,10 @@ def main() -> None:
         )
         return
     if event_name == "workflow_run":
-        pull_requests = event["workflow_run"].get("pull_requests") or []
-        if not pull_requests:
+        pr = resolve_workflow_run_pr(event["workflow_run"], github)
+        if pr is None:
+            print("Could not resolve an open PR for the approval workflow run.")
             return
-        pr = github.get_pr(int(pull_requests[0]["number"]))
         notify_authorized(
             {
                 "action": "submitted",

.github/workflows/scripts/test_run_ci_command.py

@@ -15,6 +15,7 @@ from run_ci_command import (
     COMMAND_RUN_CI_ALL,
     COMMAND_RUN_CI_NIGHTLY,
     RETRY_STATES,
+    ApiError,
     BuildkiteClient,
     HttpTransport,
     authorize,
@@ -25,6 +26,7 @@ from run_ci_command import (
     notify_authorized,
     parse_command,
     parse_trusted_users,
+    resolve_workflow_run_pr,
     run,
     select_latest_build,
 )
@@ -72,6 +74,8 @@ class FakeGitHub:
         review_decision: str = "REVIEW_REQUIRED",
         reviews: list[dict[str, Any]] | None = None,
         comments: list[str] | None = None,
+        pulls_for_commit: list[dict[str, Any]] | None = None,
+        prs: dict[int, dict[str, Any]] | None = None,
     ) -> None:
         self.comments = comments or []
         self.permission = permission
@@ -80,13 +84,21 @@ class FakeGitHub:
         self.reactions: list[str] = []
         self.review_decision = review_decision
         self.reviews = reviews or []
+        self.pulls_for_commit = pulls_for_commit or []
+        self.prs = prs or {self.pr["number"]: self.pr}
 
     def get_pr(self, number: int) -> dict[str, Any]:
-        return self.pr
+        try:
+            return self.prs[number]
+        except KeyError as error:
+            raise ApiError(404, "Not found") from error
 
     def get_permission(self, actor: str) -> str:
         return self.permissions.get(actor, self.permission)
 
+    def list_pulls_for_commit(self, commit: str) -> list[dict[str, Any]]:
+        return self.pulls_for_commit
+
     def get_review_decision(self, number: int) -> str:
         return self.review_decision
 
@@ -525,7 +537,7 @@ class RunCiCommandTest(unittest.TestCase):
         self.assertNotIn(COMMAND_RUN_CI_NIGHTLY, comment)
         self.assertIn(CI_AUTHORIZED_COMMENT_MARKER, comment)
 
-    def test_ready_label_does_not_notify_after_trusted_approval(self) -> None:
+    def test_ready_label_notifies_after_missed_approval_notification(self) -> None:
         pr = make_pr(labels=[{"name": "ready"}])
         event = {
             "action": "labeled",
@@ -547,7 +559,8 @@ class RunCiCommandTest(unittest.TestCase):
 
         notify_authorized(event, github)
 
-        self.assertEqual(github.comments, [])
+        self.assertEqual(len(github.comments), 1)
+        self.assertIn("@author", github.comments[0])
 
     def test_trusted_approval_notifies_author_once(self) -> None:
         event = {
@@ -573,6 +586,37 @@ class RunCiCommandTest(unittest.TestCase):
         self.assertEqual(len(github.comments), 1)
         self.assertIn("@author", github.comments[0])
 
+    def test_workflow_run_resolves_pr_from_head_commit(self) -> None:
+        pr = make_pr()
+        github = FakeGitHub(
+            pr=pr,
+            pulls_for_commit=[{"number": pr["number"]}],
+        )
+
+        resolved = resolve_workflow_run_pr(
+            {"head_sha": pr["head"]["sha"], "pull_requests": []},
+            github,
+        )
+
+        self.assertEqual(resolved, pr)
+
+    def test_workflow_run_ignores_unrelated_pr_association(self) -> None:
+        pr = make_pr()
+        github = FakeGitHub(
+            pr=pr,
+            pulls_for_commit=[{"number": pr["number"]}],
+        )
+
+        resolved = resolve_workflow_run_pr(
+            {
+                "head_sha": pr["head"]["sha"],
+                "pull_requests": [{"number": 1}],
+            },
+            github,
+        )
+
+        self.assertEqual(resolved, pr)
+
     def test_approval_does_not_notify_when_ready_label_exists(self) -> None:
         pr = make_pr(labels=[{"name": "ready"}])
         event = {