🔒 Add global 5rps ratelimit

Author: Paillat-dev

Dockerfile

@@ -16,7 +16,7 @@ ADD ./src /app/
 RUN mkdir /app/cache
 
 # ADDING USER
-RUN adduser -u 9263 --disabled-password --gecos "" appuser && chown -R appuser /app
+RUN adduser -u 9263 --disabled-password --gecos "" appuser && chown -R appuser /app && chown -R appuser /app/cache
 USER appuser
 
 HEALTHCHECK CMD curl --fail http://localhost:8000/v1/health || exit 1

app.py

@@ -3,11 +3,14 @@
 import secrets
 import logging
 import os
+
+from slowapi import Limiter, _rate_limit_exceeded_handler
+from slowapi.errors import RateLimitExceeded
 from pydantic import BaseModel
 from transformers import pipeline
 from fastapi import FastAPI, HTTPException, status, Depends
+from fastapi.requests import Request
 from fastapi.security import OAuth2PasswordBearer
-from concurrent.futures import ThreadPoolExecutor
 from dotenv import load_dotenv
 from cachier import cachier
 
@@ -19,6 +22,11 @@
 # Create FastAPI instance
 app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None)
 
+# Setup rate limiter
+limiter = Limiter(key_func=lambda: "global", default_limits=["1000 per day"])
+app.state.limiter = limiter  # noqa
+app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)  # noqa
+
 # Setup logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -93,8 +101,12 @@ def classify_sync(message: str, labels: list[str]) -> dict:
 
 # Route to classify message
 @app.get("/v1/classify")
+@limiter.limit("6/second")
 async def classify(
-    message: str, labels: list[str] = None, token: str = Depends(authenticate_user)
+    request: Request,
+    message: str,
+    labels: list[str] = None,
+    token: str = Depends(authenticate_user),
 ) -> Classification:
     labels = labels or DEFAULT_LABELS
     async with classification_lock:  # Ensure only one classification at a time

pdm.lock

@@ -11,6 +11,7 @@ dependencies = [
     "transformers>=4.41.2",
     "fastapi>=0.111.0",
     "cachier>=3.0.0",
+    "slowapi>=0.1.9",
 ]
 requires-python = "==3.11.*"
 readme = "README.md"

pyproject.toml

@@ -8,6 +8,7 @@ certifi==2024.6.2
 charset-normalizer==3.3.2
 click==8.1.7
 colorama==0.4.6; sys_platform == "win32" or platform_system == "Windows"
+deprecated==1.2.14
 dnspython==2.6.1
 email-validator==2.2.0
 fastapi==0.111.0
@@ -20,8 +21,10 @@ httptools==0.6.1
 httpx==0.27.0
 huggingface-hub==0.23.4
 idna==3.7
+importlib-resources==6.4.0
 intel-openmp==2021.4.0; platform_system == "Windows"
 jinja2==3.1.4
+limits==3.12.0
 markdown-it-py==3.0.0
 markupsafe==2.1.5
 mdurl==0.1.2
@@ -57,6 +60,7 @@ rich==13.7.1
 safetensors==0.4.3
 setuptools==70.1.0
 shellingham==1.5.4
+slowapi==0.1.9
 sniffio==1.3.1
 starlette==0.37.2
 sympy==1.12.1
@@ -75,5 +79,6 @@ uvloop==0.19.0; (sys_platform != "cygwin" and sys_platform != "win32") and platf
 watchdog==4.0.1
 watchfiles==0.22.0
 websockets==12.0
+wrapt==1.16.0
 --index-url https://pypi.org/simple
 --extra-index-url https://download.pytorch.org/whl/cpu