157 lines
5.5 KiB
Python
157 lines
5.5 KiB
Python
"""Rate limiter distribuido para llamadas a OpenAI.
|
|
|
|
Usa Redis como backend para compartir el estado entre múltiples workers.
|
|
Implementa el algoritmo de token bucket con ventana deslizante para respetar
|
|
los límites de RPM (requests per minute) y TPM (tokens per minute) de OpenAI.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import time
|
|
from typing import Optional
|
|
|
|
import redis.asyncio as aioredis
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Límites por defecto para OpenAI gpt-4o (tier 1).
|
|
# Sobreescribibles con variables de entorno vía config.
|
|
DEFAULT_RPM = 500 # requests per minute
|
|
DEFAULT_TPM = 800_000 # tokens per minute
|
|
|
|
|
|
class DistributedRateLimiter:
|
|
"""Token bucket distribuido respaldado por Redis.
|
|
|
|
Un solo objeto se comparte entre todos los workers de un proceso.
|
|
Múltiples réplicas del worker comparten el estado en Redis.
|
|
|
|
Uso:
|
|
limiter = DistributedRateLimiter(redis_url="redis://redis:6379/1")
|
|
async with limiter.acquire(estimated_tokens=4000):
|
|
response = await openai_client.chat.completions.create(...)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
redis_url: str,
|
|
rpm: int = DEFAULT_RPM,
|
|
tpm: int = DEFAULT_TPM,
|
|
key_prefix: str = "ips:rl",
|
|
) -> None:
|
|
self._redis_url = redis_url
|
|
self.rpm = rpm
|
|
self.tpm = tpm
|
|
self._prefix = key_prefix
|
|
self._client: Optional[aioredis.Redis] = None
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def _get_client(self) -> aioredis.Redis:
|
|
if self._client is None:
|
|
async with self._lock:
|
|
if self._client is None:
|
|
self._client = aioredis.from_url(
|
|
self._redis_url,
|
|
encoding="utf-8",
|
|
decode_responses=True,
|
|
socket_connect_timeout=5,
|
|
)
|
|
return self._client
|
|
|
|
async def _sliding_window_count(self, key: str, window_s: float = 60.0) -> int:
|
|
"""Cuenta eventos en la ventana deslizante usando un Sorted Set de Redis."""
|
|
client = await self._get_client()
|
|
now = time.monotonic()
|
|
# Eliminar eventos expirados
|
|
await client.zremrangebyscore(key, "-inf", now - window_s)
|
|
return await client.zcard(key)
|
|
|
|
async def _record_event(self, key: str, tokens: int = 0, window_s: float = 60.0) -> None:
|
|
"""Registra un evento en el sorted set. El score es el timestamp."""
|
|
client = await self._get_client()
|
|
now = time.monotonic()
|
|
member = f"{now}:{tokens}"
|
|
pipe = client.pipeline()
|
|
pipe.zadd(key, {member: now})
|
|
pipe.expire(key, int(window_s) + 5)
|
|
await pipe.execute()
|
|
|
|
async def _token_sum_in_window(self, key: str, window_s: float = 60.0) -> int:
|
|
"""Suma los tokens de todos los eventos en la ventana."""
|
|
client = await self._get_client()
|
|
now = time.monotonic()
|
|
await client.zremrangebyscore(key, "-inf", now - window_s)
|
|
members = await client.zrange(key, 0, -1)
|
|
total = 0
|
|
for m in members:
|
|
parts = m.split(":", 1)
|
|
if len(parts) == 2:
|
|
try:
|
|
total += int(parts[1])
|
|
except ValueError:
|
|
pass
|
|
return total
|
|
|
|
async def acquire(self, estimated_tokens: int = 2000, max_wait_s: float = 120.0) -> None:
|
|
"""Espera hasta que haya capacidad disponible en RPM y TPM.
|
|
|
|
Args:
|
|
estimated_tokens: Tokens que se espera consumir (input + output).
|
|
max_wait_s: Tiempo máximo de espera antes de lanzar TimeoutError.
|
|
"""
|
|
rpm_key = f"{self._prefix}:rpm"
|
|
tpm_key = f"{self._prefix}:tpm"
|
|
start = time.monotonic()
|
|
sleep_s = 0.5
|
|
|
|
while True:
|
|
rpm_count = await self._sliding_window_count(rpm_key)
|
|
tpm_used = await self._token_sum_in_window(tpm_key)
|
|
|
|
rpm_ok = rpm_count < self.rpm
|
|
tpm_ok = (tpm_used + estimated_tokens) <= self.tpm
|
|
|
|
if rpm_ok and tpm_ok:
|
|
await self._record_event(rpm_key, tokens=0)
|
|
await self._record_event(tpm_key, tokens=estimated_tokens)
|
|
return
|
|
|
|
elapsed = time.monotonic() - start
|
|
if elapsed >= max_wait_s:
|
|
raise TimeoutError(
|
|
f"Rate limiter: espera maxima alcanzada ({max_wait_s}s). "
|
|
f"rpm={rpm_count}/{self.rpm} tpm={tpm_used}/{self.tpm}"
|
|
)
|
|
|
|
reason = "RPM" if not rpm_ok else "TPM"
|
|
logger.debug(
|
|
"Rate limiter: esperando por %s (rpm=%d/%d tpm=%d/%d) — sleep %.1fs",
|
|
reason, rpm_count, self.rpm, tpm_used, self.tpm, sleep_s,
|
|
)
|
|
await asyncio.sleep(sleep_s)
|
|
sleep_s = min(sleep_s * 1.5, 10.0)
|
|
|
|
async def close(self) -> None:
|
|
if self._client is not None:
|
|
await self._client.aclose()
|
|
self._client = None
|
|
|
|
|
|
# Instancia global compartida por todos los workers del proceso.
|
|
# Se inicializa de forma lazy en el primer uso.
|
|
_limiter: Optional[DistributedRateLimiter] = None
|
|
|
|
|
|
def get_rate_limiter() -> DistributedRateLimiter:
|
|
"""Devuelve el rate limiter global, inicializándolo si es necesario."""
|
|
global _limiter
|
|
if _limiter is None:
|
|
from domain.config import settings
|
|
_limiter = DistributedRateLimiter(
|
|
redis_url=settings.redis_url,
|
|
rpm=settings.openai_rpm_limit,
|
|
tpm=settings.openai_tpm_limit,
|
|
)
|
|
return _limiter
|