Implement rate limiting

Author: S1ro1

src/kernelbot/api/api_utils.py

@@ -18,6 +18,7 @@
     SubmissionRequest,
     prepare_submission,
 )
+from libkernelbot.utils import KernelBotError
 
 
 async def _handle_discord_oauth(code: str, redirect_uri: str) -> tuple[str, str]:
@@ -147,6 +148,8 @@ async def _run_submission(
 ):
     try:
         req = prepare_submission(submission, backend)
+    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=str(e)) from e
 

src/kernelbot/api/main.py

@@ -36,6 +36,7 @@
 
 app = FastAPI()
 
+
 def json_serializer(obj):
     """JSON serializer for objects not serializable by default json code"""
     if isinstance(obj, (datetime.datetime, datetime.date, datetime.time)):
@@ -255,10 +256,16 @@ async def cli_auth(auth_provider: str, code: str, state: str, db_context=Depends
         raise e
     except Exception as e:
         # Catch unexpected errors during OAuth handling
-        raise HTTPException(status_code=500, detail=f"Error during {auth_provider} OAuth flow: {e}") from e
+        raise HTTPException(
+            status_code=500,
+            detail=f"Error during {auth_provider} OAuth flow: {e}",
+        ) from e
 
     if not user_id or not user_name:
-        raise HTTPException(status_code=500,detail="Failed to retrieve user ID or username from provider.",)
+        raise HTTPException(
+            status_code=500,
+            detail="Failed to retrieve user ID or username from provider.",
+        )
 
     try:
         with db_context as db:
@@ -268,7 +275,10 @@ async def cli_auth(auth_provider: str, code: str, state: str, db_context=Depends
                 db.create_user_from_cli(user_id, user_name, cli_id, auth_provider)
 
     except AttributeError as e:
-        raise HTTPException(status_code=500, detail=f"Database interface error during update: {e}") from e
+        raise HTTPException(
+            status_code=500,
+            detail=f"Database interface error during update: {e}",
+        ) from e
     except Exception as e:
         raise HTTPException(status_code=400, detail=f"Database update failed: {e}") from e
 
@@ -280,6 +290,7 @@ async def cli_auth(auth_provider: str, code: str, state: str, db_context=Depends
         "is_reset": is_reset,
     }
 
+
 async def _stream_submission_response(
     submission_request: SubmissionRequest,
     submission_mode_enum: SubmissionMode,
@@ -298,18 +309,22 @@ async def _stream_submission_response(
 
         while not task.done():
             elapsed_time = time.time() - start_time
-            yield f"event: status\ndata: {json.dumps({'status': 'processing',
-                                                      'elapsed_time': round(elapsed_time, 2)},
-                                                      default=json_serializer)}\n\n"
+            status_data = json.dumps(
+                {"status": "processing", "elapsed_time": round(elapsed_time, 2)},
+                default=json_serializer,
+            )
+            yield f"event: status\ndata: {status_data}\n\n"
 
             try:
                 await asyncio.wait_for(asyncio.shield(task), timeout=15.0)
             except asyncio.TimeoutError:
                 continue
             except asyncio.CancelledError:
-                yield f"event: error\ndata: {json.dumps(
-                    {'status': 'error', 'detail': 'Submission cancelled'},
-                    default=json_serializer)}\n\n"
+                error_data = json.dumps(
+                    {"status": "error", "detail": "Submission cancelled"},
+                    default=json_serializer,
+                )
+                yield f"event: error\ndata: {error_data}\n\n"
                 return
 
         result, reports = await task
@@ -343,6 +358,7 @@ async def _stream_submission_response(
             except asyncio.CancelledError:
                 pass
 
+
 @app.post("/{leaderboard_name}/{gpu_type}/{submission_mode}")
 async def run_submission(  # noqa: C901
     leaderboard_name: str,
@@ -381,13 +397,13 @@ async def run_submission(  # noqa: C901
     )
     return StreamingResponse(generator, media_type="text/event-stream")
 
+
 async def enqueue_background_job(
     req: ProcessedSubmissionRequest,
     mode: SubmissionMode,
     backend: KernelBackend,
     manager: BackgroundSubmissionManager,
 ):
-
     # pre-create the submission for api returns
     with backend.db as db:
         sub_id = db.create_submission(
@@ -401,7 +417,8 @@ async def enqueue_background_job(
         job_id = db.upsert_submission_job_status(sub_id, "initial", None)
     # put submission request in queue
     await manager.enqueue(req, mode, sub_id)
-    return sub_id,job_id
+    return sub_id, job_id
+
 
 @app.post("/submission/{leaderboard_name}/{gpu_type}/{submission_mode}")
 async def run_submission_async(
@@ -425,37 +442,49 @@ async def run_submission_async(
     Raises:
         HTTPException: If the kernelbot is not initialized, or header/input is invalid.
     Returns:
-        JSONResponse: A JSON response containing job_id and and submission_id for the client to poll for status.
+        JSONResponse: A JSON response containing job_id and submission_id.
+            The client can poll for status using these ids.
     """
     try:
-
         await simple_rate_limit()
-        logger.info(f"Received submission request for {leaderboard_name} {gpu_type} {submission_mode}")
-
+        logger.info(
+            "Received submission request for %s %s %s",
+            leaderboard_name,
+            gpu_type,
+            submission_mode,
+        )
 
         # throw error if submission request is invalid
         try:
             submission_request, submission_mode_enum = await to_submit_info(
-            user_info, submission_mode, file, leaderboard_name, gpu_type, db_context
+                user_info, submission_mode, file, leaderboard_name, gpu_type, db_context
             )
 
             req = prepare_submission(submission_request, backend_instance)
 
+        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)}") from e
+            raise HTTPException(
+                status_code=400,
+                detail=f"failed to prepare submission request: {str(e)}",
+            ) from e
 
         # prepare submission request before the submission is started
         if not req.gpus or len(req.gpus) != 1:
             raise HTTPException(status_code=400, detail="Invalid GPU type")
 
         # put submission request to background manager to run in background
-        sub_id,job_status_id = await enqueue_background_job(
+        sub_id, job_status_id = await enqueue_background_job(
             req, submission_mode_enum, backend_instance, background_submission_manager
         )
 
         return JSONResponse(
             status_code=202,
-            content={"details":{"id": sub_id, "job_status_id": job_status_id}, "status": "accepted"},
+            content={
+                "details": {"id": sub_id, "job_status_id": job_status_id},
+                "status": "accepted",
+            },
         )
         # Preserve FastAPI HTTPException as-is
     except HTTPException:
@@ -470,6 +499,7 @@ async def run_submission_async(
         logger.error(f"Unexpected error in api submissoin: {e}")
         raise HTTPException(status_code=500, detail="Internal server error") from e
 
+
 @app.get("/leaderboards")
 async def get_leaderboards(db_context=Depends(get_db)):
     """An endpoint that returns all leaderboards.

src/kernelbot/cogs/admin_cog.py

@@ -122,6 +122,16 @@ def __init__(self, bot: "ClusterBot"):
             name="set-forum-ids", description="Sets forum IDs"
         )(self.set_forum_ids)
 
+        self.set_submission_rate_limit = bot.admin_group.command(
+            name="set-submission-rate-limit",
+            description="Set default or per-user submission rate limit (submissions/minute).",
+        )(self.set_submission_rate_limit)
+
+        self.get_submission_rate_limit = bot.admin_group.command(
+            name="get-submission-rate-limit",
+            description="Get default or per-user submission rate limit (submissions/minute).",
+        )(self.get_submission_rate_limit)
+
         self._scheduled_cleanup_temp_users.start()
 
     # --------------------------------------------------------------------------
@@ -512,6 +522,153 @@ async def start(self, interaction: discord.Interaction):
             interaction, "Bot will accept submissions again!", ephemeral=True
         )
 
+    def _parse_user_id_arg(self, user_id: str) -> str:
+        """Accepts a raw id or a discord mention and returns the id string."""
+        s = (user_id or "").strip()
+        if s.startswith("<@") and s.endswith(">"):
+            s = s[2:-1].strip()
+            if s.startswith("!"):
+                s = s[1:].strip()
+        return s
+
+    def _format_rate(self, rate: float | None) -> str:
+        if rate is None:
+            return "unlimited"
+        r = float(rate)
+        if r == 0:
+            return "blocked"
+        return f"{r:g}/min"
+
+    @app_commands.describe(
+        rate_per_minute="Rate in submissions/minute. Use 'none' for unlimited; 'default' clears a user override.",  # noqa: E501
+        user_id="Optional user id or mention. If omitted, sets the default.",
+    )
+    @with_error_handling
+    async def set_submission_rate_limit(
+        self,
+        interaction: discord.Interaction,
+        rate_per_minute: str,
+        user_id: Optional[str] = None,
+    ):
+        is_admin = await self.admin_check(interaction)
+        if not is_admin:
+            await send_discord_message(
+                interaction,
+                "You need to be Admin to use this command.",
+                ephemeral=True,
+            )
+            return
+
+        rate_s = (rate_per_minute or "").strip().lower()
+        if rate_s in {"none", "unlimited", "off"}:
+            parsed_rate: float | None = None
+            clear_override = False
+        elif rate_s == "default":
+            parsed_rate = None
+            clear_override = True
+        else:
+            try:
+                parsed_rate = float(rate_s)
+            except ValueError:
+                await send_discord_message(
+                    interaction,
+                    "Invalid rate. Use a number like `1` or `0.5`, or `none`.",
+                    ephemeral=True,
+                )
+                return
+            if parsed_rate < 0:
+                await send_discord_message(
+                    interaction,
+                    "Invalid rate. Must be >= 0 (or `none`).",
+                    ephemeral=True,
+                )
+                return
+            clear_override = False
+
+        with self.bot.leaderboard_db as db:
+            if user_id is None or user_id.strip() == "":
+                if clear_override:
+                    await send_discord_message(
+                        interaction,
+                        "For default limit, use a number or `none` (not `default`).",
+                        ephemeral=True,
+                    )
+                    return
+                db.set_default_submission_rate_limit(parsed_rate)
+                await send_discord_message(
+                    interaction,
+                    f"Default submission rate limit set to **{self._format_rate(parsed_rate)}**.",
+                    ephemeral=True,
+                )
+                return
+
+            uid = self._parse_user_id_arg(user_id)
+            if uid == "":
+                await send_discord_message(
+                    interaction,
+                    "Invalid user id.",
+                    ephemeral=True,
+                )
+                return
+
+            if clear_override:
+                db.clear_user_submission_rate_limit(uid)
+                await send_discord_message(
+                    interaction,
+                    f"Cleared per-user submission rate limit override for `{uid}` (default now applies).",  # noqa: E501
+                    ephemeral=True,
+                )
+                return
+
+            db.set_user_submission_rate_limit(uid, parsed_rate)
+            await send_discord_message(
+                interaction,
+                f"Submission rate limit for `{uid}` set to **{self._format_rate(parsed_rate)}**.",
+                ephemeral=True,
+            )
+
+    @app_commands.describe(
+        user_id="Optional user id or mention. If omitted, shows the default.",
+    )
+    @with_error_handling
+    async def get_submission_rate_limit(
+        self,
+        interaction: discord.Interaction,
+        user_id: Optional[str] = None,
+    ):
+        is_admin = await self.admin_check(interaction)
+        if not is_admin:
+            await send_discord_message(
+                interaction,
+                "You need to be Admin to use this command.",
+                ephemeral=True,
+            )
+            return
+
+        with self.bot.leaderboard_db as db:
+            if user_id is None or user_id.strip() == "":
+                default_rate, capacity = db.get_default_submission_rate_limit()
+                msg = (
+                    f"Default submission rate limit: **{self._format_rate(default_rate)}**\n"
+                    f"Bucket capacity: **{capacity:g}**"
+                )
+                await send_discord_message(interaction, msg, ephemeral=True)
+                return
+
+            uid = self._parse_user_id_arg(user_id)
+            effective, has_override, user_rate, default_rate, capacity = (
+                db.get_submission_rate_limits(uid)
+            )
+            override_text = "default" if not has_override else self._format_rate(user_rate)
+            msg = (
+                f"User `{uid}` submission rate limit:\n"
+                f"- Effective: **{self._format_rate(effective)}**\n"
+                f"- User override: **{override_text}**\n"
+                f"- Default: **{self._format_rate(default_rate)}**\n"
+                f"- Bucket capacity: **{capacity:g}**"
+            )
+            await send_discord_message(interaction, msg, ephemeral=True)
+
     @app_commands.describe(
         problem_set="Which problem set to load.",
         repository_name="Name of the repository to load problems from (in format: user/repo)",