Feat: rate limiting (#470)

S1ro1 · web-flow · commit 01a0072acaa3 · 2026-03-17T21:13:51.000+01:00
diff --git a/src/kernelbot/api/api_utils.py b/src/kernelbot/api/api_utils.py
@@ -146,7 +146,7 @@ async def _run_submission(
     submission: SubmissionRequest, mode: SubmissionMode, backend: KernelBackend
 ):
     try:
-        req = prepare_submission(submission, backend)
+        req = prepare_submission(submission, backend, mode)
     except Exception as e:
         raise HTTPException(status_code=400, detail=str(e)) from e
 
diff --git a/src/kernelbot/api/main.py b/src/kernelbot/api/main.py
@@ -478,8 +478,9 @@ async def enqueue_background_job(
             file_name=req.file_name,
             code=req.code,
             user_id=req.user_id,
-            time=datetime.datetime.now(),
+            time=datetime.datetime.now(datetime.timezone.utc),
             user_name=req.user_name,
+            mode_category=req.mode_category,
         )
         job_id = db.upsert_submission_job_status(sub_id, "initial", None)
     # put submission request in queue
@@ -523,8 +524,10 @@ async def run_submission_async(
                 user_info, submission_mode, file, leaderboard_name, gpu_type, db_context
             )
 
-            req = prepare_submission(submission_request, backend_instance)
+            req = prepare_submission(submission_request, backend_instance, submission_mode_enum)
 
+        except KernelBotError as e:
+            raise HTTPException(status_code=e.http_code, detail=str(e)) from e
         except Exception as e:
             raise HTTPException(
                 status_code=400, detail=f"failed to prepare submission request: {str(e)}"
@@ -1041,6 +1044,58 @@ async def admin_set_visibility(
     return {"status": "ok", "leaderboard": leaderboard_name, "visibility": visibility}
 
 
+@app.put("/admin/leaderboards/{leaderboard_name}/rate-limits")
+async def admin_set_rate_limit(
+    leaderboard_name: str,
+    payload: dict,
+    _: Annotated[None, Depends(require_admin)],
+    db_context=Depends(get_db),
+) -> dict:
+    """Create or update a rate limit for a leaderboard."""
+    mode_category = payload.get("mode_category")
+    if mode_category not in ("test", "leaderboard"):
+        raise HTTPException(status_code=400, detail="mode_category must be 'test' or 'leaderboard'")
+    max_per_hour = payload.get("max_submissions_per_hour")
+    if not isinstance(max_per_hour, int) or max_per_hour < 1:
+        raise HTTPException(status_code=400, detail="max_submissions_per_hour must be a positive integer")
+    try:
+        with db_context as db:
+            result = db.set_rate_limit(leaderboard_name, mode_category, max_per_hour)
+        return {"status": "ok", "rate_limit": dict(result)}
+    except KernelBotError as e:
+        raise HTTPException(status_code=e.http_code, detail=str(e)) from e
+
+
+@app.get("/admin/leaderboards/{leaderboard_name}/rate-limits")
+async def admin_get_rate_limits(
+    leaderboard_name: str,
+    _: Annotated[None, Depends(require_admin)],
+    db_context=Depends(get_db),
+) -> dict:
+    """List rate limits for a leaderboard."""
+    with db_context as db:
+        limits = db.get_rate_limits(leaderboard_name)
+    return {"status": "ok", "leaderboard": leaderboard_name, "rate_limits": [dict(r) for r in limits]}
+
+
+@app.delete("/admin/leaderboards/{leaderboard_name}/rate-limits/{mode_category}")
+async def admin_delete_rate_limit(
+    leaderboard_name: str,
+    mode_category: str,
+    _: Annotated[None, Depends(require_admin)],
+    db_context=Depends(get_db),
+) -> dict:
+    """Delete a rate limit for a leaderboard."""
+    if mode_category not in ("test", "leaderboard"):
+        raise HTTPException(status_code=400, detail="mode_category must be 'test' or 'leaderboard'")
+    try:
+        with db_context as db:
+            db.delete_rate_limit(leaderboard_name, mode_category)
+        return {"status": "ok", "leaderboard": leaderboard_name, "mode_category": mode_category}
+    except KernelBotError as e:
+        raise HTTPException(status_code=e.http_code, detail=str(e)) from e
+
+
 @app.post("/user/join")
 async def user_join_leaderboard(
     payload: dict,
diff --git a/src/kernelbot/cogs/leaderboard_cog.py b/src/kernelbot/cogs/leaderboard_cog.py
@@ -126,7 +126,7 @@ async def submit(
             leaderboard=leaderboard_name,
             identity_type="discord",
         )
-        req = prepare_submission(req, self.bot.backend)
+        req = prepare_submission(req, self.bot.backend, mode)
 
         if req.gpus is None:
             view = await self.select_gpu_view(interaction, leaderboard_name, req.task_gpus)
diff --git a/src/libkernelbot/backend.py b/src/libkernelbot/backend.py
@@ -1,10 +1,10 @@
 import asyncio
 import copy
-from datetime import datetime
+import datetime
 from types import SimpleNamespace
 from typing import Optional
 
-from libkernelbot.consts import GPU, GPU_TO_SM, SubmissionMode, get_gpu_by_name
+from libkernelbot.consts import GPU, GPU_TO_SM, SubmissionMode, get_gpu_by_name, get_mode_category
 from libkernelbot.launchers import Launcher
 from libkernelbot.leaderboard_db import LeaderboardDB
 from libkernelbot.report import (
@@ -67,8 +67,9 @@ async def submit_full(
                     file_name=req.file_name,
                     code=req.code,
                     user_id=req.user_id,
-                    time=datetime.now(),
+                    time=datetime.datetime.now(datetime.timezone.utc),
                     user_name=req.user_name,
+                    mode_category=req.mode_category or get_mode_category(mode),
                 )
         selected_gpus = [get_gpu_by_name(gpu) for gpu in req.gpus]
         try:
diff --git a/src/libkernelbot/consts.py b/src/libkernelbot/consts.py
@@ -101,6 +101,13 @@ class SubmissionMode(Enum):
     PRIVATE = "private"
 
 
+def get_mode_category(mode: "SubmissionMode") -> str:
+    """Map a SubmissionMode to its rate limit category ('test' or 'leaderboard')."""
+    if mode == SubmissionMode.LEADERBOARD:
+        return "leaderboard"
+    return "test"
+
+
 class Language(Enum):
     Python = "py"
     CUDA = "cu"
diff --git a/src/libkernelbot/db_types.py b/src/libkernelbot/db_types.py
@@ -62,4 +62,12 @@ class SubmissionItem(TypedDict):
     runs: List[RunItem]
 
 
-__all__ = [LeaderboardItem, LeaderboardRankedEntry, RunItem, SubmissionItem]
+class RateLimitItem(TypedDict):
+    id: int
+    leaderboard_id: int
+    leaderboard_name: str
+    mode_category: str
+    max_submissions_per_hour: int
+
+
+__all__ = [LeaderboardItem, LeaderboardRankedEntry, RunItem, SubmissionItem, RateLimitItem]
diff --git a/src/libkernelbot/leaderboard_db.py b/src/libkernelbot/leaderboard_db.py
@@ -9,6 +9,7 @@
     IdentityType,
     LeaderboardItem,
     LeaderboardRankedEntry,
+    RateLimitItem,
     RunItem,
     SubmissionItem,
 )
@@ -276,8 +277,13 @@ def create_submission(
         code: str,
         time: datetime.datetime,
         user_name: str = None,
+        mode_category: str = None,
     ) -> Optional[int]:
         try:
+            if time.tzinfo is None:
+                time = time.astimezone()
+            time = time.astimezone(datetime.timezone.utc)
+
             # check if we already have the code
             self.cursor.execute(
                 """
@@ -323,10 +329,10 @@ def create_submission(
             self.cursor.execute(
                 """
                 INSERT INTO leaderboard.submission (leaderboard_id, file_name,
-                    user_id, code_id, submission_time)
+                    user_id, code_id, submission_time, mode_category)
                 VALUES (
                     (SELECT id FROM leaderboard.leaderboard WHERE name = %s),
-                    %s, %s, %s, %s)
+                    %s, %s, %s, %s, %s)
                 RETURNING id
                 """,
                 (
@@ -335,6 +341,7 @@ def create_submission(
                     user_id,
                     code_id,
                     time,
+                    mode_category,
                 ),
             )
             submission_id = self.cursor.fetchone()[0]
@@ -1438,6 +1445,154 @@ def validate_cli_id(self, cli_id: str) -> Optional[dict[str, str]]:
             raise KernelBotError("Error validating CLI ID") from e
 
 
+    def set_rate_limit(self, leaderboard_name: str, mode_category: str, max_per_hour: int) -> RateLimitItem:
+        try:
+            self.cursor.execute(
+                """
+                INSERT INTO leaderboard.rate_limit (leaderboard_id, mode_category, max_submissions_per_hour)
+                VALUES (
+                    (SELECT id FROM leaderboard.leaderboard WHERE name = %s),
+                    %s, %s
+                )
+                ON CONFLICT (leaderboard_id, mode_category)
+                DO UPDATE SET max_submissions_per_hour = EXCLUDED.max_submissions_per_hour
+                RETURNING id, leaderboard_id, mode_category, max_submissions_per_hour
+                """,
+                (leaderboard_name, mode_category, max_per_hour),
+            )
+            row = self.cursor.fetchone()
+            if row is None:
+                raise LeaderboardDoesNotExist(leaderboard_name)
+            self.connection.commit()
+            return RateLimitItem(
+                id=row[0],
+                leaderboard_id=row[1],
+                leaderboard_name=leaderboard_name,
+                mode_category=row[2],
+                max_submissions_per_hour=row[3],
+            )
+        except psycopg2.Error as e:
+            self.connection.rollback()
+            if "null value in column" in str(e):
+                raise LeaderboardDoesNotExist(leaderboard_name) from e
+            logger.exception("Error setting rate limit", exc_info=e)
+            raise KernelBotError("Error setting rate limit") from e
+
+    def get_rate_limits(self, leaderboard_name: str) -> List[RateLimitItem]:
+        try:
+            self.cursor.execute(
+                """
+                SELECT rl.id, rl.leaderboard_id, rl.mode_category, rl.max_submissions_per_hour
+                FROM leaderboard.rate_limit rl
+                JOIN leaderboard.leaderboard lb ON rl.leaderboard_id = lb.id
+                WHERE lb.name = %s
+                """,
+                (leaderboard_name,),
+            )
+            return [
+                RateLimitItem(
+                    id=row[0],
+                    leaderboard_id=row[1],
+                    leaderboard_name=leaderboard_name,
+                    mode_category=row[2],
+                    max_submissions_per_hour=row[3],
+                )
+                for row in self.cursor.fetchall()
+            ]
+        except psycopg2.Error as e:
+            self.connection.rollback()
+            logger.exception("Error getting rate limits", exc_info=e)
+            raise KernelBotError("Error getting rate limits") from e
+
+    def delete_rate_limit(self, leaderboard_name: str, mode_category: str) -> None:
+        try:
+            self.cursor.execute(
+                """
+                DELETE FROM leaderboard.rate_limit
+                WHERE leaderboard_id = (SELECT id FROM leaderboard.leaderboard WHERE name = %s)
+                    AND mode_category = %s
+                """,
+                (leaderboard_name, mode_category),
+            )
+            if self.cursor.rowcount == 0:
+                raise KernelBotError(
+                    f"No rate limit found for '{leaderboard_name}' with category '{mode_category}'",
+                    code=404,
+                )
+            self.connection.commit()
+        except KernelBotError:
+            raise
+        except psycopg2.Error as e:
+            self.connection.rollback()
+            logger.exception("Error deleting rate limit", exc_info=e)
+            raise KernelBotError("Error deleting rate limit") from e
+
+    def check_rate_limit(
+        self, leaderboard_name: str, user_id: str, mode_category: str
+    ) -> Optional[dict]:
+        """Check if a user has exceeded the rate limit for a leaderboard+category.
+
+        Returns None if no rate limit is configured, otherwise returns a dict with:
+        - allowed: bool
+        - current_count: int
+        - max_per_hour: int
+        - retry_after_seconds: int (0 if allowed)
+        """
+        try:
+            # Get the rate limit config
+            self.cursor.execute(
+                """
+                SELECT rl.max_submissions_per_hour
+                FROM leaderboard.rate_limit rl
+                JOIN leaderboard.leaderboard lb ON rl.leaderboard_id = lb.id
+                WHERE lb.name = %s AND rl.mode_category = %s
+                """,
+                (leaderboard_name, mode_category),
+            )
+            row = self.cursor.fetchone()
+            if row is None:
+                return None
+
+            max_per_hour = row[0]
+
+            # Count submissions in the last hour
+            self.cursor.execute(
+                """
+                SELECT COUNT(*), MIN(s.submission_time)
+                FROM leaderboard.submission s
+                JOIN leaderboard.leaderboard lb ON s.leaderboard_id = lb.id
+                WHERE lb.name = %s
+                    AND s.user_id = %s
+                    AND s.mode_category = %s
+                    AND s.submission_time > NOW() - INTERVAL '1 hour'
+                """,
+                (leaderboard_name, user_id, mode_category),
+            )
+            count_row = self.cursor.fetchone()
+            current_count = count_row[0]
+            oldest_time = count_row[1]
+
+            allowed = current_count < max_per_hour
+            retry_after = 0
+            if not allowed and oldest_time is not None:
+                import datetime as dt
+
+                expiry = oldest_time + dt.timedelta(hours=1)
+                now = dt.datetime.now(dt.timezone.utc)
+                retry_after = max(0, int((expiry - now).total_seconds()))
+
+            return {
+                "allowed": allowed,
+                "current_count": current_count,
+                "max_per_hour": max_per_hour,
+                "retry_after_seconds": retry_after,
+            }
+        except psycopg2.Error as e:
+            self.connection.rollback()
+            logger.exception("Error checking rate limit", exc_info=e)
+            raise KernelBotError("Error checking rate limit") from e
+
+
 class LeaderboardDoesNotExist(KernelBotError):
     def __init__(self, name: str):
         super().__init__(message=f"Leaderboard `{name}` does not exist.", code=404)
diff --git a/src/libkernelbot/submission.py b/src/libkernelbot/submission.py
@@ -7,7 +7,7 @@
 
 from better_profanity import profanity
 
-from libkernelbot.consts import RankCriterion
+from libkernelbot.consts import RankCriterion, SubmissionMode, get_mode_category
 from libkernelbot.db_types import RunItem, SubmissionItem
 from libkernelbot.leaderboard_db import LeaderboardDB, LeaderboardItem
 from libkernelbot.run_eval import FullResult
@@ -38,10 +38,11 @@ class ProcessedSubmissionRequest(SubmissionRequest):
     task: LeaderboardTask = None
     secret_seed: int = None
     task_gpus: list = None
+    mode_category: str = None
 
 
 def prepare_submission(  # noqa: C901
-    req: SubmissionRequest, backend: "KernelBackend"
+    req: SubmissionRequest, backend: "KernelBackend", mode: SubmissionMode = None
 ) -> ProcessedSubmissionRequest:
     if not backend.accepts_jobs:
         raise KernelBotError(
@@ -70,6 +71,19 @@ def prepare_submission(  # noqa: C901
                 )
             if not db.check_leaderboard_access(req.leaderboard, str(req.user_id)):
                 raise KernelBotError("You do not have access to this leaderboard", code=403)
+
+        mode_category = get_mode_category(mode) if mode else None
+        if mode_category is not None:
+            with backend.db as db:
+                rate_check = db.check_rate_limit(req.leaderboard, str(req.user_id), mode_category)
+                if rate_check and not rate_check["allowed"]:
+                    raise KernelBotError(
+                        f"Rate limit exceeded: {rate_check['current_count']}/{rate_check['max_per_hour']} "
+                        f"{mode_category} submissions per hour. "
+                        f"Try again in {rate_check['retry_after_seconds']}s.",
+                        code=429,
+                    )
+
     check_deadline(leaderboard)
 
     task_gpus = get_avail_gpus(req.leaderboard, backend.db)
@@ -91,6 +105,7 @@ def prepare_submission(  # noqa: C901
         task=leaderboard["task"],
         secret_seed=leaderboard["secret_seed"],
         task_gpus=task_gpus,
+        mode_category=mode_category,
     )
 
 
diff --git a/src/migrations/20260317_01_rate-limits.py b/src/migrations/20260317_01_rate-limits.py
diff --git a/tests/conftest.py b/tests/conftest.py
diff --git a/tests/test_admin_api.py b/tests/test_admin_api.py
diff --git a/tests/test_leaderboard_db.py b/tests/test_leaderboard_db.py

Original file line number	Diff line number	Diff line change
`@@ -126,7 +126,7 @@ async def submit(`
`126`	`126`	`leaderboard=leaderboard_name,`
`127`	`127`	`identity_type="discord",`
`128`	`128`	`)`
`129`		`- req = prepare_submission(req, self.bot.backend)`
	`129`	`+ req = prepare_submission(req, self.bot.backend, mode)`
`130`	`130`
`131`	`131`	`if req.gpus is None:`
`132`	`132`	`view = await self.select_gpu_view(interaction, leaderboard_name, req.task_gpus)`