[CI] Retry Buildkite API rate limits (#50481)

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

.github/workflows/scripts/run_ci_command.py

@@ -3,11 +3,13 @@
 
 import json
 import os
+import random
 import sys
+import time
 import urllib.error
 import urllib.parse
 import urllib.request
-from collections.abc import Mapping, Sequence
+from collections.abc import Callable, Mapping, Sequence
 from typing import Any
 
 COMMAND_RUN_CI = "/ci run"
@@ -39,7 +41,22 @@ class ApiError(RuntimeError):
         self.status = status
 
 
+def rate_limit_jitter() -> float:
+    return random.uniform(1, 5)
+
+
 class HttpTransport:
+    def __init__(
+        self,
+        *,
+        max_retries: int = 3,
+        jitter: Callable[[], float] = rate_limit_jitter,
+        sleep: Callable[[float], None] = time.sleep,
+    ) -> None:
+        self.max_retries = max_retries
+        self.jitter = jitter
+        self.sleep = sleep
+
     def request(
         self,
         url: str,
@@ -55,18 +72,33 @@ class HttpTransport:
             headers=dict(headers or {}),
             method=method,
         )
-        try:
-            with urllib.request.urlopen(request, timeout=30) as response:
-                response_body = response.read().decode()
-        except urllib.error.HTTPError as error:
-            response_body = error.read().decode()
-            message = self._error_message(response_body, error.reason)
-            raise ApiError(
-                error.code,
-                f"API returned {error.code}: {message}",
-            ) from error
-        except urllib.error.URLError as error:
-            raise ApiError(None, f"API request failed: {error.reason}") from error
+        for attempt in range(self.max_retries + 1):
+            try:
+                with urllib.request.urlopen(request, timeout=30) as response:
+                    response_body = response.read().decode()
+                break
+            except urllib.error.HTTPError as error:
+                response_body = error.read().decode()
+                if error.code == 429 and attempt < self.max_retries:
+                    delay = self._rate_limit_delay(error, response_body)
+                    print(
+                        "API rate limit reached; "
+                        f"retry {attempt + 1}/{self.max_retries} "
+                        f"in {delay:g} seconds.",
+                        file=sys.stderr,
+                    )
+                    self.sleep(delay)
+                    continue
+                message = self._error_message(response_body, error.reason)
+                raise ApiError(
+                    error.code,
+                    f"API returned {error.code}: {message}",
+                ) from error
+            except urllib.error.URLError as error:
+                raise ApiError(
+                    None,
+                    f"API request failed: {error.reason}",
+                ) from error
 
         if not response_body:
             return None
@@ -83,6 +115,33 @@ class HttpTransport:
             return fallback
         return str(parsed.get("message", fallback))
 
+    def _rate_limit_delay(
+        self,
+        error: urllib.error.HTTPError,
+        response_body: str,
+    ) -> float:
+        try:
+            parsed = json.loads(response_body)
+        except json.JSONDecodeError:
+            parsed = {}
+
+        scope = parsed.get("scope")
+        reset_header = (
+            "RateLimit-User-Reset" if scope == "rest_user" else "RateLimit-Reset"
+        )
+        candidates = [
+            error.headers.get(reset_header),
+            parsed.get("reset"),
+        ]
+        for candidate in candidates:
+            try:
+                delay = float(candidate)
+            except (TypeError, ValueError):
+                continue
+            if delay >= 0:
+                return delay + self.jitter()
+        return 60 + self.jitter()
+
 
 class GitHubClient:
     def __init__(

.github/workflows/scripts/test_run_ci_command.py

@@ -1,9 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import io
 import json
 import unittest
+import urllib.error
 from typing import Any
+from unittest.mock import patch
 
 from run_ci_command import (
     CI_AUTHORIZED_COMMENT_MARKER,
@@ -11,6 +14,7 @@ from run_ci_command import (
     COMMAND_RUN_CI,
     RETRY_STATES,
     BuildkiteClient,
+    HttpTransport,
     authorize,
     create_build_payload,
     has_trusted_approval,
@@ -160,7 +164,109 @@ class FakeTransport:
         return self.response
 
 
+class FakeHttpResponse:
+    def __init__(self, response: Any) -> None:
+        self.response = response
+
+    def __enter__(self) -> "FakeHttpResponse":
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        return None
+
+    def read(self) -> bytes:
+        return json.dumps(self.response).encode()
+
+
 class RunCiCommandTest(unittest.TestCase):
+    @patch("run_ci_command.urllib.request.urlopen")
+    def test_http_transport_retries_buildkite_rate_limit(self, urlopen: Any) -> None:
+        body = {
+            "message": "Please wait 9 seconds before making more requests.",
+            "reset": 9,
+            "scope": "rest",
+        }
+        rate_limit_error = urllib.error.HTTPError(
+            "https://api.buildkite.com/v2/builds",
+            429,
+            "Too Many Requests",
+            {
+                "RateLimit-Limit": "400",
+                "RateLimit-Remaining": "0",
+                "RateLimit-Reset": "9",
+            },
+            io.BytesIO(json.dumps(body).encode()),
+        )
+        urlopen.side_effect = [rate_limit_error, FakeHttpResponse({"ok": True})]
+        delays: list[float] = []
+
+        response = HttpTransport(
+            jitter=lambda: 2.5,
+            sleep=delays.append,
+        ).request("https://api.buildkite.com/v2/builds")
+
+        self.assertEqual(response, {"ok": True})
+        self.assertEqual(delays, [11.5])
+        self.assertEqual(urlopen.call_count, 2)
+
+    @patch("run_ci_command.urllib.request.urlopen")
+    def test_http_transport_retries_rate_limit_three_times(self, urlopen: Any) -> None:
+        def rate_limit_error() -> urllib.error.HTTPError:
+            body = {
+                "message": "Please wait 9 seconds before making more requests.",
+                "reset": 9,
+                "scope": "rest",
+            }
+            return urllib.error.HTTPError(
+                "https://api.buildkite.com/v2/builds",
+                429,
+                "Too Many Requests",
+                {
+                    "RateLimit-Limit": "400",
+                    "RateLimit-Remaining": "0",
+                    "RateLimit-Reset": "9",
+                },
+                io.BytesIO(json.dumps(body).encode()),
+            )
+
+        urlopen.side_effect = [rate_limit_error() for _ in range(4)]
+        delays: list[float] = []
+
+        with self.assertRaisesRegex(RuntimeError, "API returned 429"):
+            HttpTransport(
+                jitter=lambda: 2,
+                sleep=delays.append,
+            ).request("https://api.buildkite.com/v2/builds")
+
+        self.assertEqual(delays, [11, 11, 11])
+        self.assertEqual(urlopen.call_count, 4)
+
+    @patch("run_ci_command.urllib.request.urlopen")
+    def test_http_transport_does_not_retry_permission_error(self, urlopen: Any) -> None:
+        permission_error = urllib.error.HTTPError(
+            "https://api.github.com/repos/vllm-project/vllm/issues/1/comments",
+            403,
+            "Forbidden",
+            {},
+            io.BytesIO(b'{"message":"Resource not accessible by integration"}'),
+        )
+        urlopen.side_effect = permission_error
+        delays: list[float] = []
+
+        with self.assertRaisesRegex(
+            RuntimeError,
+            "Resource not accessible by integration",
+        ):
+            HttpTransport(
+                jitter=lambda: 2,
+                sleep=delays.append,
+            ).request(
+                "https://api.github.com/repos/vllm-project/vllm/issues/1/comments"
+            )
+
+        self.assertEqual(delays, [])
+        self.assertEqual(urlopen.call_count, 1)
+
     def test_only_exact_ci_commands_are_accepted(self) -> None:
         self.assertEqual(parse_command(COMMAND_RUN_CI), COMMAND_RUN_CI)
         self.assertEqual(