Spaces:
Sleeping
Sleeping
File size: 2,621 Bytes
5db30ac | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | from __future__ import annotations
import asyncio
from collections import deque
from time import monotonic
from elbench.schemas.config import RateLimitConfig
class SlidingWindowRateLimiter:
def __init__(self, config: RateLimitConfig) -> None:
self._config = config
self._lock = asyncio.Lock()
self._request_events_1s: deque[float] = deque()
self._request_events_60s: deque[float] = deque()
self._token_events_60s: deque[tuple[float, int]] = deque()
concurrency = config.max_concurrency or 100
self._semaphore = asyncio.Semaphore(concurrency)
async def acquire(self, estimated_tokens: int = 0) -> None:
await self._semaphore.acquire()
await self._wait_for_slot(estimated_tokens)
def release(self) -> None:
self._semaphore.release()
async def _wait_for_slot(self, estimated_tokens: int) -> None:
while True:
async with self._lock:
now = monotonic()
self._prune(now)
delay = self._compute_delay(now, estimated_tokens)
if delay <= 0:
self._request_events_1s.append(now)
self._request_events_60s.append(now)
if estimated_tokens > 0:
self._token_events_60s.append((now, estimated_tokens))
return
await asyncio.sleep(delay)
def _compute_delay(self, now: float, estimated_tokens: int) -> float:
delays: list[float] = []
if self._config.qps and len(self._request_events_1s) >= self._config.qps:
delays.append(max(0.0, 1.0 - (now - self._request_events_1s[0])))
if self._config.rpm and len(self._request_events_60s) >= self._config.rpm:
delays.append(max(0.0, 60.0 - (now - self._request_events_60s[0])))
if self._config.tpm:
used_tokens = sum(tokens for _, tokens in self._token_events_60s)
if used_tokens + estimated_tokens > self._config.tpm and self._token_events_60s:
delays.append(max(0.0, 60.0 - (now - self._token_events_60s[0][0])))
return max(delays) if delays else 0.0
def _prune(self, now: float) -> None:
while self._request_events_1s and now - self._request_events_1s[0] >= 1.0:
self._request_events_1s.popleft()
while self._request_events_60s and now - self._request_events_60s[0] >= 60.0:
self._request_events_60s.popleft()
while self._token_events_60s and now - self._token_events_60s[0][0] >= 60.0:
self._token_events_60s.popleft()
|