96 lines
3.1 KiB
Python
96 lines
3.1 KiB
Python
"""Overload protection: per-IP rate limits and the inference thread pool."""
|
|
|
|
import asyncio
|
|
from collections import defaultdict, deque
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
import time
|
|
|
|
from fastapi import HTTPException, Request, WebSocket
|
|
import torch
|
|
|
|
from src.production import settings
|
|
|
|
|
|
def client_ip(conn: Request | WebSocket) -> str:
|
|
# Caddy is the only proxy in front of this container and appends the real
|
|
# client address as the last X-Forwarded-For entry.
|
|
forwarded = conn.headers.get("x-forwarded-for")
|
|
if forwarded:
|
|
return forwarded.split(",")[-1].strip()
|
|
return conn.client.host if conn.client else "unknown"
|
|
|
|
|
|
class RateLimiter:
|
|
"""Sliding-window limiter per key (in memory, single process)."""
|
|
|
|
def __init__(self, limit: int, window: float):
|
|
self.limit = limit
|
|
self.window = window
|
|
self.hits = defaultdict(deque)
|
|
|
|
def allow(self, key: str) -> bool:
|
|
now = time.monotonic()
|
|
if len(self.hits) > 10_000:
|
|
self.hits = defaultdict(deque, {
|
|
k: v for k, v in self.hits.items() if v and now - v[-1] < self.window
|
|
})
|
|
q = self.hits[key]
|
|
while q and now - q[0] >= self.window:
|
|
q.popleft()
|
|
if len(q) >= self.limit:
|
|
return False
|
|
q.append(now)
|
|
return True
|
|
|
|
|
|
predict_limiter = RateLimiter(settings.PREDICT_RATE_LIMIT, settings.PREDICT_RATE_WINDOW)
|
|
review_limiter = RateLimiter(settings.REVIEW_RATE_LIMIT, settings.REVIEW_RATE_WINDOW)
|
|
|
|
|
|
def rate_limit(limiter: RateLimiter):
|
|
"""FastAPI dependency that answers 429 when the client is over the limit."""
|
|
def dependency(request: Request):
|
|
if not limiter.allow(client_ip(request)):
|
|
raise HTTPException(status_code=429, detail="Too many requests", headers={"Retry-After": "60"})
|
|
return dependency
|
|
|
|
|
|
class Busy(Exception):
|
|
pass
|
|
|
|
|
|
def init_inference_thread():
|
|
# Must run inside each worker thread: with OpenMP the thread count is a
|
|
# per-thread setting, so setting it once in the main thread is not enough.
|
|
torch.set_num_threads(settings.TORCH_THREADS)
|
|
|
|
|
|
class InferenceGate:
|
|
"""Runs predictions on `parallel` worker threads. At most
|
|
`parallel + queued` jobs (running + waiting) are admitted; everything
|
|
else is rejected at once instead of queueing up and eating memory.
|
|
|
|
The models are in eval mode under torch.inference_mode(), so several
|
|
threads can safely run forward passes on the same model object."""
|
|
|
|
def __init__(self, parallel: int, queued: int):
|
|
self.max_pending = parallel + queued
|
|
self.pending = 0
|
|
self.executor = ThreadPoolExecutor(
|
|
max_workers=parallel,
|
|
thread_name_prefix="inference",
|
|
initializer=init_inference_thread,
|
|
)
|
|
|
|
async def run(self, fn, *args):
|
|
if self.pending >= self.max_pending:
|
|
raise Busy
|
|
self.pending += 1
|
|
try:
|
|
return await asyncio.get_running_loop().run_in_executor(self.executor, fn, *args)
|
|
finally:
|
|
self.pending -= 1
|
|
|
|
|
|
gate = InferenceGate(settings.PARALLEL_INFERENCES, settings.QUEUED_INFERENCES)
|