From 746089eceab6825df287838a40552c4ce4300aac Mon Sep 17 00:00:00 2001 From: Alishahryar1 Date: Thu, 29 Jan 2026 23:34:41 -0800 Subject: [PATCH] Migrated provider ratelimiter to aiolimiter --- providers/nvidia_nim.py | 33 +++++------- providers/rate_limit.py | 47 ++++++++++++----- providers/utils/__init__.py | 2 - providers/utils/rate_limiter.py | 72 -------------------------- tests/verify_provider_limiter.py | 87 ++++++++++++++++++++++++++++++++ 5 files changed, 135 insertions(+), 106 deletions(-) delete mode 100644 providers/utils/rate_limiter.py create mode 100644 tests/verify_provider_limiter.py diff --git a/providers/nvidia_nim.py b/providers/nvidia_nim.py index 2033a6f7..d36acc2b 100644 --- a/providers/nvidia_nim.py +++ b/providers/nvidia_nim.py @@ -7,7 +7,7 @@ import uuid from typing import Dict, Any, AsyncIterator import httpx -from httpx import TimeoutException, ReadTimeout, ConnectTimeout +from httpx import TimeoutException, ConnectTimeout from .base import BaseProvider, ProviderConfig from .utils import ( @@ -107,21 +107,19 @@ class NvidiaNimProvider(BaseProvider): self, request: Any, input_tokens: int = 0 ) -> AsyncIterator[str]: """Stream response in Anthropic SSE format.""" - # Wait if globally rate limited - blocked = await self._global_rate_limiter.wait_if_blocked() - if blocked: - # Yield error event for rate limit blocking + # Wait if globally rate limited (proactive throttle + reactive block) + waited_reactively = await self._global_rate_limiter.wait_if_blocked() + + if waited_reactively: + # Yield error event for reactive rate limit blocking (user feedback) message_id = f"msg_{uuid.uuid4()}" sse = SSEBuilder(message_id, request.model, input_tokens) - error_msg = "⏱️ Rate limit exceeded. Please try again in a minute." - logger.warning(f"NIM_STREAM: Rate limit blocked, yielding error event") + error_msg = "⏱️ Global rate limit active. Resuming now..." + logger.info(f"NIM_STREAM: Reactive block detected, notified user") yield sse.message_start() for event in sse.emit_error(error_msg): yield event - yield sse.message_delta("stop", 0) - yield sse.message_stop() - yield sse.done() - return + # After notification, we continue to the actual request body = self._build_request_body(request, stream=True) # Log compact request summary @@ -228,13 +226,13 @@ class NvidiaNimProvider(BaseProvider): logger.error(f"NIM_ERROR: {type(e).__name__}: {e}") error_occurred = True error_message = str(e) - logger.info(f"NIM_STREAM: Emitting SSE error event for exception") + logger.info("NIM_STREAM: Emitting SSE error event for exception") # Handle errors if error_occurred: for event in sse.emit_error(error_message): yield event - logger.info(f"NIM_STREAM: Error event yielded, total events emitted") + logger.info("NIM_STREAM: Error event yielded, total events emitted") # Flush remaining content from parsers remaining = think_parser.flush() @@ -406,13 +404,8 @@ class NvidiaNimProvider(BaseProvider): async def complete(self, request: Any) -> dict: """Make a non-streaming completion request.""" - # Wait if globally rate limited - wait_time = await self._global_rate_limiter.wait_if_blocked() - if wait_time > 0: - # Raise error for rate limit blocking in non-streaming mode - error_msg = "⏱️ Rate limit exceeded. Please try again in a minute." - logger.warning(f"NIM_COMPLETE: Rate limit blocked for {wait_time:.1f}s, raising error") - raise APIError(error_msg, status_code=429) + # Wait if globally rate limited (proactive throttle + reactive block) + await self._global_rate_limiter.wait_if_blocked() body = self._build_request_body(request, stream=False) # Log compact request summary diff --git a/providers/rate_limit.py b/providers/rate_limit.py index 610e39ff..d2d28d15 100644 --- a/providers/rate_limit.py +++ b/providers/rate_limit.py @@ -3,7 +3,9 @@ import asyncio import time import logging +import os from typing import Optional +from aiolimiter import AsyncLimiter logger = logging.getLogger(__name__) @@ -11,17 +13,31 @@ logger = logging.getLogger(__name__) class GlobalRateLimiter: """ Global singleton rate limiter that blocks all requests - when a rate limit error is encountered. + when a rate limit error is encountered (reactive) and + throttles requests (proactive) using aiolimiter. - No proactive limits - only reactive when 429 is hit. - No retry logic - just pauses all requests until cooldown expires. + Proactive limits - throttles requests to stay within API limits. + Reactive limits - pauses all requests when a 429 is hit. """ _instance: Optional["GlobalRateLimiter"] = None def __init__(self): + # Prevent double initialization in singleton + if hasattr(self, "_initialized"): + return + + rate_limit = int(os.getenv("NVIDIA_NIM_RATE_LIMIT", "40")) + rate_window = float(os.getenv("NVIDIA_NIM_RATE_WINDOW", "60.0")) + + self.limiter = AsyncLimiter(rate_limit, rate_window) self._blocked_until: float = 0 self._lock = asyncio.Lock() + self._initialized = True + + logger.info( + f"GlobalRateLimiter (Provider) initialized ({rate_limit} req / {rate_window}s)" + ) @classmethod def get_instance(cls) -> "GlobalRateLimiter": @@ -37,33 +53,40 @@ class GlobalRateLimiter: async def wait_if_blocked(self) -> bool: """ - Wait if currently rate limited. + Wait if currently rate limited or throttle to meet quota. Returns: - True if was blocked and waited, False if not blocked + True if was reactively blocked and waited, False otherwise. """ + # 1. Reactive check: Wait if someone hit a 429 + waited_reactively = False now = time.time() if now < self._blocked_until: wait_time = self._blocked_until - now - logger.warning(f"Global rate limit active, waiting {wait_time:.1f}s...") + logger.warning( + f"Global provider rate limit active (reactive), waiting {wait_time:.1f}s..." + ) await asyncio.sleep(wait_time) - return True - return False + waited_reactively = True + + # 2. Proactive check: Acquire slot from aiolimiter + async with self.limiter: + return waited_reactively def set_blocked(self, seconds: float = 60) -> None: """ - Set global block for specified seconds. + Set global block for specified seconds (reactive). Args: seconds: How long to block (default 60s) """ self._blocked_until = time.time() + seconds - logger.warning(f"Global rate limit set for {seconds:.1f}s") + logger.warning(f"Global provider rate limit set for {seconds:.1f}s (reactive)") def is_blocked(self) -> bool: - """Check if currently blocked.""" + """Check if currently reactively blocked.""" return time.time() < self._blocked_until def remaining_wait(self) -> float: - """Get remaining wait time in seconds.""" + """Get remaining reactive wait time in seconds.""" return max(0, self._blocked_until - time.time()) diff --git a/providers/utils/__init__.py b/providers/utils/__init__.py index c535c56e..4cfa0bba 100644 --- a/providers/utils/__init__.py +++ b/providers/utils/__init__.py @@ -1,6 +1,5 @@ """Utility modules for providers.""" -from .rate_limiter import SlidingWindowRateLimiter from .sse_builder import SSEBuilder, ContentBlockManager, map_stop_reason from .think_parser import ( ThinkTagParser, @@ -17,7 +16,6 @@ from .message_converter import ( ) __all__ = [ - "SlidingWindowRateLimiter", "SSEBuilder", "ContentBlockManager", "map_stop_reason", diff --git a/providers/utils/rate_limiter.py b/providers/utils/rate_limiter.py deleted file mode 100644 index a029fca4..00000000 --- a/providers/utils/rate_limiter.py +++ /dev/null @@ -1,72 +0,0 @@ -"""Reusable async rate limiter with sliding window.""" - -import asyncio -import time -import threading -from collections import deque -from typing import Deque - - -class SlidingWindowRateLimiter: - """ - Async rate limiter using sliding window algorithm. - - Thread-safe for use across multiple async contexts. - """ - - def __init__( - self, - rate_limit: int = 40, - window_seconds: int = 60, - max_retries: int = 120, - ): - self.rate_limit = rate_limit - self.window_seconds = window_seconds - self.max_retries = max_retries - self._timestamps: Deque[float] = deque() - self._lock = threading.Condition() - - async def acquire(self) -> None: - """ - Acquire a rate limit slot, waiting if necessary. - - Uses exponential backoff up to max_retries attempts. - """ - for _ in range(self.max_retries): - now = time.time() - - with self._lock: - # Remove expired timestamps - while self._timestamps and now - self._timestamps[0] > self.window_seconds: - self._timestamps.popleft() - - # Check if we can proceed - if len(self._timestamps) < self.rate_limit: - self._timestamps.append(now) - return - - # Calculate wait time - wait_time = self._timestamps[0] + self.window_seconds - now - - if wait_time <= 0: - continue - - await asyncio.sleep(min(wait_time, 1.0)) - - # Fallback: allow request after max retries - with self._lock: - self._timestamps.append(time.time()) - - def reset(self) -> None: - """Clear all timestamps.""" - with self._lock: - self._timestamps.clear() - - @property - def current_count(self) -> int: - """Current number of requests in the window.""" - now = time.time() - with self._lock: - while self._timestamps and now - self._timestamps[0] > self.window_seconds: - self._timestamps.popleft() - return len(self._timestamps) diff --git a/tests/verify_provider_limiter.py b/tests/verify_provider_limiter.py new file mode 100644 index 00000000..50c2bc82 --- /dev/null +++ b/tests/verify_provider_limiter.py @@ -0,0 +1,87 @@ +import asyncio +import os +import time +import logging +import sys + +# Add current directory to path +sys.path.append(os.getcwd()) + +from providers.rate_limit import GlobalRateLimiter + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + + +async def test_proactive_throttling(): + print("\n--- Testing Proactive Throttling (3 req / 1 sec) ---") + os.environ["NVIDIA_NIM_RATE_LIMIT"] = "3" + os.environ["NVIDIA_NIM_RATE_WINDOW"] = "1.0" + + GlobalRateLimiter.reset_instance() + limiter = GlobalRateLimiter.get_instance() + + start_time = time.time() + + async def call_limiter(i): + await limiter.wait_if_blocked() + print(f"[{time.time() - start_time:.2f}s] Request {i} passed") + + print("Sending 5 requests sequentially...") + # 3 should pass immediately, 4th and 5th should wait + for i in range(5): + await call_limiter(i) + + total_time = time.time() - start_time + print(f"All requests completed in {total_time:.2f}s") + + # 5 requests with limit 3 per 1s: + # R0, R1, R2 at ~0s + # R3 waits until R0 is 1s old -> ~1s + # R4 waits until R1 is 1s old -> ~1s + if total_time >= 0.9: # Allow some jitter + print("[SUCCESS] Proactive throttling working!") + else: + print(f"[FAILURE] Proactive throttling failed! Took only {total_time:.2f}s") + + +async def test_reactive_blocking(): + print("\n--- Testing Reactive Blocking ---") + GlobalRateLimiter.reset_instance() + limiter = GlobalRateLimiter.get_instance() + + start_time = time.time() + + print("Setting manual block for 2s...") + limiter.set_blocked(2) + + async def call_limiter(i): + waited = await limiter.wait_if_blocked() + print( + f"[{time.time() - start_time:.2f}s] Request {i} passed (waited reactively: {waited})" + ) + return waited + + results = await asyncio.gather(*[call_limiter(i) for i in range(2)]) + + total_time = time.time() - start_time + print(f"Requests completed in {total_time:.2f}s") + + if total_time >= 1.9 and any(results): + print("[SUCCESS] Reactive blocking working!") + else: + print(f"[FAILURE] Reactive blocking failed! Took only {total_time:.2f}s") + + +async def main(): + try: + await test_proactive_throttling() + await test_reactive_blocking() + except Exception as e: + print(f"Error during verification: {e}") + finally: + print("\nVerification complete.") + + +if __name__ == "__main__": + asyncio.run(main())