From 0176848d90c8fbb9328bc5f185bec8fb0dd89dd1 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Tue, 13 Oct 2026 12:44:39 -0700 Subject: [PATCH 1/3] fix: coordinate REST API rate limits and bound retries --- docs/api/handling_errors.md | 35 ++ runpod/api/_rate_limit.py | 166 ++++++++++ runpod/api/rest.py | 170 +++++++--- runpod/error.py | 14 + tests/test_api/test_rate_limit.py | 510 ++++++++++++++++++++++++++++++ 5 files changed, 858 insertions(+), 37 deletions(-) create mode 100644 runpod/api/_rate_limit.py create mode 100644 tests/test_api/test_rate_limit.py diff --git a/docs/api/handling_errors.md b/docs/api/handling_errors.md index e8b1a28b1..7d5d3517b 100644 --- a/docs/api/handling_errors.md +++ b/docs/api/handling_errors.md @@ -26,3 +26,38 @@ except runpod.error.QueryError as err: print(err.query) print(err.errors) ``` + +## REST rate limits + +REST calls share Runpod's minute, hour, and day quotas across routes within the +process, keyed by API origin and API key. Both synchronous and asynchronous calls +reserve the remaining budget advertised by the `RateLimit` response header and +wait when a window is exhausted. Separate processes and different API keys for the +same account cannot coordinate their local budgets. + +HTTP 429 responses honor `Retry-After` in seconds or HTTP-date format. Missing or +invalid retry hints use capped exponential backoff with jitter. By default, a call +allows three retries and at most 60 seconds of cumulative rate-limit sleeps. +Network failures and other HTTP errors are not retried, since a failed response +can follow a successful resource mutation. + +An exhausted retry or wait budget raises `RateLimitError`, a `QueryError` subclass +with `status_code=429` and a `retry_after` delay in seconds. A successful call that +uses the last available token still returns immediately; subsequent calls wait or +raise if the required delay exceeds their wait budget. + +```python +try: + pods = runpod.get_pods() +except runpod.error.RateLimitError as err: + print(f"Quota exhausted. Retry after {err.retry_after:.1f} seconds.") +``` + +The low-level `runpod.api.rest.run_rest_request` and `run_rest_request_async` +functions accept `max_retries` and `max_rate_limit_wait` keyword arguments. Set +`max_retries=0` to surface the first 429, or `max_rate_limit_wait=0` to fail rather +than sleep for an exhausted quota. `timeout` applies separately to each HTTP +attempt. Async quota waits are cancellable and do not block the event loop. + +These protections apply to the management REST API transport. GraphQL calls and +Serverless job execution use separate transports. diff --git a/runpod/api/_rate_limit.py b/runpod/api/_rate_limit.py new file mode 100644 index 000000000..2f1e0ab36 --- /dev/null +++ b/runpod/api/_rate_limit.py @@ -0,0 +1,166 @@ +"""Shared REST quotas for threads and asynchronous callers.""" + +import hashlib +import math +import random +import re +import threading +import time +from contextlib import contextmanager +from email.utils import parsedate_to_datetime +from typing import Mapping, Optional +from urllib.parse import urlsplit + + +def _retry_after(value: Optional[str]) -> Optional[float]: + if value is None: + return None + try: + delay = float(value) + except ValueError: + try: + delay = parsedate_to_datetime(value).timestamp() - time.time() + except (ValueError, TypeError, OverflowError): + return None + delay = max(0.0, delay) + return delay if math.isfinite(delay) and delay >= 0 else None + + +def _windows(value: str) -> dict[str, tuple[int, float]]: + windows = {} + for member in value.split(","): + parts = member.strip().split(";") + if parts[0] not in {'"minute"', '"hour"', '"day"'}: + continue + parameters = {} + for part in parts[1:]: + match = re.fullmatch(r"\s*([a-z]+)\s*=\s*([0-9]{1,15})\s*", part) + if match: + parameters[match[1]] = int(match[2]) + if "r" in parameters and "t" in parameters: + windows[parts[0]] = (parameters["r"], float(parameters["t"])) + return windows + + +class _Bucket: + def __init__(self): + self.lock = threading.Lock() + self.windows: dict[str, tuple[Optional[int], float]] = {} + self.blocked_until = 0.0 + self.discovered = False + self.in_flight = 0 + self.users = 0 + self.last_used = time.monotonic() + + def acquire(self) -> float: + with self.lock: + now = time.monotonic() + delay = max(0.0, self.blocked_until - now) + if self.in_flight and ( + not self.discovered or (self.blocked_until and not delay) + ): + delay = 0.05 + for name, (remaining, reset_at) in self.windows.items(): + if now >= reset_at: + self.windows[name] = (None, reset_at) + if self.in_flight: + delay = max(delay, 0.05) + elif remaining == 0: + delay = max(delay, reset_at - now) + if delay: + return delay + self.in_flight += 1 + for name, (remaining, reset_at) in self.windows.items(): + if remaining is not None: + self.windows[name] = (max(0, remaining - 1), reset_at) + return 0.0 + + def update( + self, status: int, headers: Mapping[str, str], attempt: int + ) -> Optional[float]: + headers = {key.lower(): value for key, value in headers.items()} + windows = _windows(headers.get("ratelimit", "")) + retry_after = None + if status == 429: + retry_after = _retry_after(headers.get("retry-after")) + if retry_after is None: + retry_after = random.uniform(0.5, 1.0) * 2 ** min(attempt, 6) + with self.lock: + now = time.monotonic() + self.discovered = True + if self.blocked_until <= now: + self.blocked_until = 0.0 + for name, (remaining, reset_after) in windows.items(): + reset_at = now + reset_after + remaining = max(0, remaining - max(0, self.in_flight - 1)) + previous, previous_reset = self.windows.get(name, (None, 0.0)) + if previous is not None and now < previous_reset: + remaining = min(remaining, previous) + reset_at = max(reset_at, previous_reset) + self.windows[name] = (remaining, reset_at) + for name in list(self.windows): + if name not in windows and self.windows[name][1] <= now: + del self.windows[name] + if retry_after is not None: + retry_after = max( + retry_after, + *( + reset - now + for remaining, reset in self.windows.values() + if remaining == 0 + ), + 0.0, + ) + self.blocked_until = max(self.blocked_until, now + retry_after) + return retry_after + + def release(self): + with self.lock: + self.in_flight -= 1 + + +class _RateLimits: + def __init__(self): + self.lock = threading.Lock() + self.buckets: dict[tuple[str, str, bytes], _Bucket] = {} + + @contextmanager + def bucket(self, url: str, api_key: str): + origin = urlsplit(url) + key = ( + origin.scheme.lower(), + origin.netloc.lower(), + hashlib.sha256(api_key.encode()).digest(), + ) + with self.lock: + now = time.monotonic() + for cached_key, bucket in list(self.buckets.items()): + if ( + bucket.users == 0 + and now - bucket.last_used >= 300 + and bucket.blocked_until <= now + and all(reset <= now for _, reset in bucket.windows.values()) + ): + del self.buckets[cached_key] + bucket = self.buckets.setdefault(key, _Bucket()) + bucket.users += 1 + try: + yield bucket + finally: + with self.lock: + bucket.users -= 1 + bucket.last_used = time.monotonic() + + +def _validate_retry_options(max_retries: int, max_rate_limit_wait: float) -> None: + if ( + isinstance(max_retries, bool) + or not isinstance(max_retries, int) + or max_retries < 0 + ): + raise ValueError("max_retries must be a non-negative integer") + if not math.isfinite(max_rate_limit_wait) or max_rate_limit_wait < 0: + raise ValueError("max_rate_limit_wait must be finite and non-negative") + + +rate_limits = _RateLimits() diff --git a/runpod/api/rest.py b/runpod/api/rest.py index 22e519c72..c74ad74a3 100644 --- a/runpod/api/rest.py +++ b/runpod/api/rest.py @@ -1,7 +1,9 @@ """Runpod REST API transport.""" +import asyncio import json as json_module import os +import time from typing import Any, Mapping, Optional import aiohttp @@ -10,6 +12,8 @@ from runpod import error from runpod.user_agent import USER_AGENT +from ._rate_limit import _validate_retry_options, rate_limits + HTTP_STATUS_NO_CONTENT = 204 HTTP_STATUS_BAD_REQUEST = 400 HTTP_STATUS_UNAUTHORIZED = 401 @@ -17,7 +21,9 @@ def _resolve_api_key(api_key: Optional[str]) -> str: - from runpod import api_key as global_api_key # pylint: disable=import-outside-toplevel,cyclic-import + from runpod import ( + api_key as global_api_key, + ) # pylint: disable=import-outside-toplevel,cyclic-import effective_api_key = api_key or global_api_key if not effective_api_key: @@ -40,7 +46,11 @@ def _build_headers(api_key: str) -> dict[str, str]: def _raise_for_error( - status_code: int, method: str, path: str, text: str = "" + status_code: int, + method: str, + path: str, + text: str = "", + retry_after: Optional[float] = None, ) -> None: if status_code == HTTP_STATUS_UNAUTHORIZED: raise error.AuthenticationError( @@ -60,6 +70,14 @@ def _raise_for_error( if not message: message = text or f"Request failed with status {status_code}" + if status_code == 429: + raise error.RateLimitError( + str(message), + f"{method.upper()} {path}", + retry_after=retry_after or 0.0, + errors=payload.get("errors"), + ) + raise error.QueryError( str(message), f"{method.upper()} {path}", @@ -76,22 +94,67 @@ def run_rest_request( params: Optional[Mapping[str, Any]] = None, json: Optional[Mapping[str, Any]] = None, timeout: float = 30, + max_retries: int = 3, + max_rate_limit_wait: float = 60, ) -> Optional[dict[str, Any]]: - """Send an authenticated request to the Runpod REST API.""" - response = requests.request( - method, - _build_url(path), - headers=_build_headers(_resolve_api_key(api_key)), - params=params, - json=json, - timeout=timeout, - ) - if response.status_code >= HTTP_STATUS_BAD_REQUEST: - _raise_for_error(response.status_code, method, path, response.text) - - if response.status_code == HTTP_STATUS_NO_CONTENT or not response.content: - return None - return response.json() + """Send a REST request with bounded 429 retries and shared quota waits. + + ``timeout`` applies to each HTTP attempt. ``max_rate_limit_wait`` bounds + cumulative quota sleeps in seconds; ``max_retries`` counts 429 retries. + Network failures and other HTTP errors are not retried. + """ + _validate_retry_options(max_retries, max_rate_limit_wait) + key = _resolve_api_key(api_key) + url = _build_url(path) + headers = _build_headers(key) + with rate_limits.bucket(url, key) as bucket: + for attempt in range(max_retries + 1): + while delay := bucket.acquire(): + if delay > max_rate_limit_wait: + raise error.RateLimitError( + "Rate limit wait budget exceeded", + f"{method.upper()} {path}", + retry_after=delay, + ) + time.sleep(delay) + max_rate_limit_wait -= delay + try: + response = requests.request( + method, + url, + headers=headers, + params=params, + json=json, + timeout=timeout, + ) + try: + retry_after = bucket.update( + response.status_code, response.headers, attempt + ) + if ( + response.status_code == 429 + and attempt < max_retries + and retry_after <= max_rate_limit_wait + ): + continue + if response.status_code >= HTTP_STATUS_BAD_REQUEST: + _raise_for_error( + response.status_code, + method, + path, + response.text, + retry_after, + ) + if ( + response.status_code == HTTP_STATUS_NO_CONTENT + or not response.content + ): + return None + return response.json() + finally: + response.close() + finally: + bucket.release() async def run_rest_request_async( @@ -102,28 +165,61 @@ async def run_rest_request_async( params: Optional[Mapping[str, Any]] = None, json: Optional[Mapping[str, Any]] = None, timeout: float = 30, + max_retries: int = 3, + max_rate_limit_wait: float = 60, ) -> Optional[dict[str, Any]]: - """Send an authenticated REST request without blocking the event loop.""" - headers = _build_headers(_resolve_api_key(api_key)) + """Send a REST request with cancellable quota waits and bounded 429 retries. + + ``timeout`` applies to each HTTP attempt. ``max_rate_limit_wait`` bounds + cumulative quota sleeps in seconds; ``max_retries`` counts 429 retries. + Network failures and other HTTP errors are not retried. + """ + _validate_retry_options(max_retries, max_rate_limit_wait) + key = _resolve_api_key(api_key) + url = _build_url(path) + headers = _build_headers(key) client_timeout = aiohttp.ClientTimeout(total=timeout) async with aiohttp.ClientSession(timeout=client_timeout) as session: # aiohttp otherwise replays put/delete requests after an ambiguous disconnect. session._retry_connection = False # pylint: disable=protected-access - async with session.request( - method, - _build_url(path), - headers=headers, - params=params, - json=json, - ) as response: - if response.status >= HTTP_STATUS_BAD_REQUEST: - text = ( - "" - if response.status == HTTP_STATUS_UNAUTHORIZED - else await response.text(errors="replace") - ) - _raise_for_error(response.status, method, path, text) - - if response.status == HTTP_STATUS_NO_CONTENT or not await response.read(): - return None - return await response.json(content_type=None) + with rate_limits.bucket(url, key) as bucket: + for attempt in range(max_retries + 1): + while delay := bucket.acquire(): + if delay > max_rate_limit_wait: + raise error.RateLimitError( + "Rate limit wait budget exceeded", + f"{method.upper()} {path}", + retry_after=delay, + ) + await asyncio.sleep(delay) + max_rate_limit_wait -= delay + try: + async with session.request( + method, url, headers=headers, params=params, json=json + ) as response: + retry_after = bucket.update( + response.status, response.headers, attempt + ) + if ( + response.status == 429 + and attempt < max_retries + and retry_after <= max_rate_limit_wait + ): + continue + if response.status >= HTTP_STATUS_BAD_REQUEST: + text = ( + "" + if response.status == HTTP_STATUS_UNAUTHORIZED + else await response.text(errors="replace") + ) + _raise_for_error( + response.status, method, path, text, retry_after + ) + if ( + response.status == HTTP_STATUS_NO_CONTENT + or not await response.read() + ): + return None + return await response.json(content_type=None) + finally: + bucket.release() diff --git a/runpod/error.py b/runpod/error.py index 35a1a8f31..2410edf61 100644 --- a/runpod/error.py +++ b/runpod/error.py @@ -44,3 +44,17 @@ def __init__( self.query = query self.status_code = status_code self.errors = errors or [] + + +class RateLimitError(QueryError): + """A REST quota exceeds the request's retry or wait budget.""" + + def __init__( + self, + message: str, + query: str, + retry_after: float, + errors: Optional[list[str]] = None, + ): + super().__init__(message, query, status_code=429, errors=errors) + self.retry_after = retry_after diff --git a/tests/test_api/test_rate_limit.py b/tests/test_api/test_rate_limit.py new file mode 100644 index 000000000..bab05273c --- /dev/null +++ b/tests/test_api/test_rate_limit.py @@ -0,0 +1,510 @@ +"""REST quota parsing, reservations, and bounded retries.""" + +import asyncio +import json +import threading +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timezone +from email.utils import format_datetime +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from unittest.mock import Mock + +import pytest +import requests + +from runpod.api import _rate_limit, rest +from runpod.error import QueryError, RateLimitError + + +@pytest.fixture +def clock(monkeypatch): + class Clock: + now = 100.0 + + def sleep(self, seconds): + self.now += seconds + + clock = Clock() + monkeypatch.setattr(_rate_limit.time, "monotonic", lambda: clock.now) + return clock + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + (None, None), + ("2", 2), + (" 0.25 ", 0.25), + ("0", 0), + ("-1", None), + ("NaN", None), + ("inf", None), + ("1e309", None), + ("garbage", None), + ("", None), + ], +) +def test_retry_after_seconds(value, expected): + assert _rate_limit._retry_after(value) == expected + + +def test_retry_after_http_date(monkeypatch): + monkeypatch.setattr(_rate_limit.time, "time", lambda: 100) + future = format_datetime(datetime.fromtimestamp(105, timezone.utc), usegmt=True) + past = format_datetime(datetime.fromtimestamp(90, timezone.utc), usegmt=True) + assert _rate_limit._retry_after(future) == 5 + assert _rate_limit._retry_after(past) == 0 + + +def test_parse_runpod_windows(): + assert _rate_limit._windows( + '"minute";r=45;t=12, "hour";t=1812;r=2800, "day";r=49500;t=45012' + ) == {'"minute"': (45, 12), '"hour"': (2800, 1812), '"day"': (49500, 45012)} + + +@pytest.mark.parametrize( + "value", + [ + "", + "garbage", + '"minute";r=0', + '"minute";r=-1;t=30', + '"minute";r=0;t=NaN', + '"minute";r=0;t=inf', + '"minute";r=0;t=999999999999999999999999999', + '"unknown";r=0;t=30', + ], +) +def test_malformed_windows_are_ignored(value): + assert _rate_limit._windows(value) == {} + + +def test_budget_reservations_and_single_reset_probe(clock): + bucket = _rate_limit._Bucket() + assert bucket.acquire() == 0 + bucket.update(200, {"RateLimit": '"minute";r=2;t=10'}, 0) + bucket.release() + assert bucket.acquire() == 0 + assert bucket.acquire() == 0 + assert bucket.acquire() == 10 + bucket.release() + bucket.release() + clock.sleep(10) + assert bucket.acquire() == 0 + assert bucket.acquire() > 0 + bucket.update(200, {"RateLimit": '"minute";r=5;t=60'}, 0) + bucket.release() + assert bucket.acquire() == 0 + bucket.release() + + +def test_longest_exhausted_window_controls_wait(clock): + bucket = _rate_limit._Bucket() + assert bucket.acquire() == 0 + assert ( + bucket.update( + 429, + { + "Retry-After": "1", + "RateLimit": '"minute";r=0;t=10, "hour";r=0;t=100, "day";r=50;t=1000', + }, + 0, + ) + == 100 + ) + bucket.release() + assert bucket.acquire() == 100 + + +def test_out_of_order_success_cannot_restore_spent_quota_or_clear_cooldown(clock): + bucket = _rate_limit._Bucket() + bucket.update(200, {}, 0) + assert bucket.acquire() == 0 + assert bucket.acquire() == 0 + bucket.update(429, {"Retry-After": "10", "RateLimit": '"minute";r=0;t=5'}, 0) + bucket.release() + bucket.update(200, {"RateLimit": '"minute";r=10;t=5'}, 0) + bucket.release() + assert bucket.acquire() == 10 + assert bucket.windows['"minute"'][0] == 0 + + +def test_overlapping_cooldowns_only_extend(clock): + bucket = _rate_limit._Bucket() + bucket.update(429, {"Retry-After": "10"}, 0) + clock.sleep(1) + bucket.update(429, {"Retry-After": "2"}, 0) + assert bucket.acquire() == 9 + bucket.update(429, {"Retry-After": "20"}, 0) + assert bucket.acquire() == 20 + + +def test_headerless_cooldown_allows_one_probe(clock): + bucket = _rate_limit._Bucket() + bucket.update(429, {"Retry-After": "1"}, 0) + clock.sleep(1) + assert bucket.acquire() == 0 + assert bucket.acquire() > 0 + bucket.update(200, {}, 0) + bucket.release() + assert bucket.acquire() == 0 + assert bucket.acquire() == 0 + bucket.release() + bucket.release() + + +def test_expired_windows_without_response_headers_do_not_deadlock(clock): + bucket = _rate_limit._Bucket() + bucket.update(200, {"RateLimit": '"minute";r=0;t=1'}, 0) + clock.sleep(1) + assert bucket.acquire() == 0 + bucket.update(200, {}, 0) + bucket.release() + assert bucket.windows == {} + + +def test_in_flight_requests_are_subtracted_from_learned_budget(clock): + bucket = _rate_limit._Bucket() + bucket.update(200, {}, 0) + assert bucket.acquire() == 0 + assert bucket.acquire() == 0 + bucket.update(200, {"RateLimit": '"minute";r=1;t=10'}, 0) + bucket.release() + assert bucket.acquire() == 10 + bucket.release() + + +def test_initial_discovery_allows_one_request_then_headerless_concurrency(clock): + bucket = _rate_limit._Bucket() + assert bucket.acquire() == 0 + assert bucket.acquire() > 0 + bucket.update(200, {}, 0) + bucket.release() + assert bucket.acquire() == 0 + assert bucket.acquire() == 0 + bucket.release() + bucket.release() + + +def test_failed_discovery_releases_probe(clock): + bucket = _rate_limit._Bucket() + assert bucket.acquire() == 0 + bucket.release() + assert bucket.acquire() == 0 + bucket.release() + + +def test_atomic_reservations_across_threads(): + bucket = _rate_limit._Bucket() + bucket.update(200, {"RateLimit": '"minute";r=3;t=60'}, 0) + with ThreadPoolExecutor(max_workers=10) as executor: + delays = list(executor.map(lambda _: bucket.acquire(), range(20))) + assert delays.count(0) == 3 + assert bucket.in_flight == 3 + for _ in range(3): + bucket.release() + + +def test_registry_isolates_credentials_and_origins_but_shares_routes(clock): + limits = _rate_limit._RateLimits() + with limits.bucket("https://example.com/v2/pods", "key") as first: + with limits.bucket("https://example.com/v2/templates?all=1", "key") as same: + assert same is first + with limits.bucket("https://example.com/v2/pods", "other") as other: + assert other is not first + with limits.bucket("http://example.com/v2/pods", "key") as other: + assert other is not first + with limits.bucket("https://other.example/v2/pods", "key") as other: + assert other is not first + assert all(isinstance(key[2], bytes) for key in limits.buckets) + assert "key" not in repr(limits.buckets) + + +def test_registry_retains_active_cooldowns_and_evicts_idle_state(clock): + limits = _rate_limit._RateLimits() + with limits.bucket("https://example.com", "key") as active: + with limits.bucket("https://example.com", "limited") as limited: + limited.update(429, {"Retry-After": "500"}, 0) + with limits.bucket("https://example.com", "idle") as idle: + pass + clock.sleep(301) + with limits.bucket("https://example.com", "new"): + assert active in limits.buckets.values() + assert limited in limits.buckets.values() + assert idle not in limits.buckets.values() + clock.sleep(300) + with limits.bucket("https://example.com", "new"): + assert limited not in limits.buckets.values() + + +@pytest.mark.parametrize("attempt", [0, 1, 5, 1000]) +def test_fallback_is_positive_jittered_and_capped(attempt, clock, monkeypatch): + monkeypatch.setattr(_rate_limit.random, "uniform", lambda low, high: 0.75) + bucket = _rate_limit._Bucket() + assert bucket.update(429, {"Retry-After": "NaN"}, attempt) == ( + 0.75 * 2 ** min(attempt, 6) + ) + + +@pytest.fixture +def quota_server(monkeypatch): + state = {"requests": [], "responses": [], "lock": threading.Lock()} + + class Handler(BaseHTTPRequestHandler): + def do_POST(self): + body = self.rfile.read(int(self.headers.get("Content-Length", 0))) + with state["lock"]: + state["requests"].append( + (self.path, body, self.headers["Authorization"]) + ) + status, headers, payload = ( + state["responses"].pop(0) + if state["responses"] + else (200, {}, {"id": "pod-1"}) + ) + body = json.dumps(payload).encode() + self.send_response(status) + self.send_header("Content-Length", str(len(body))) + self.send_header("Content-Type", "application/json") + for key, value in headers.items(): + self.send_header(key, value) + self.end_headers() + self.wfile.write(body) + + do_GET = do_PUT = do_DELETE = do_POST + + def log_message(self, *args): + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + monkeypatch.setenv("RUNPOD_API_BASE_URL", f"http://127.0.0.1:{server.server_port}") + monkeypatch.setattr(rest, "rate_limits", _rate_limit._RateLimits()) + try: + yield state + finally: + server.shutdown() + server.server_close() + thread.join() + + +@pytest.fixture(params=["sync", "async"]) +def request_rest(request): + if request.param == "sync": + return rest.run_rest_request + + def run(*args, **kwargs): + return asyncio.run(rest.run_rest_request_async(*args, **kwargs)) + + return run + + +@pytest.mark.parametrize("method", ["GET", "POST", "PUT", "DELETE"]) +def test_retry_preserves_request(quota_server, request_rest, method): + quota_server["responses"] = [ + (429, {"Retry-After": "0.01"}, {"detail": "slow down"}) + ] + assert request_rest( + method, "/v2/pods", params={"all": "true"}, json={"name": "test"}, api_key="key" + ) == {"id": "pod-1"} + assert len(quota_server["requests"]) == 2 + assert quota_server["requests"][0] == quota_server["requests"][1] + assert quota_server["requests"][0][0] == "/v2/pods?all=true" + + +def test_retries_are_bounded_and_preserve_error_details(quota_server, request_rest): + quota_server["responses"] = [ + (429, {"Retry-After": "0"}, {"detail": "slow down", "errors": ["quota"]}) + ] * 5 + with pytest.raises(RateLimitError) as raised: + request_rest("POST", "/v2/pods", api_key="key") + assert len(quota_server["requests"]) == 4 + assert isinstance(raised.value, QueryError) + assert raised.value.status_code == 429 + assert raised.value.retry_after == 0 + assert raised.value.errors == ["quota"] + assert raised.value.query == "POST /v2/pods" + assert str(raised.value) == "slow down" + + +def test_zero_retries_surfaces_first_rejection(quota_server, request_rest): + quota_server["responses"] = [(429, {"Retry-After": "1"}, {"detail": "slow down"})] + with pytest.raises(RateLimitError): + request_rest("GET", "/v2/pods", api_key="key", max_retries=0) + assert len(quota_server["requests"]) == 1 + + +def test_long_wait_fails_without_sleep_or_retry(quota_server, request_rest): + quota_server["responses"] = [ + (429, {"Retry-After": "3600"}, {"detail": "hour quota"}) + ] + with pytest.raises(RateLimitError) as raised: + request_rest("POST", "/v2/pods", api_key="key", max_rate_limit_wait=1) + assert raised.value.retry_after == 3600 + assert str(raised.value) == "hour quota" + assert len(quota_server["requests"]) == 1 + + +def test_wait_budget_is_cumulative_across_retries(quota_server, request_rest): + quota_server["responses"] = [ + (429, {"Retry-After": "0.04"}, {"detail": "quota"}) + ] * 3 + with pytest.raises(RateLimitError) as raised: + request_rest("GET", "/v2/pods", api_key="key", max_rate_limit_wait=0.06) + assert raised.value.retry_after == 0.04 + assert len(quota_server["requests"]) == 2 + + +@pytest.mark.parametrize("retry_header", [{}, {"Retry-After": "invalid"}]) +def test_missing_or_malformed_retry_hint_uses_backoff( + quota_server, request_rest, monkeypatch, retry_header +): + monkeypatch.setattr(_rate_limit.random, "uniform", lambda low, high: 0.01) + quota_server["responses"] = [(429, retry_header, {"detail": "quota"})] + assert request_rest("GET", "/v2/pods", api_key="key") == {"id": "pod-1"} + assert len(quota_server["requests"]) == 2 + + +def test_sync_response_is_closed_before_sleep(monkeypatch, clock): + first = requests.Response() + first.status_code = 429 + first.headers["Retry-After"] = "1" + first._content = b"quota" + first._content_consumed = True + first.close = Mock() + second = requests.Response() + second.status_code = 200 + second._content = b'{"id":"pod-1"}' + second._content_consumed = True + second.close = Mock() + monkeypatch.setattr(rest, "rate_limits", _rate_limit._RateLimits()) + monkeypatch.setattr(rest.requests, "request", Mock(side_effect=[first, second])) + + def sleep(seconds): + first.close.assert_called_once() + clock.sleep(seconds) + + monkeypatch.setattr(rest.time, "sleep", sleep) + assert rest.run_rest_request("GET", "/v2/pods", api_key="key") == {"id": "pod-1"} + second.close.assert_called_once() + + +def test_successful_last_token_blocks_other_routes(quota_server, request_rest): + quota_server["responses"] = [ + (200, {"RateLimit": '"minute";r=0;t=60'}, {"id": "pod-1"}) + ] + assert request_rest("GET", "/v2/pods", api_key="key") == {"id": "pod-1"} + with pytest.raises(RateLimitError) as raised: + request_rest("GET", "/v2/templates", api_key="key", max_rate_limit_wait=0) + assert 0 < raised.value.retry_after <= 60 + assert len(quota_server["requests"]) == 1 + assert request_rest("GET", "/v2/templates", api_key="other") == {"id": "pod-1"} + + +def test_sync_and_async_share_cooldown(quota_server): + quota_server["responses"] = [(429, {"Retry-After": "60"}, {"detail": "quota"})] + with pytest.raises(RateLimitError): + rest.run_rest_request("GET", "/v2/pods", api_key="key", max_retries=0) + with pytest.raises(RateLimitError): + asyncio.run( + rest.run_rest_request_async( + "GET", "/v2/templates", api_key="key", max_rate_limit_wait=0 + ) + ) + assert len(quota_server["requests"]) == 1 + + +@pytest.mark.parametrize("status", [400, 401, 403, 404, 500, 503]) +def test_other_errors_are_not_retried(quota_server, request_rest, status): + from runpod.error import RunPodError + + quota_server["responses"] = [(status, {"Retry-After": "0"}, {"detail": "failure"})] + with pytest.raises(RunPodError): + request_rest("POST", "/v2/pods", api_key="key") + assert len(quota_server["requests"]) == 1 + + +@pytest.mark.parametrize( + "options", + [ + {"max_retries": -1}, + {"max_retries": 1.5}, + {"max_retries": True}, + {"max_rate_limit_wait": -1}, + {"max_rate_limit_wait": float("nan")}, + {"max_rate_limit_wait": float("inf")}, + ], +) +def test_invalid_retry_options_fail_before_network(request_rest, options): + with pytest.raises(ValueError): + request_rest("GET", "/v2/pods", api_key="key", **options) + + +@pytest.mark.asyncio +async def test_cancellation_during_wait_does_not_hold_reservation( + quota_server, monkeypatch +): + quota_server["responses"] = [(429, {"Retry-After": "60"}, {"detail": "quota"})] + sleeping = asyncio.Event() + original_sleep = asyncio.sleep + + async def sleep(delay): + sleeping.set() + await original_sleep(delay) + + monkeypatch.setattr(rest.asyncio, "sleep", sleep) + pending = asyncio.create_task( + rest.run_rest_request_async( + "POST", "/v2/pods", api_key="key", max_rate_limit_wait=120 + ) + ) + try: + await asyncio.wait_for(sleeping.wait(), 2) + bucket = next(iter(rest.rate_limits.buckets.values())) + assert bucket.in_flight == 0 + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + assert bucket.users == 0 + assert bucket.in_flight == 0 + assert bucket.blocked_until > _rate_limit.time.monotonic() + assert len(quota_server["requests"]) == 1 + assert await rest.run_rest_request_async( + "GET", "/v2/pods", api_key="other" + ) == {"id": "pod-1"} + finally: + pending.cancel() + await asyncio.gather(pending, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_shared_wait_is_nonblocking_and_refreshes_quota(quota_server): + quota_server["responses"] = [ + (200, {"RateLimit": '"minute";r=0;t=1'}, {"id": "pod-1"}) + ] + await rest.run_rest_request_async("GET", "/v2/pods", api_key="key") + pending = [ + asyncio.create_task( + rest.run_rest_request_async("GET", "/v2/templates", api_key="key") + ) + for _ in range(3) + ] + try: + await asyncio.sleep(0.02) + assert len(quota_server["requests"]) == 1 + assert not any(task.done() for task in pending) + assert await rest.run_rest_request_async( + "GET", "/v2/pods", api_key="other" + ) == {"id": "pod-1"} + quota_server["responses"] = [ + (200, {"RateLimit": '"minute";r=10;t=60'}, {"id": "pod-1"}) + ] + assert ( + await asyncio.wait_for(asyncio.gather(*pending), 3) == [{"id": "pod-1"}] * 3 + ) + assert len(quota_server["requests"]) == 5 + finally: + for task in pending: + task.cancel() + await asyncio.gather(*pending, return_exceptions=True) From 2a8cd83d010ece56e8631b76507e31c9ec18bc29 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Tue, 13 Oct 2026 12:51:26 -0700 Subject: [PATCH 2/3] fix: harden rate-limit credential fingerprints and fixtures --- runpod/api/_rate_limit.py | 6 ++++-- tests/test_api/test_rate_limit.py | 15 +++++++++++++-- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/runpod/api/_rate_limit.py b/runpod/api/_rate_limit.py index 2f1e0ab36..0ba9bebd5 100644 --- a/runpod/api/_rate_limit.py +++ b/runpod/api/_rate_limit.py @@ -1,7 +1,8 @@ """Shared REST quotas for threads and asynchronous callers.""" -import hashlib +import hmac import math +import os import random import re import threading @@ -122,6 +123,7 @@ def release(self): class _RateLimits: def __init__(self): self.lock = threading.Lock() + self.fingerprint_key = os.urandom(32) self.buckets: dict[tuple[str, str, bytes], _Bucket] = {} @contextmanager @@ -130,7 +132,7 @@ def bucket(self, url: str, api_key: str): key = ( origin.scheme.lower(), origin.netloc.lower(), - hashlib.sha256(api_key.encode()).digest(), + hmac.digest(self.fingerprint_key, api_key.encode(), "sha256"), ) with self.lock: now = time.monotonic() diff --git a/tests/test_api/test_rate_limit.py b/tests/test_api/test_rate_limit.py index bab05273c..78c06311b 100644 --- a/tests/test_api/test_rate_limit.py +++ b/tests/test_api/test_rate_limit.py @@ -237,6 +237,14 @@ def test_registry_retains_active_cooldowns_and_evicts_idle_state(clock): assert limited not in limits.buckets.values() +def test_credential_fingerprints_are_private_to_each_registry(): + first = _rate_limit._RateLimits() + second = _rate_limit._RateLimits() + with first.bucket("https://example.com", "key"): + with second.bucket("https://example.com", "key"): + assert next(iter(first.buckets)) != next(iter(second.buckets)) + + @pytest.mark.parametrize("attempt", [0, 1, 5, 1000]) def test_fallback_is_positive_jittered_and_capped(attempt, clock, monkeypatch): monkeypatch.setattr(_rate_limit.random, "uniform", lambda low, high: 0.75) @@ -267,7 +275,10 @@ def do_POST(self): self.send_header("Content-Length", str(len(body))) self.send_header("Content-Type", "application/json") for key, value in headers.items(): - self.send_header(key, value) + self.send_header( + key.replace("\r", "").replace("\n", ""), + value.replace("\r", "").replace("\n", ""), + ) self.end_headers() self.wfile.write(body) @@ -465,7 +476,7 @@ async def sleep(delay): assert bucket.in_flight == 0 pending.cancel() with pytest.raises(asyncio.CancelledError): - await pending + await asyncio.wait_for(pending, 1) assert bucket.users == 0 assert bucket.in_flight == 0 assert bucket.blocked_until > _rate_limit.time.monotonic() From 3d1c1a25b19989be184959f5d0d3d0bb9bcf88fa Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Wed, 14 Oct 2026 18:58:57 -0700 Subject: [PATCH 3/3] fix: reconcile REST quota snapshots and retry budgets --- docs/api/handling_errors.md | 6 + runpod/api/_rate_limit.py | 122 +++++++++++++---- runpod/api/rest.py | 50 +++---- tests/test_api/test_rate_limit.py | 217 ++++++++++++++++++++++++++++++ 4 files changed, 340 insertions(+), 55 deletions(-) diff --git a/docs/api/handling_errors.md b/docs/api/handling_errors.md index 7d5d3517b..5957b8224 100644 --- a/docs/api/handling_errors.md +++ b/docs/api/handling_errors.md @@ -35,6 +35,12 @@ reserve the remaining budget advertised by the `RateLimit` response header and wait when a window is exhausted. Separate processes and different API keys for the same account cannot coordinate their local budgets. +Quota estimates account for outstanding requests separately from response +snapshots. A window refreshes when the response timing establishes a later reset; +out-of-order snapshots within a window cannot restore spent quota. Discovery and +reset probes coordinate concurrent callers without consuming their rate-limit +wait budgets. + HTTP 429 responses honor `Retry-After` in seconds or HTTP-date format. Missing or invalid retry hints use capped exponential backoff with jitter. By default, a call allows three retries and at most 60 seconds of cumulative rate-limit sleeps. diff --git a/runpod/api/_rate_limit.py b/runpod/api/_rate_limit.py index 0ba9bebd5..527f37070 100644 --- a/runpod/api/_rate_limit.py +++ b/runpod/api/_rate_limit.py @@ -8,10 +8,13 @@ import threading import time from contextlib import contextmanager +from datetime import timezone from email.utils import parsedate_to_datetime from typing import Mapping, Optional from urllib.parse import urlsplit +from runpod.error import RateLimitError + def _retry_after(value: Optional[str]) -> Optional[float]: if value is None: @@ -20,7 +23,10 @@ def _retry_after(value: Optional[str]) -> Optional[float]: delay = float(value) except ValueError: try: - delay = parsedate_to_datetime(value).timestamp() - time.time() + date = parsedate_to_datetime(value) + if date.tzinfo is None: + date = date.replace(tzinfo=timezone.utc) + delay = date.timestamp() - time.time() except (ValueError, TypeError, OverflowError): return None delay = max(0.0, delay) @@ -43,44 +49,80 @@ def _windows(value: str) -> dict[str, tuple[int, float]]: return windows +class _RetryBudget: + def __init__(self, max_retries: int, max_wait: float, query: str): + self.max_retries = max_retries + self.remaining_wait = max_wait + self.query = query + + def wait(self, delay: float) -> None: + if delay > self.remaining_wait: + raise RateLimitError( + "Rate limit wait budget exceeded", self.query, retry_after=delay + ) + self.remaining_wait -= delay + + def can_retry(self, attempt: int, retry_after: Optional[float]) -> bool: + return ( + retry_after is not None + and attempt < self.max_retries + and retry_after <= self.remaining_wait + ) + + class _Bucket: def __init__(self): self.lock = threading.Lock() self.windows: dict[str, tuple[Optional[int], float]] = {} + self.estimates: dict[str, int] = {} + self.reset_lower_bounds: dict[str, float] = {} self.blocked_until = 0.0 self.discovered = False self.in_flight = 0 self.users = 0 self.last_used = time.monotonic() - def acquire(self) -> float: + def acquire(self, budget: Optional[_RetryBudget] = None) -> float: with self.lock: now = time.monotonic() delay = max(0.0, self.blocked_until - now) - if self.in_flight and ( + probing = self.in_flight and ( not self.discovered or (self.blocked_until and not delay) - ): - delay = 0.05 + ) for name, (remaining, reset_at) in self.windows.items(): if now >= reset_at: self.windows[name] = (None, reset_at) if self.in_flight: - delay = max(delay, 0.05) + probing = True elif remaining == 0: delay = max(delay, reset_at - now) if delay: + if budget is not None: + budget.wait(delay) return delay + if probing: + return 0.05 self.in_flight += 1 for name, (remaining, reset_at) in self.windows.items(): if remaining is not None: self.windows[name] = (max(0, remaining - 1), reset_at) + self.estimates[name] = max(0, self.estimates[name] - 1) return 0.0 def update( - self, status: int, headers: Mapping[str, str], attempt: int + self, + status: int, + headers: Mapping[str, str], + attempt: int, + *, + started_at: Optional[float] = None, ) -> Optional[float]: + windows = _windows( + ", ".join( + value for key, value in headers.items() if key.lower() == "ratelimit" + ) + ) headers = {key.lower(): value for key, value in headers.items()} - windows = _windows(headers.get("ratelimit", "")) retry_after = None if status == 429: retry_after = _retry_after(headers.get("retry-after")) @@ -93,22 +135,40 @@ def update( self.blocked_until = 0.0 for name, (remaining, reset_after) in windows.items(): reset_at = now + reset_after - remaining = max(0, remaining - max(0, self.in_flight - 1)) + reset_lower = (now if started_at is None else started_at) + max( + 0.0, reset_after - 1.0 + ) previous, previous_reset = self.windows.get(name, (None, 0.0)) - if previous is not None and now < previous_reset: - remaining = min(remaining, previous) - reset_at = max(reset_at, previous_reset) - self.windows[name] = (remaining, reset_at) + previous_lower = self.reset_lower_bounds.get(name, 0.0) + if reset_at < previous_lower: + continue + if ( + previous is not None + and now < previous_reset + and reset_lower <= previous_reset + ): + remaining = min(remaining, self.estimates[name]) + reset_at = min(reset_at, previous_reset) + reset_lower = max(reset_lower, previous_lower) + self.estimates[name] = remaining + self.reset_lower_bounds[name] = reset_lower + self.windows[name] = ( + max(0, remaining - max(0, self.in_flight - 1)), + reset_at, + ) for name in list(self.windows): if name not in windows and self.windows[name][1] <= now: del self.windows[name] + del self.estimates[name] + del self.reset_lower_bounds[name] if retry_after is not None: retry_after = max( retry_after, + self.blocked_until - now, *( reset - now - for remaining, reset in self.windows.values() - if remaining == 0 + for name, (_, reset) in self.windows.items() + if self.estimates[name] == 0 ), 0.0, ) @@ -118,6 +178,12 @@ def update( def release(self): with self.lock: self.in_flight -= 1 + for name, (remaining, reset_at) in self.windows.items(): + if remaining is not None: + self.windows[name] = ( + max(0, self.estimates[name] - self.in_flight), + reset_at, + ) class _RateLimits: @@ -125,6 +191,7 @@ def __init__(self): self.lock = threading.Lock() self.fingerprint_key = os.urandom(32) self.buckets: dict[tuple[str, str, bytes], _Bucket] = {} + self.next_cleanup = 0.0 @contextmanager def bucket(self, url: str, api_key: str): @@ -136,14 +203,16 @@ def bucket(self, url: str, api_key: str): ) with self.lock: now = time.monotonic() - for cached_key, bucket in list(self.buckets.items()): - if ( - bucket.users == 0 - and now - bucket.last_used >= 300 - and bucket.blocked_until <= now - and all(reset <= now for _, reset in bucket.windows.values()) - ): - del self.buckets[cached_key] + if now >= self.next_cleanup: + for cached_key, bucket in list(self.buckets.items()): + if bucket.users or now - bucket.last_used < 300: + continue + with bucket.lock: + if bucket.blocked_until <= now and all( + reset <= now for _, reset in bucket.windows.values() + ): + del self.buckets[cached_key] + self.next_cleanup = now + 60 bucket = self.buckets.setdefault(key, _Bucket()) bucket.users += 1 try: @@ -161,7 +230,12 @@ def _validate_retry_options(max_retries: int, max_rate_limit_wait: float) -> Non or max_retries < 0 ): raise ValueError("max_retries must be a non-negative integer") - if not math.isfinite(max_rate_limit_wait) or max_rate_limit_wait < 0: + if ( + isinstance(max_rate_limit_wait, bool) + or not isinstance(max_rate_limit_wait, (int, float)) + or not math.isfinite(max_rate_limit_wait) + or max_rate_limit_wait < 0 + ): raise ValueError("max_rate_limit_wait must be finite and non-negative") diff --git a/runpod/api/rest.py b/runpod/api/rest.py index c74ad74a3..b0c9ced43 100644 --- a/runpod/api/rest.py +++ b/runpod/api/rest.py @@ -12,7 +12,7 @@ from runpod import error from runpod.user_agent import USER_AGENT -from ._rate_limit import _validate_retry_options, rate_limits +from ._rate_limit import _RetryBudget, _validate_retry_options, rate_limits HTTP_STATUS_NO_CONTENT = 204 HTTP_STATUS_BAD_REQUEST = 400 @@ -21,9 +21,9 @@ def _resolve_api_key(api_key: Optional[str]) -> str: - from runpod import ( + from runpod import ( # pylint: disable=import-outside-toplevel,cyclic-import api_key as global_api_key, - ) # pylint: disable=import-outside-toplevel,cyclic-import + ) effective_api_key = api_key or global_api_key if not effective_api_key: @@ -107,18 +107,13 @@ def run_rest_request( key = _resolve_api_key(api_key) url = _build_url(path) headers = _build_headers(key) + budget = _RetryBudget(max_retries, max_rate_limit_wait, f"{method.upper()} {path}") with rate_limits.bucket(url, key) as bucket: for attempt in range(max_retries + 1): - while delay := bucket.acquire(): - if delay > max_rate_limit_wait: - raise error.RateLimitError( - "Rate limit wait budget exceeded", - f"{method.upper()} {path}", - retry_after=delay, - ) + while delay := bucket.acquire(budget): time.sleep(delay) - max_rate_limit_wait -= delay try: + started_at = time.monotonic() response = requests.request( method, url, @@ -129,13 +124,12 @@ def run_rest_request( ) try: retry_after = bucket.update( - response.status_code, response.headers, attempt + response.status_code, + response.headers, + attempt, + started_at=started_at, ) - if ( - response.status_code == 429 - and attempt < max_retries - and retry_after <= max_rate_limit_wait - ): + if budget.can_retry(attempt, retry_after): continue if response.status_code >= HTTP_STATUS_BAD_REQUEST: _raise_for_error( @@ -178,33 +172,27 @@ async def run_rest_request_async( key = _resolve_api_key(api_key) url = _build_url(path) headers = _build_headers(key) + budget = _RetryBudget(max_retries, max_rate_limit_wait, f"{method.upper()} {path}") client_timeout = aiohttp.ClientTimeout(total=timeout) async with aiohttp.ClientSession(timeout=client_timeout) as session: # aiohttp otherwise replays put/delete requests after an ambiguous disconnect. session._retry_connection = False # pylint: disable=protected-access with rate_limits.bucket(url, key) as bucket: for attempt in range(max_retries + 1): - while delay := bucket.acquire(): - if delay > max_rate_limit_wait: - raise error.RateLimitError( - "Rate limit wait budget exceeded", - f"{method.upper()} {path}", - retry_after=delay, - ) + while delay := bucket.acquire(budget): await asyncio.sleep(delay) - max_rate_limit_wait -= delay try: + started_at = time.monotonic() async with session.request( method, url, headers=headers, params=params, json=json ) as response: retry_after = bucket.update( - response.status, response.headers, attempt + response.status, + response.headers, + attempt, + started_at=started_at, ) - if ( - response.status == 429 - and attempt < max_retries - and retry_after <= max_rate_limit_wait - ): + if budget.can_retry(attempt, retry_after): continue if response.status >= HTTP_STATUS_BAD_REQUEST: text = ( diff --git a/tests/test_api/test_rate_limit.py b/tests/test_api/test_rate_limit.py index 78c06311b..ffac0df7a 100644 --- a/tests/test_api/test_rate_limit.py +++ b/tests/test_api/test_rate_limit.py @@ -2,7 +2,9 @@ import asyncio import json +import os import threading +import time from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timezone from email.utils import format_datetime @@ -11,6 +13,7 @@ import pytest import requests +from multidict import CIMultiDict from runpod.api import _rate_limit, rest from runpod.error import QueryError, RateLimitError @@ -56,6 +59,121 @@ def test_retry_after_http_date(monkeypatch): assert _rate_limit._retry_after(past) == 0 +@pytest.mark.skipif(not hasattr(time, "tzset"), reason="requires time.tzset") +@pytest.mark.parametrize("zone", ["EST5", "JST-9"]) +@pytest.mark.parametrize( + "value", ["Thu Jan 1 00:01:45 1970", "Thu, 01 Jan 1970 00:01:45 -0000"] +) +def test_naive_retry_dates_are_utc(monkeypatch, zone, value): + previous = os.environ.get("TZ") + try: + os.environ["TZ"] = zone + time.tzset() + monkeypatch.setattr(_rate_limit.time, "time", lambda: 100) + assert _rate_limit._retry_after(value) == 5 + finally: + if previous is None: + os.environ.pop("TZ", None) + else: + os.environ["TZ"] = previous + time.tzset() + + +@pytest.mark.parametrize("remaining", [(1, 2), (2, 1)]) +def test_concurrent_snapshots_preserve_available_quota(clock, remaining): + bucket = _rate_limit._Bucket() + bucket.update(200, {"RateLimit": '"minute";r=3;t=60'}, 0) + assert bucket.acquire() == 0 + assert bucket.acquire() == 0 + for value in remaining: + bucket.update(200, {"RateLimit": f'"minute";r={value};t=60'}, 0) + bucket.release() + assert bucket.windows['"minute"'][0] == 1 + assert bucket.acquire() == 0 + assert bucket.acquire() == 60 + bucket.release() + + +def test_reservations_between_snapshots_remain_spent(clock): + bucket = _rate_limit._Bucket() + bucket.update(200, {"RateLimit": '"minute";r=4;t=60'}, 0) + assert bucket.acquire() == 0 + assert bucket.acquire() == 0 + bucket.update(200, {"RateLimit": '"minute";r=3;t=60'}, 0) + bucket.release() + assert bucket.acquire() == 0 + bucket.update(200, {"RateLimit": '"minute";r=2;t=60'}, 0) + bucket.release() + bucket.update(200, {"RateLimit": '"minute";r=1;t=60'}, 0) + bucket.release() + assert bucket.windows['"minute"'][0] == 1 + + +@pytest.mark.parametrize("window, period", [("minute", 60), ("day", 86400)]) +def test_verified_rollover_refreshes_quota_before_local_deadline(clock, window, period): + bucket = _rate_limit._Bucket() + bucket.update(200, {"RateLimit": f'"{window}";r=1;t=1'}, 0) + clock.sleep(0.5) + assert bucket.acquire() == 0 + bucket.update( + 200, + {"RateLimit": f'"{window}";r=59;t={period}'}, + 0, + started_at=clock.now, + ) + bucket.release() + assert bucket.windows[f'"{window}"'] == (59, clock.now + period) + assert bucket.acquire() == 0 + bucket.release() + + +def test_delayed_snapshot_cannot_impersonate_rollover(clock): + bucket = _rate_limit._Bucket() + started_at = clock.now + bucket.update(200, {"RateLimit": '"minute";r=3;t=60'}, 0) + assert bucket.acquire() == 0 + assert bucket.acquire() == 0 + clock.sleep(5) + bucket.update(200, {"RateLimit": '"minute";r=1;t=55'}, 0, started_at=started_at) + bucket.release() + bucket.update(200, {"RateLimit": '"minute";r=2;t=60'}, 0, started_at=started_at) + bucket.release() + assert bucket.windows['"minute"'] == (1, 160) + + +def test_previous_window_snapshot_cannot_overwrite_verified_rollover(clock): + bucket = _rate_limit._Bucket() + bucket.update(200, {"RateLimit": '"minute";r=2;t=1'}, 0) + assert bucket.acquire() == 0 + assert bucket.acquire() == 0 + bucket.update(200, {"RateLimit": '"minute";r=59;t=60'}, 0, started_at=clock.now) + bucket.release() + bucket.update(200, {"RateLimit": '"minute";r=0;t=1'}, 0, started_at=clock.now) + bucket.release() + assert bucket.windows['"minute"'] == (59, 160) + + +def test_repeated_ratelimit_headers_preserve_all_windows(clock): + bucket = _rate_limit._Bucket() + bucket.update( + 200, + CIMultiDict( + [("RateLimit", '"minute";r=0;t=30'), ("ratelimit", '"day";r=500;t=80000')] + ), + 0, + ) + assert set(bucket.windows) == {'"minute"', '"day"'} + assert bucket.acquire() == 30 + + +def test_retry_delay_includes_shared_cooldown(clock): + bucket = _rate_limit._Bucket() + assert bucket.update(429, {"Retry-After": "60"}, 0) == 60 + clock.sleep(1) + assert bucket.update(429, {"Retry-After": "1"}, 0) == 59 + assert bucket.acquire() == 59 + + def test_parse_runpod_windows(): assert _rate_limit._windows( '"minute";r=45;t=12, "hour";t=1812;r=2800, "day";r=49500;t=45012' @@ -413,6 +531,47 @@ def test_successful_last_token_blocks_other_routes(quota_server, request_rest): assert request_rest("GET", "/v2/templates", api_key="other") == {"id": "pod-1"} +def test_repeated_headers_block_exhausted_window_in_both_transports( + quota_server, request_rest +): + quota_server["responses"] = [ + ( + 200, + CIMultiDict( + [ + ("RateLimit", '"minute";r=0;t=30'), + ("RateLimit", '"day";r=500;t=80000'), + ] + ), + {"id": "pod-1"}, + ) + ] + assert request_rest("GET", "/v2/pods", api_key="key") == {"id": "pod-1"} + with pytest.raises(RateLimitError) as raised: + request_rest("GET", "/v2/templates", api_key="key", max_rate_limit_wait=0) + assert 0 < raised.value.retry_after <= 30 + assert len(quota_server["requests"]) == 1 + + +def test_retry_error_reports_effective_shared_cooldown( + quota_server, request_rest, monkeypatch, clock +): + quota_server["responses"] = [(429, {"Retry-After": "1"}, {"detail": "quota"})] + with rest.rate_limits.bucket(rest._build_url("/v2/pods"), "key") as bucket: + update = bucket.update + + def overlapping_update(*args, **kwargs): + update(429, {"Retry-After": "60"}, 0) + clock.sleep(1) + return update(*args, **kwargs) + + monkeypatch.setattr(bucket, "update", overlapping_update) + with pytest.raises(RateLimitError) as raised: + request_rest("GET", "/v2/pods", api_key="key", max_retries=0) + assert raised.value.retry_after == 59 + assert len(quota_server["requests"]) == 1 + + def test_sync_and_async_share_cooldown(quota_server): quota_server["responses"] = [(429, {"Retry-After": "60"}, {"detail": "quota"})] with pytest.raises(RateLimitError): @@ -452,6 +611,64 @@ def test_invalid_retry_options_fail_before_network(request_rest, options): request_rest("GET", "/v2/pods", api_key="key", **options) +@pytest.mark.parametrize("value", [None, True, False, "60", [], {}]) +def test_invalid_wait_types_raise_value_error(request_rest, value): + with pytest.raises(ValueError, match="max_rate_limit_wait"): + request_rest("GET", "/v2/pods", api_key="key", max_rate_limit_wait=value) + + +@pytest.mark.parametrize("state", ["discovery", "window", "cooldown"]) +def test_coordination_waits_do_not_consume_budget( + quota_server, request_rest, monkeypatch, clock, state +): + with rest.rate_limits.bucket(rest._build_url("/v2/pods"), "key") as bucket: + if state == "window": + bucket.update(200, {"RateLimit": '"minute";r=0;t=1'}, 0) + elif state == "cooldown": + bucket.update(429, {"Retry-After": "1"}, 0) + clock.sleep(1) + assert bucket.acquire() == 0 + sleeps = [] + + def sleep(delay): + sleeps.append(delay) + bucket.update(200, {}, 0) + bucket.release() + clock.sleep(delay) + + async def async_sleep(delay): + sleep(delay) + + monkeypatch.setattr(rest.time, "sleep", sleep) + monkeypatch.setattr(rest.asyncio, "sleep", async_sleep) + assert request_rest( + "GET", "/v2/pods", api_key="key", max_rate_limit_wait=0 + ) == {"id": "pod-1"} + assert sleeps == [0.05] + assert bucket.in_flight == 0 + assert len(quota_server["requests"]) == 1 + + +def test_failed_request_keeps_its_quota_reservation(clock): + bucket = _rate_limit._Bucket() + bucket.update(200, {"RateLimit": '"minute";r=1;t=60'}, 0) + assert bucket.acquire() == 0 + bucket.release() + assert bucket.acquire() == 60 + + +def test_registry_cleanup_is_periodic(clock): + limits = _rate_limit._RateLimits() + with limits.bucket("https://example.com", "key"): + deadline = limits.next_cleanup + clock.sleep(1) + with limits.bucket("https://example.com", "other"): + assert limits.next_cleanup == deadline + clock.sleep(60) + with limits.bucket("https://example.com", "key"): + assert limits.next_cleanup > deadline + + @pytest.mark.asyncio async def test_cancellation_during_wait_does_not_hold_reservation( quota_server, monkeypatch