refactor api

api/src/production/limits.py

@@ -0,0 +1,95 @@
+"""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)