From a9a086e6b2d988fa589be3e498fc524e5e88a1d4 Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Sat, 7 Feb 2026 21:54:00 +0100 Subject: [PATCH 1/3] Feat: rate limiting --- .ruff.toml | 2 +- AGENTS.md | 1 + src/kernelbot/api/api_utils.py | 27 +-- src/kernelbot/api/main.py | 82 ++++++--- src/kernelbot/cogs/admin_cog.py | 161 ++++++++-------- src/kernelbot/cogs/misc_cog.py | 32 +++- src/kernelbot/cogs/verify_run_cog.py | 22 +-- src/libkernelbot/backend.py | 16 +- src/libkernelbot/leaderboard_db.py | 173 +++++++++++++----- src/libkernelbot/run_eval.py | 35 ++-- src/libkernelbot/submission.py | 27 ++- src/libkernelbot/utils.py | 4 +- ...260207_01_EW7ve-leaderboard-rate-limits.py | 14 ++ 13 files changed, 340 insertions(+), 256 deletions(-) create mode 100644 AGENTS.md create mode 100644 src/migrations/20260207_01_EW7ve-leaderboard-rate-limits.py diff --git a/.ruff.toml b/.ruff.toml index 91078d5c..00b68b31 100644 --- a/.ruff.toml +++ b/.ruff.toml @@ -1,4 +1,4 @@ -line-length = 100 # ideally I want this to be less than 100 but don't wanna test and change files with longer lines +line-length = 120 # ideally I want this to be less than 100 but don't wanna test and change files with longer lines target-version = "py313" lint.select = [ "E", # pycodestyle errors diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 00000000..61769c98 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1 @@ +@CLAUDE.md diff --git a/src/kernelbot/api/api_utils.py b/src/kernelbot/api/api_utils.py index f26ccd3b..41661c31 100644 --- a/src/kernelbot/api/api_utils.py +++ b/src/kernelbot/api/api_utils.py @@ -74,9 +74,7 @@ async def _handle_discord_oauth(code: str, redirect_uri: str) -> tuple[str, str] user_name = user_json.get("username") if not user_id or not user_name: - raise HTTPException( - status_code=500, detail="Failed to retrieve user ID or username from Discord." - ) + raise HTTPException(status_code=500, detail="Failed to retrieve user ID or username from Discord.") return user_id, user_name @@ -135,16 +133,12 @@ async def _handle_github_oauth(code: str, redirect_uri: str) -> tuple[str, str]: user_name = user_json.get("login") # GitHub uses 'login' for username if not user_id or not user_name: - raise HTTPException( - status_code=500, detail="Failed to retrieve user ID or username from GitHub." - ) + raise HTTPException(status_code=500, detail="Failed to retrieve user ID or username from GitHub.") return user_id, user_name -async def _run_submission( - submission: SubmissionRequest, mode: SubmissionMode, backend: KernelBackend -): +async def _run_submission(submission: SubmissionRequest, mode: SubmissionMode, backend: KernelBackend): try: req = prepare_submission(submission, backend) except Exception as e: @@ -225,21 +219,6 @@ async def to_submit_info( try: with db_context as db: - # Per-user rate limit: max 1 submission per hour on Modal B200 for leaderboard 730 - if gpu_type == "B200": - lb_id = db.get_leaderboard_id(leaderboard_name) - if lb_id == 730: - last_submission_time = db.check_user_rate_limit(user_id) - if last_submission_time: - raise HTTPException( - status_code=429, - detail=( - f"Rate limit exceeded. You can submit once per hour. " - f"Last submission: {last_submission_time.isoformat()}. " - f"Consider using the NVIDIA runner instead of Modal for faster iteration." - ), - ) - leaderboard_item = db.get_leaderboard(leaderboard_name) gpus = leaderboard_item.get("gpu_types", []) if gpu_type not in gpus: diff --git a/src/kernelbot/api/main.py b/src/kernelbot/api/main.py index 2ae2bf97..f26ca91b 100644 --- a/src/kernelbot/api/main.py +++ b/src/kernelbot/api/main.py @@ -43,6 +43,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)): @@ -185,9 +186,7 @@ def require_admin( @app.get("/auth/init") async def auth_init(provider: str, db_context=Depends(get_db)) -> dict: if provider not in ["discord", "github"]: - raise HTTPException( - status_code=400, detail="Invalid provider, must be 'discord' or 'github'" - ) + raise HTTPException(status_code=400, detail="Invalid provider, must be 'discord' or 'github'") """ Initialize authentication flow for the specified provider. @@ -230,9 +229,7 @@ async def cli_auth(auth_provider: str, code: str, state: str, db_context=Depends """ if auth_provider not in ["discord", "github"]: - raise HTTPException( - status_code=400, detail="Invalid provider, must be 'discord' or 'github'" - ) + raise HTTPException(status_code=400, detail="Invalid provider, must be 'discord' or 'github'") if not code or not state: raise HTTPException(status_code=400, detail="Missing authorization code or state") @@ -252,8 +249,7 @@ async def cli_auth(auth_provider: str, code: str, state: str, db_context=Depends if not api_base_url: raise HTTPException( status_code=500, - detail="Redirect URI base not configured." - "Set HEROKU_APP_DEFAULT_DOMAIN_NAME or POPCORN_API_URL.", + detail="Redirect URI base not configured.Set HEROKU_APP_DEFAULT_DOMAIN_NAME or POPCORN_API_URL.", ) redirect_uri_base = api_base_url.rstrip("/") redirect_uri = f"https://{redirect_uri_base}/auth/cli/{auth_provider}" @@ -275,7 +271,10 @@ async def cli_auth(auth_provider: str, code: str, state: str, db_context=Depends 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: @@ -297,6 +296,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, @@ -315,18 +315,18 @@ 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" + yield f"event: status\ndata: { + json.dumps({'status': 'processing', 'elapsed_time': round(elapsed_time, 2)}, default=json_serializer) + }\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" + yield f"event: error\ndata: { + json.dumps({'status': 'error', 'detail': 'Submission cancelled'}, default=json_serializer) + }\n\n" return result, reports = await task @@ -360,6 +360,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, @@ -398,13 +399,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( @@ -412,13 +413,14 @@ 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, ) 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( @@ -445,15 +447,13 @@ async def run_submission_async( JSONResponse: A JSON response containing job_id and and submission_id for the client to poll for status. """ try: - await simple_rate_limit() logger.info(f"Received submission request for {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) @@ -466,13 +466,13 @@ async def run_submission_async( 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: @@ -536,8 +536,7 @@ async def create_dev_leaderboard( # GPUs must be specified in task.yml if not definition.gpus: raise HTTPException( - status_code=400, - detail="No gpus specified in task.yml. Add 'gpus:' field with list of GPU types." + status_code=400, detail="No gpus specified in task.yml. Add 'gpus:' field with list of GPU types." ) with db_context as db: @@ -629,7 +628,7 @@ async def admin_update_problems( branch=branch, force=force, creator_id=0, # API-created - forum_id=-1, # No Discord forum + forum_id=-1, # No Discord forum ) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) from e @@ -643,6 +642,33 @@ async def admin_update_problems( } +@app.get("/leaderboard/rate-limits/{leaderboard_name}") +async def get_leaderboard_rate_limits(leaderboard_name: str, db_context=Depends(get_db)) -> dict: + with db_context as db: + rate_limits = db.get_leaderboard_rate_limits(leaderboard_name) + return {"status": "ok", "rate_limits": rate_limits} + + +@app.post("/leaderboard/rate-limits/{leaderboard_name}/{gpu_type}") +async def set_leaderboard_gpu_rate_limit( + leaderboard_name: str, + gpu_type: str, + rate_limit_seconds: int, + _: Annotated[None, Depends(require_admin)], + db_context=Depends(get_db), +) -> dict: + if rate_limit_seconds <= 0: + rate_limit_seconds = None + with db_context as db: + db.set_leaderboard_gpu_rate_limit(leaderboard_name, gpu_type, rate_limit_seconds) + return { + "status": "ok", + "leaderboard_name": leaderboard_name, + "gpu_type": gpu_type, + "rate_limit_seconds": rate_limit_seconds, + } + + @app.get("/leaderboards") async def get_leaderboards(db_context=Depends(get_db)): """An endpoint that returns all leaderboards. @@ -692,9 +718,7 @@ async def get_submissions( try: with db_context as db: # Add validation for leaderboard and GPU? Might be redundant if DB handles it. - return db.get_leaderboard_submissions( - leaderboard_name, gpu_name, limit=limit, offset=offset - ) + return db.get_leaderboard_submissions(leaderboard_name, gpu_name, limit=limit, offset=offset) except Exception as e: raise HTTPException(status_code=500, detail=f"Error fetching submissions: {e}") from e diff --git a/src/kernelbot/cogs/admin_cog.py b/src/kernelbot/cogs/admin_cog.py index a2d0f85b..c55fbf22 100644 --- a/src/kernelbot/cogs/admin_cog.py +++ b/src/kernelbot/cogs/admin_cog.py @@ -54,9 +54,7 @@ async def leaderboard_dir_autocomplete( ) -> list[discord.app_commands.Choice[str]]: """Return leaderboard names that match the current typed name""" root = Path(env.PROBLEM_DEV_DIR) - return [ - discord.app_commands.Choice(name=x.name, value=x.name) for x in root.iterdir() if x.is_dir() - ] + return [discord.app_commands.Choice(name=x.name, value=x.name) for x in root.iterdir() if x.is_dir()] # ensure valid serialization @@ -83,9 +81,9 @@ def __init__(self, bot: "ClusterBot"): name="delete-leaderboard", description="Delete a leaderboard" )(self.delete_leaderboard) - self.delete_submission = bot.admin_group.command( - name="delete-submission", description="Delete a submission" - )(self.delete_submission) + self.delete_submission = bot.admin_group.command(name="delete-submission", description="Delete a submission")( + self.delete_submission + ) self.accept_jobs = bot.admin_group.command( name="start", description="Make the kernelbot accept new submissions" @@ -99,9 +97,9 @@ def __init__(self, bot: "ClusterBot"): name="update-problems", description="Reload all problem definitions" )(self.update_problems) - self.show_bot_stats = bot.admin_group.command( - name="show-stats", description="Show stats for the kernelbot" - )(self.show_bot_stats) + self.show_bot_stats = bot.admin_group.command(name="show-stats", description="Show stats for the kernelbot")( + self.show_bot_stats + ) self.resync = bot.admin_group.command( name="resync", description="Trigger re-synchronization of slash commands" @@ -111,17 +109,21 @@ def __init__(self, bot: "ClusterBot"): name="get-submission", description="Retrieve one of past submissions" )(self.get_submission_by_id) - self.get_user_names = bot.admin_group.command( - name="get-user-names", description="Get user names" - )(self.get_user_names) + self.get_user_names = bot.admin_group.command(name="get-user-names", description="Get user names")( + self.get_user_names + ) - self.update_user_names = bot.admin_group.command( - name="update-user-names", description="Update user names" - )(self.update_user_names) + self.update_user_names = bot.admin_group.command(name="update-user-names", description="Update user names")( + self.update_user_names + ) - self.set_forum_ids = bot.admin_group.command( - name="set-forum-ids", description="Sets forum IDs" - )(self.set_forum_ids) + self.set_forum_ids = bot.admin_group.command(name="set-forum-ids", description="Sets forum IDs")( + self.set_forum_ids + ) + + self.set_leaderboard_gpu_rate_limit = bot.admin_group.command( + name="set-leaderboard-gpu-rate-limit", description="Set a rate limit for a leaderboard GPU type" + )(self.set_leaderboard_gpu_rate_limit) self._scheduled_cleanup_temp_users.start() @@ -139,9 +141,7 @@ async def creator_check(self, interaction: discord.Interaction) -> bool: return True return False - async def is_creator_check( - self, interaction: discord.Interaction, leaderboard_name: str - ) -> bool: + async def is_creator_check(self, interaction: discord.Interaction, leaderboard_name: str) -> bool: with self.bot.leaderboard_db as db: leaderboard_item = db.get_leaderboard(leaderboard_name) if leaderboard_item["creator_id"] == interaction.user.id: @@ -218,9 +218,7 @@ async def leaderboard_create_local( f"Leaderboard '{leaderboard_name}' created.", ) - def _leaderboard_opening_message( - self, leaderboard_name: str, deadline: datetime, description: str - ): + def _leaderboard_opening_message(self, leaderboard_name: str, deadline: datetime, description: str): return f""" # New Leaderboard: {leaderboard_name}\n **Deadline**: {deadline.strftime("%Y-%m-%d %H:%M")}\n @@ -252,7 +250,7 @@ async def leaderboard_create_impl( # noqa: C901 ephemeral=True, ) - if date_value < datetime.now(): + if date_value < datetime.now(timezone.utc): await send_discord_message( interaction, f"Deadline {date_value} has already passed.", @@ -268,9 +266,7 @@ async def leaderboard_create_impl( # noqa: C901 content=self._leaderboard_opening_message( leaderboard_name, date_value, - definition.description[:1500] - if len(definition.description) > 1500 - else definition.description, + definition.description[:1500] if len(definition.description) > 1500 else definition.description, ), auto_archive_duration=10080, # 7 days ) @@ -293,8 +289,7 @@ async def leaderboard_create_impl( # noqa: C901 except discord.Forbidden: await send_discord_message( interaction, - "Error: Bot doesn't have permission to create forum threads." - " Leaderboard was not created.", + "Error: Bot doesn't have permission to create forum threads." " Leaderboard was not created.", ephemeral=True, ) except discord.HTTPException: @@ -328,9 +323,7 @@ async def create_leaderboard_in_db( ) -> bool: if gpu is None: # Ask the user to select GPUs - view = GPUSelectionView( - [gpu.name for gpu in GitHubGPU] + [gpu.name for gpu in ModalGPU] - ) + view = GPUSelectionView([gpu.name for gpu in GitHubGPU] + [gpu.name for gpu in ModalGPU]) await send_discord_message( interaction, @@ -368,9 +361,7 @@ async def create_leaderboard_in_db( @discord.app_commands.describe(leaderboard_name="Name of the leaderboard") @discord.app_commands.autocomplete(leaderboard_name=leaderboard_name_autocomplete) @with_error_handling - async def delete_leaderboard( - self, interaction: discord.Interaction, leaderboard_name: str, force: bool = False - ): + async def delete_leaderboard(self, interaction: discord.Interaction, leaderboard_name: str, force: bool = False): is_admin = await self.admin_check(interaction) is_creator = await self.creator_check(interaction) is_creator_of_leaderboard = await self.is_creator_check(interaction, leaderboard_name) @@ -391,9 +382,7 @@ async def delete_leaderboard( ) return - modal = DeleteConfirmationModal( - "leaderboard", leaderboard_name, self.bot.leaderboard_db, force=force - ) + modal = DeleteConfirmationModal("leaderboard", leaderboard_name, self.bot.leaderboard_db, force=force) forum_channel = self.bot.get_channel(self.bot.leaderboard_forum_id) @@ -404,7 +393,7 @@ async def delete_leaderboard( if threads: thread = threads[0] new_name = ( - f"{leaderboard_name} - archived at {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}" + f"{leaderboard_name} - archived at {datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S')} UTC" ) await thread.edit(name=new_name, archived=True) @@ -459,9 +448,7 @@ async def no_delete(): reject_text="Keep", reject_callback=no_delete, ) - await send_discord_message( - interaction, "# Attention\nYou are about to **delete** the following submission:\n" - ) + await send_discord_message(interaction, "# Attention\nYou are about to **delete** the following submission:\n") await send_discord_message(interaction, msg, files=files) await send_discord_message( interaction, @@ -482,9 +469,7 @@ async def stop(self, interaction: discord.Interaction): return self.bot.backend.accepts_jobs = False - await send_discord_message( - interaction, "Bot will refuse all future submissions!", ephemeral=True - ) + await send_discord_message(interaction, "Bot will refuse all future submissions!", ephemeral=True) @with_error_handling async def start(self, interaction: discord.Interaction): @@ -498,8 +483,40 @@ async def start(self, interaction: discord.Interaction): return self.bot.backend.accepts_jobs = True + await send_discord_message(interaction, "Bot will accept submissions again!", ephemeral=True) + + @app_commands.describe( + leaderboard_name="Name of the leaderboard", + gpu_type="Name of the GPU type", + rate_limit_seconds="Rate limit in seconds", + ) + @with_error_handling + async def set_leaderboard_gpu_rate_limit( + self, + interaction: discord.Interaction, + leaderboard_name: str, + gpu_type: str, + rate_limit_seconds: int, + ): + await interaction.response.defer(ephemeral=True) + is_admin = await self.admin_check(interaction) + if not is_admin: + await send_discord_message( + interaction, + "You need to have Admin permissions to run this command", + ephemeral=True, + ) + return + if rate_limit_seconds <= 0: + rate_limit_seconds = None + + with self.bot.leaderboard_db as db: + db.set_leaderboard_gpu_rate_limit(leaderboard_name, gpu_type, rate_limit_seconds) + await send_discord_message( - interaction, "Bot will accept submissions again!", ephemeral=True + interaction, + f"Leaderboard GPU rate limit set to {rate_limit_seconds} seconds for {gpu_type} on {leaderboard_name}", + ephemeral=True, ) @app_commands.describe( @@ -643,9 +660,7 @@ async def _create_update_plan( # noqa: C901 # we can only change things that have no bearing on existing # runs (like description and templates) if ot.files != new_task.files: - file_list = set.symmetric_difference( - set(ot.files.keys()), set(new_task.files) - ) + file_list = set.symmetric_difference(set(ot.files.keys()), set(new_task.files)) if len(file_list) != 0: await send_discord_message( interaction, @@ -653,9 +668,7 @@ async def _create_update_plan( # noqa: C901 f" is currently not possible. File list difference: {file_list}", ) else: - diff_files = { - key for key in ot.files if ot.files[key] != new_task.files[key] - } + diff_files = {key for key in ot.files if ot.files[key] != new_task.files[key]} await send_discord_message( interaction, f"Changing task files of existing problem `{name}`" @@ -692,9 +705,7 @@ async def _create_update_plan( # noqa: C901 return update_list, create_list - async def update_competition( - self, interaction: discord.Interaction, spec_file: Path, force: bool = False - ): + async def update_competition(self, interaction: discord.Interaction, spec_file: Path, force: bool = False): try: root = spec_file.parent with open(spec_file) as f: @@ -703,9 +714,7 @@ async def update_competition( header = f"Handling `{competition['name']}`..." await send_discord_message(interaction, header) - update_list, create_list = await self._create_update_plan( - interaction, competition, root, force - ) + update_list, create_list = await self._create_update_plan(interaction, competition, root, force) # OK, now we know what we want to do plan = "" @@ -738,9 +747,7 @@ async def update_competition( for entry in update_list: with self.bot.leaderboard_db as db: task = make_task_definition(root / entry["directory"]) - db.update_leaderboard( - entry["name"], parse_deadline(entry["deadline"]), task - ) + db.update_leaderboard(entry["name"], parse_deadline(entry["deadline"]), task) new_lb: LeaderboardItem = db.get_leaderboard(entry["name"]) forum_id = new_lb["forum_id"] @@ -753,9 +760,7 @@ async def update_competition( ) ) except (discord.errors.NotFound, discord.errors.HTTPException): - logger.warning( - "Could not find forum thread %s for lb %s", forum_id, entry["name"] - ) + logger.warning("Could not find forum thread %s for lb %s", forum_id, entry["name"]) pass header += " DONE" @@ -802,9 +807,7 @@ async def resync(self, interaction: discord.Interaction): logger.error(f"Error in resync command: {str(e)}", exc_info=True) await send_discord_message(interaction, f"Error: {str(e)}") else: - await send_discord_message( - interaction, "You need administrator permissions to use this command" - ) + await send_discord_message(interaction, "You need administrator permissions to use this command") # admin version of this command; less restricted @discord.app_commands.describe(submission_id="ID of the submission") @@ -828,9 +831,7 @@ async def get_submission_by_id( # allowed/possible to see submission if sub is None: - await send_discord_message( - interaction, f"Submission {submission_id} does not exist", ephemeral=True - ) + await send_discord_message(interaction, f"Submission {submission_id} does not exist", ephemeral=True) return msg, files = self._make_submission_message(submission_id, sub) @@ -917,9 +918,7 @@ async def get_user_names(self, interaction: discord.Interaction): await send_discord_message(interaction, error_message) @app_commands.describe(attachment="The JSON file containing user ID to name mapping") - async def update_user_names( - self, interaction: discord.Interaction, attachment: discord.Attachment - ): + async def update_user_names(self, interaction: discord.Interaction, attachment: discord.Attachment): """Update the database with user names from a JSON file""" if not await self.admin_check(interaction): await send_discord_message( @@ -932,9 +931,7 @@ async def update_user_names( try: if not attachment.filename.endswith(".json"): - await send_discord_message( - interaction, "Please attach a JSON file with .json extension." - ) + await send_discord_message(interaction, "Please attach a JSON file with .json extension.") return json_content = await attachment.read() @@ -982,9 +979,7 @@ async def update_user_names( ) except json.JSONDecodeError: - await send_discord_message( - interaction, "Invalid JSON format in the attached file.", ephemeral=True - ) + await send_discord_message(interaction, "Invalid JSON format in the attached file.", ephemeral=True) except Exception as e: error_message = f"Error updating database with user names: {str(e)}" logger.error(error_message, exc_info=True) @@ -1007,11 +1002,7 @@ async def set_forum_ids(self, interaction: discord.Interaction): threads = [thread for thread in forum_channel.threads if thread.name == name] if len(threads) == 0: # is it an archived thread? - threads = [ - thread - async for thread in forum_channel.archived_threads() - if thread.name == name - ] + threads = [thread async for thread in forum_channel.archived_threads() if thread.name == name] if len(threads) != 1: await send_discord_message( interaction, f"Could not set forum thread for {name}", ephemeral=True diff --git a/src/kernelbot/cogs/misc_cog.py b/src/kernelbot/cogs/misc_cog.py index 09e78bc8..28c6d93e 100644 --- a/src/kernelbot/cogs/misc_cog.py +++ b/src/kernelbot/cogs/misc_cog.py @@ -6,7 +6,7 @@ from discord import app_commands from discord.ext import commands -from kernelbot.discord_utils import send_discord_message +from kernelbot.discord_utils import leaderboard_name_autocomplete, send_discord_message from kernelbot.env import env from libkernelbot.utils import setup_logging @@ -25,6 +25,29 @@ async def ping(self, interaction: discord.Interaction): """Simple ping command to check if the bot is responsive""" await send_discord_message(interaction, "pong") + @app_commands.describe(leaderboard_name="Name of the leaderboard") + @app_commands.autocomplete(leaderboard_name=leaderboard_name_autocomplete) + @app_commands.command(name="get-leaderboard-rate-limit") + async def get_leaderboard_rate_limit(self, interaction: discord.Interaction, leaderboard_name: str): + """Get the rate limit for a leaderboard GPU type""" + await interaction.response.defer(ephemeral=True) + + with self.bot.leaderboard_db as db: + rate_limits = db.get_leaderboard_rate_limits(leaderboard_name) + + rate_limit_msg = f"## Leaderboard: {leaderboard_name}\n" + + def format_msg(gpu_type: str, rate_limit: int): + nonlocal rate_limit_msg + rate_limit_str = f"{rate_limit} seconds" if rate_limit is not None else "no rate limit" + rate_limit_msg += f"- {gpu_type}: {rate_limit_str}\n" + return rate_limit_msg + + for gpu_type, rate_limit in rate_limits.items(): + format_msg(gpu_type, rate_limit) + + await send_discord_message(interaction, rate_limit_msg, ephemeral=True) + @app_commands.command(name="verifydb") async def verify_db(self, interaction: discord.Interaction): """Command to verify database connectivity""" @@ -42,9 +65,7 @@ async def verify_db(self, interaction: discord.Interaction): result = cursor.fetchone() if result: random_value = result[0] - await send_discord_message( - interaction, f"Your lucky number is {random_value}." - ) + await send_discord_message(interaction, f"Your lucky number is {random_value}.") else: await send_discord_message(interaction, "No result returned.") except Exception as e: @@ -57,8 +78,7 @@ async def get_api_url(self, interaction: discord.Interaction): if not os.environ.get("HEROKU_APP_DEFAULT_DOMAIN_NAME"): await send_discord_message( interaction, - "No `HEROKU_APP_DEFAULT_DOMAIN_NAME` present," - " are you sure you aren't running locally?", + "No `HEROKU_APP_DEFAULT_DOMAIN_NAME` present, are you sure you aren't running locally?", ephemeral=True, ) else: diff --git a/src/kernelbot/cogs/verify_run_cog.py b/src/kernelbot/cogs/verify_run_cog.py index 53102682..1a6c2dd9 100644 --- a/src/kernelbot/cogs/verify_run_cog.py +++ b/src/kernelbot/cogs/verify_run_cog.py @@ -50,14 +50,10 @@ async def trigger_run(self, interaction: discord.Interaction, gpu: GPU, reporter submit_leaderboard = self.bot.backend.submit_leaderboard if lang == "py": - sub_code = create_mock_attachment( - "submission.py", Path("examples/identity_py/submission.py").read_text() - ) + sub_code = create_mock_attachment("submission.py", Path("examples/identity_py/submission.py").read_text()) leaderboard = make_task_definition("examples/identity_py") else: - sub_code = create_mock_attachment( - "test.cu", Path("examples/identity_cuda/submission.cu").read_text() - ) + sub_code = create_mock_attachment("test.cu", Path("examples/identity_cuda/submission.cu").read_text()) leaderboard = make_task_definition("examples/identity_cuda") return await submit_leaderboard( @@ -177,9 +173,7 @@ async def verify_modal_run( ] ) @with_error_handling - async def verify_task( - self, interaction: discord.Interaction, task: str, mode: Choice[str] = None - ): + async def verify_task(self, interaction: discord.Interaction, task: str, mode: Choice[str] = None): directory = Path(env.PROBLEM_DEV_DIR) / task if not directory.resolve().is_relative_to(Path.cwd() / env.PROBLEM_DEV_DIR): await send_discord_message(interaction, f"Invalid path {directory.resolve()}") @@ -206,7 +200,7 @@ async def verify_task( db.create_leaderboard( { "name": lb_name, - "deadline": datetime.datetime.now() + datetime.timedelta(days=1), + "deadline": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(days=1), "task": task, "gpu_types": "T4", "creator_id": interaction.user.id, @@ -218,9 +212,7 @@ async def verify_task( reports = [] for sub in directory.glob("solutions/*/*"): for mode in modes: - submissions.append( - self.verify_submission(interaction, lb_name, sub, mode, reports) - ) + submissions.append(self.verify_submission(interaction, lb_name, sub, mode, reports)) await asyncio.gather(*submissions) except Exception as E: logger.exception("Error in LB test", exc_info=E) @@ -313,6 +305,4 @@ async def verify_runs(self, interaction: discord.Interaction): except Exception as e: logger.error(f"Error starting verification runs: {e}", exc_info=True) - await send_discord_message( - interaction, f"❌ Problem performing verification runs: {str(e)}" - ) + await send_discord_message(interaction, f"❌ Problem performing verification runs: {str(e)}") diff --git a/src/libkernelbot/backend.py b/src/libkernelbot/backend.py index f3b68bb0..79a2f2f0 100644 --- a/src/libkernelbot/backend.py +++ b/src/libkernelbot/backend.py @@ -1,6 +1,6 @@ import asyncio import copy -from datetime import datetime +from datetime import datetime, timezone from types import SimpleNamespace from typing import Optional @@ -67,7 +67,7 @@ async def submit_full( file_name=req.file_name, code=req.code, user_id=req.user_id, - time=datetime.now(), + time=datetime.now(timezone.utc), user_name=req.user_name, ) selected_gpus = [get_gpu_by_name(gpu) for gpu in req.gpus] @@ -100,9 +100,7 @@ async def submit_full( ) for gpu in selected_gpus ] - await reporter.show( - f"Submission **{sub_id}**: `{req.file_name}` for `{req.leaderboard}`" - ) + await reporter.show(f"Submission **{sub_id}**: `{req.file_name}` for `{req.leaderboard}`") results = await asyncio.gather(*tasks) finally: with self.db as db: @@ -191,9 +189,7 @@ async def handle_submission( if successful, returns the result of the run. """ launcher = self.launcher_map[gpu_type.value] - config = build_task_config( - task=task, submission_content=code, arch=self._get_arch(gpu_type), mode=mode - ) + config = build_task_config(task=task, submission_content=code, arch=self._get_arch(gpu_type), mode=mode) logger.info("submitting task to runner %s", launcher.name) @@ -206,9 +202,7 @@ async def handle_submission( else: await reporter.update_title(reporter.title + " ✅ success") - short_report = make_short_report( - result.runs, full=mode in [SubmissionMode.PRIVATE, SubmissionMode.LEADERBOARD] - ) + short_report = make_short_report(result.runs, full=mode in [SubmissionMode.PRIVATE, SubmissionMode.LEADERBOARD]) stream_msg = ( """ diff --git a/src/libkernelbot/leaderboard_db.py b/src/libkernelbot/leaderboard_db.py index f1d07cb8..b4bfb95d 100644 --- a/src/libkernelbot/leaderboard_db.py +++ b/src/libkernelbot/leaderboard_db.py @@ -90,7 +90,14 @@ def create_leaderboard( VALUES (%s, %s, %s, %s, %s, %s) RETURNING id """, - (name, deadline, task.to_str(), creator_id, forum_id, definition.description), + ( + name, + deadline.astimezone(datetime.timezone.utc), + task.to_str(), + creator_id, + forum_id, + definition.description, + ), ) leaderboard_id = self.cursor.fetchone()[0] @@ -123,15 +130,11 @@ def create_leaderboard( except psycopg2.Error as e: logger.exception("Error in leaderboard creation.", exc_info=e) if isinstance(e, psycopg2.errors.UniqueViolation): - raise KernelBotError( - f"Error: Tried to create a leaderboard '{name}' that already exists." - ) from e + raise KernelBotError(f"Error: Tried to create a leaderboard '{name}' that already exists.") from e self.connection.rollback() # Ensure rollback if error occurs raise KernelBotError("Error in leaderboard creation.") from e - def update_leaderboard( - self, name, deadline: datetime.datetime, definition: LeaderboardDefinition - ): + def update_leaderboard(self, name, deadline: datetime.datetime, definition: LeaderboardDefinition): task = definition.task try: lb_id = self.get_leaderboard_id(name) @@ -259,9 +262,7 @@ def validate_identity( (identifier,), ) row = self.cursor.fetchone() - return ( - {"user_id": row[0], "user_name": row[1], "id_type": id_type.value} if row else None - ) + return {"user_id": row[0], "user_name": row[1], "id_type": id_type.value} if row else None except psycopg2.Error as e: self.connection.rollback() logger.exception("Error validating %s %s", human_label, identifier, exc_info=e) @@ -447,15 +448,9 @@ def create_submission_run( submission, mode, ) - raise KernelBotError( - "Internal error: Attempted to add run, " - "but submission was already marked as done." - ) + raise KernelBotError("Internal error: Attempted to add run, but submission was already marked as done.") - meta = { - k: result.__dict__[k] - for k in ["stdout", "stderr", "success", "exit_code", "command", "duration"] - } + meta = {k: result.__dict__[k] for k in ["stdout", "stderr", "success", "exit_code", "command", "duration"]} self.cursor.execute( """ INSERT INTO leaderboard.runs (submission_id, start_time, end_time, mode, @@ -494,7 +489,7 @@ def get_leaderboard_names(self, active_only: bool = False) -> list[str]: if active_only: self.cursor.execute( "SELECT name FROM leaderboard.leaderboard WHERE leaderboard.deadline > %s", - (datetime.datetime.now().astimezone(datetime.timezone.utc),), + (datetime.datetime.now(datetime.timezone.utc),), ) else: self.cursor.execute("SELECT name FROM leaderboard.leaderboard") @@ -512,9 +507,7 @@ def get_leaderboards(self) -> list["LeaderboardItem"]: leaderboards = [] for lb in lbs: - self.cursor.execute( - "SELECT * from leaderboard.gpu_type WHERE leaderboard_id = %s", [lb[0]] - ) + self.cursor.execute("SELECT * from leaderboard.gpu_type WHERE leaderboard_id = %s", [lb[0]]) gpu_types = [x[1] for x in self.cursor.fetchall()] leaderboards.append( @@ -707,9 +700,7 @@ def get_leaderboard_submissions( # did we specify a valid GPU? gpus = self.get_leaderboard_gpu_types(leaderboard_name) if gpu_name not in gpus: - raise KernelBotError( - f"Invalid GPU type '{gpu_name}' for leaderboard '{leaderboard_name}'" - ) + raise KernelBotError(f"Invalid GPU type '{gpu_name}' for leaderboard '{leaderboard_name}'") return result @@ -909,23 +900,27 @@ def get_user_submissions( sub_id = run_row[0] if sub_id not in runs_by_submission: runs_by_submission[sub_id] = [] - runs_by_submission[sub_id].append({ - "gpu_type": run_row[1], - "score": run_row[2], - }) + runs_by_submission[sub_id].append( + { + "gpu_type": run_row[1], + "score": run_row[2], + } + ) # Build result with runs grouped by submission results = [] for row in submissions: sub_id = row[0] - results.append({ - "id": sub_id, - "leaderboard_name": row[1], - "file_name": row[2], - "submission_time": row[3], - "done": row[4], - "runs": runs_by_submission.get(sub_id, []), - }) + results.append( + { + "id": sub_id, + "leaderboard_name": row[1], + "file_name": row[2], + "submission_time": row[3], + "done": row[4], + "runs": runs_by_submission.get(sub_id, []), + } + ) return results except psycopg2.Error as e: self.connection.rollback() @@ -1030,9 +1025,7 @@ def get_leaderboard_submission_count( # did we specify a valid GPU? gpus = self.get_leaderboard_gpu_types(leaderboard_name) if gpu_name not in gpus: - raise KernelBotError( - f"Invalid GPU type '{gpu_name}' for leaderboard '{leaderboard_name}'" - ) + raise KernelBotError(f"Invalid GPU type '{gpu_name}' for leaderboard '{leaderboard_name}'") return count @@ -1139,9 +1132,7 @@ def reset_user_from_cli(self, user_id: str, cli_id: str, auth_provider: str): (user_id,), ) if not self.cursor.fetchone(): - raise Exception( - "User not found. Please use the register command to create an account." - ) + raise Exception("User not found. Please use the register command to create an account.") self.cursor.execute( """ @@ -1214,6 +1205,102 @@ def validate_cli_id(self, cli_id: str) -> Optional[dict[str, str]]: logger.exception("Error validating CLI ID %s", cli_id, exc_info=e) raise KernelBotError("Error validating CLI ID") from e + def is_user_rate_limited(self, user_id: int, leaderboard_id: int, gpu_type: str) -> bool: + try: + self.cursor.execute( + """ + SELECT rate_limit_seconds + FROM leaderboard.gpu_type + WHERE leaderboard_id = %s AND gpu_type = %s + """, + (leaderboard_id, gpu_type), + ) + row = self.cursor.fetchone() + if row is None or row[0] is None: + return False, None + + rate_limit_seconds = row[0] + self.cursor.execute( + """ + SELECT submission_time + FROM leaderboard.submission + WHERE user_id = %s + AND leaderboard_id = %s + AND submission_time > NOW() - make_interval(secs => %s) + ORDER BY submission_time DESC + LIMIT 1 + """, + (str(user_id), leaderboard_id, rate_limit_seconds), + ) + last = self.cursor.fetchone() + if last is not None: + last_time = last[0] + return True, ( + f"Rate limit exceeded for {gpu_type}`. " + f"You can submit once every {rate_limit_seconds} seconds. " + f"Last submission: {last_time.strftime('%Y-%m-%d %H:%M:%S UTC')}." + ) + except psycopg2.Error as e: + self.connection.rollback() + logger.exception("Error checking rate limit for user %s", user_id, exc_info=e) + raise KernelBotError("Error checking submission rate limit") from e + return False, None + + def is_user_allowed_to_submit(self, user_id: int, leaderboard: str, gpus: list) -> tuple[bool, str | None]: + """ + Check if a user is allowed to submit to a leaderboard. + Enforces per-(leaderboard, gpu) rate limits stored in leaderboard.gpu_type.rate_limit_seconds. + """ + try: + lb_id = self.get_leaderboard_id(leaderboard) + + for gpu in gpus: + is_user_rate_limited, reason = self.is_user_rate_limited(user_id, lb_id, gpu) + if is_user_rate_limited: + return False, reason + + return True, None + except psycopg2.Error as e: + self.connection.rollback() + logger.exception("Error checking rate limit for user %s", user_id, exc_info=e) + raise KernelBotError("Error checking submission rate limit") from e + + def set_leaderboard_gpu_rate_limit(self, leaderboard_name: str, gpu_type: str, rate_limit_seconds: int): + try: + leaderboard_id = self.get_leaderboard_id(leaderboard_name) + self.cursor.execute( + """ + UPDATE leaderboard.gpu_type + SET rate_limit_seconds = %s + WHERE leaderboard_id = %s AND gpu_type = %s + """, + (rate_limit_seconds, leaderboard_id, gpu_type), + ) + + self.connection.commit() + except psycopg2.Error as e: + self.connection.rollback() + logger.exception( + "Error setting leaderboard GPU rate limit for %s %s", leaderboard_name, gpu_type, exc_info=e + ) + raise KernelBotError("Error setting leaderboard GPU rate limit") from e + + def get_leaderboard_rate_limits(self, leaderboard_name: str) -> dict[str, int]: + try: + leaderboard_id = self.get_leaderboard_id(leaderboard_name) + self.cursor.execute( + """ + SELECT gpu_type, rate_limit_seconds + FROM leaderboard.gpu_type + WHERE leaderboard_id = %s + """, + (leaderboard_id,), + ) + return {x[0]: x[1] for x in self.cursor.fetchall()} + except psycopg2.Error as e: + logger.exception("Error getting leaderboard rate limits for %s", leaderboard_name, exc_info=e) + self.connection.rollback() + class LeaderboardDoesNotExist(KernelBotError): def __init__(self, name: str): diff --git a/src/libkernelbot/run_eval.py b/src/libkernelbot/run_eval.py index aec59f95..917a7577 100644 --- a/src/libkernelbot/run_eval.py +++ b/src/libkernelbot/run_eval.py @@ -273,9 +273,7 @@ def compile_cuda_script( # # noqa: C901 print_("[Compiling]") try: - compile_process = subprocess.run( - command, capture_output=True, text=True, check=True, timeout=Timeout.COMPILE - ) + compile_process = subprocess.run(command, capture_output=True, text=True, check=True, timeout=Timeout.COMPILE) except subprocess.CalledProcessError as e: return CompileResult( nvcc_found=True, @@ -357,10 +355,7 @@ def run_program( result_dict[key.strip()] = value.strip() return RunResult( - success=( - run_process.returncode == ExitCode.SUCCESS - or run_process.returncode == ExitCode.VALIDATE_FAIL - ), + success=(run_process.returncode == ExitCode.SUCCESS or run_process.returncode == ExitCode.VALIDATE_FAIL), passed=result_dict.get("check", None) == "pass", command=_make_cmd(run_process.args), stdout=_limit_length(run_process.stdout), @@ -474,9 +469,7 @@ def profile_program_ncu( "--", ] + call - run_result = run_program( - call, seed=seed, timeout=timeout, multi_gpu=multi_gpu, extra_env={"POPCORN_NCU": "1"} - ) + run_result = run_program(call, seed=seed, timeout=timeout, multi_gpu=multi_gpu, extra_env={"POPCORN_NCU": "1"}) profile_result = None try: @@ -494,9 +487,7 @@ def profile_program_ncu( ] report = subprocess.check_output(ncu_cmd, text=True) report = _filter_ncu_report(report, get_tables) - run_result.result["benchmark.0.report"] = base64.b64encode(report.encode("utf-8")).decode( - "utf-8" - ) + run_result.result["benchmark.0.report"] = base64.b64encode(report.encode("utf-8")).decode("utf-8") except subprocess.CalledProcessError: pass @@ -602,9 +593,7 @@ def make_system_info() -> SystemInfo: # noqa: C901 # try again for HIP try: rocm_info = json.loads( - subprocess.check_output( - ["rocm-smi", "--showproductname", "--json"], encoding="utf-8" - ) + subprocess.check_output(["rocm-smi", "--showproductname", "--json"], encoding="utf-8") ) if len(rocm_info) > 0: info.gpu = next(rocm_info.__iter__())["Card Series"] @@ -668,7 +657,7 @@ def run_cuda_script( # # noqa: C901 Returns: tuple[CompileResult, RunResult]: CUDA compile/eval result information """ - start = datetime.datetime.now() + start = datetime.datetime.now(datetime.timezone.utc) try: # Write submission files to directory _create_files(sources) @@ -687,7 +676,7 @@ def run_cuda_script( # # noqa: C901 if not compile_result.success: return EvalResult( start=start, - end=datetime.datetime.now(), + end=datetime.datetime.now(datetime.timezone.utc), compilation=compile_result, run=None, profile=None, @@ -704,7 +693,7 @@ def run_cuda_script( # # noqa: C901 run_result, profile_result = run_single_evaluation(["./eval.out"], **kwargs) return EvalResult( start=start, - end=datetime.datetime.now(), + end=datetime.datetime.now(datetime.timezone.utc), compilation=compile_result, run=run_result, profile=profile_result, @@ -727,7 +716,7 @@ def run_pytorch_script( # noqa: C901 Returns: RunResult """ - start = datetime.datetime.now() + start = datetime.datetime.now(datetime.timezone.utc) try: assert main in sources.keys() @@ -767,7 +756,7 @@ def run_pytorch_script( # noqa: C901 return EvalResult( start=start, - end=datetime.datetime.now(), + end=datetime.datetime.now(datetime.timezone.utc), compilation=comp, run=run, profile=profile, @@ -847,9 +836,7 @@ def run_config(config: dict): "multi_gpu": config.get("multi_gpu", False), } if config["lang"] == "py": - runner = functools.partial( - run_pytorch_script, sources=config["sources"], main=config["main"] - ) + runner = functools.partial(run_pytorch_script, sources=config["sources"], main=config["main"]) elif config["lang"] == "cu": runner = functools.partial( run_cuda_script, diff --git a/src/libkernelbot/submission.py b/src/libkernelbot/submission.py index 805f7435..2f8b04aa 100644 --- a/src/libkernelbot/submission.py +++ b/src/libkernelbot/submission.py @@ -39,13 +39,9 @@ class ProcessedSubmissionRequest(SubmissionRequest): task_gpus: list -def prepare_submission( - req: SubmissionRequest, backend: "KernelBackend" -) -> ProcessedSubmissionRequest: +def prepare_submission(req: SubmissionRequest, backend: "KernelBackend") -> ProcessedSubmissionRequest: if not backend.accepts_jobs: - raise KernelBotError( - "The bot is currently not accepting any new submissions, please try again later." - ) + raise KernelBotError("The bot is currently not accepting any new submissions, please try again later.") if profanity.contains_profanity(req.file_name): raise KernelBotError("Please provide a non-rude filename") @@ -72,12 +68,17 @@ def prepare_submission( task_gpu_list = "".join([f" * {t}\n" for t in task_gpus]) raise KernelBotError( - f"GPU {g} not available for `{req.leaderboard}`\n" - f"Choose one of: {task_gpu_list}", + f"GPU {g} not available for `{req.leaderboard}`\n" f"Choose one of: {task_gpu_list}", ) elif len(task_gpus) == 1: req.gpus = task_gpus + with backend.db as db: + is_user_allowed, potential_reason = db.is_user_allowed_to_submit(req.user_id, req.leaderboard, req.gpus) + + if not is_user_allowed: + raise KernelBotError(potential_reason) + return ProcessedSubmissionRequest( **dataclasses.asdict(req), task=leaderboard["task"], @@ -92,8 +93,7 @@ def check_deadline(leaderboard: LeaderboardItem): if now > deadline: raise KernelBotError( - f"The deadline to submit to {leaderboard['name']} has passed.\n" - f"It was {deadline} and today is {now}." + f"The deadline to submit to {leaderboard['name']} has passed.\n" f"It was {deadline} and today is {now}." ) @@ -174,14 +174,11 @@ def compute_score(result: FullResult, task: LeaderboardTask, submission_id: int) if task.ranking_by == RankCriterion.LAST: if num_benchmarks != 1: logger.error( - "Ranked submission error for submission %d ranking_by is `last`, " - "but got %d benchmarks", + "Ranked submission error for submission %d ranking_by is `last`, " "but got %d benchmarks", submission_id, num_benchmarks, ) - raise KernelBotError( - f"Expected submission to have exactly one benchmark, got {num_benchmarks}." - ) + raise KernelBotError(f"Expected submission to have exactly one benchmark, got {num_benchmarks}.") score = float(result.runs["leaderboard"].run.result["benchmark.0.mean"]) / 1e9 else: scores = [] diff --git a/src/libkernelbot/utils.py b/src/libkernelbot/utils.py index fc1a3356..d345d985 100644 --- a/src/libkernelbot/utils.py +++ b/src/libkernelbot/utils.py @@ -1,7 +1,7 @@ import logging import os import subprocess -from datetime import datetime +from datetime import datetime, timezone from typing import Any, Optional @@ -63,7 +63,7 @@ def parse_deadline(deadline: str) -> Optional[datetime]: """ for fmt in ("%Y-%m-%d %H:%M", "%Y-%m-%d"): try: - return datetime.strptime(deadline, fmt) + return datetime.strptime(deadline, fmt).replace(tzinfo=timezone.utc) except ValueError: continue return None diff --git a/src/migrations/20260207_01_EW7ve-leaderboard-rate-limits.py b/src/migrations/20260207_01_EW7ve-leaderboard-rate-limits.py new file mode 100644 index 00000000..0f76889a --- /dev/null +++ b/src/migrations/20260207_01_EW7ve-leaderboard-rate-limits.py @@ -0,0 +1,14 @@ +""" +leaderboard-rate-limits +""" + +from yoyo import step + +__depends__ = {"20260108_01_gzSm3-add-submission-status"} + +steps = [ + step( + "ALTER TABLE leaderboard.gpu_type ADD COLUMN rate_limit_seconds INTEGER DEFAULT NULL", + "ALTER TABLE leaderboard.gpu_type DROP COLUMN rate_limit_seconds", + ), +] From cd046a03b74e1e705c5a6fe1e8daf07671ce8a2f Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Sat, 7 Feb 2026 22:21:54 +0100 Subject: [PATCH 2/3] Fixes --- src/kernelbot/api/main.py | 2 +- src/kernelbot/cogs/admin_cog.py | 5 +- src/libkernelbot/leaderboard_db.py | 24 +-- tests/test_admin_api.py | 261 +++++++++++++---------------- tests/test_leaderboard_db.py | 259 ++++++++++++++++++++++------ tests/test_submission.py | 108 ++++++++++-- 6 files changed, 433 insertions(+), 226 deletions(-) diff --git a/src/kernelbot/api/main.py b/src/kernelbot/api/main.py index f26ca91b..afba8f43 100644 --- a/src/kernelbot/api/main.py +++ b/src/kernelbot/api/main.py @@ -249,7 +249,7 @@ async def cli_auth(auth_provider: str, code: str, state: str, db_context=Depends if not api_base_url: raise HTTPException( status_code=500, - detail="Redirect URI base not configured.Set HEROKU_APP_DEFAULT_DOMAIN_NAME or POPCORN_API_URL.", + detail="Redirect URI base not configured. Set HEROKU_APP_DEFAULT_DOMAIN_NAME or POPCORN_API_URL.", ) redirect_uri_base = api_base_url.rstrip("/") redirect_uri = f"https://{redirect_uri_base}/auth/cli/{auth_provider}" diff --git a/src/kernelbot/cogs/admin_cog.py b/src/kernelbot/cogs/admin_cog.py index c55fbf22..44502bd8 100644 --- a/src/kernelbot/cogs/admin_cog.py +++ b/src/kernelbot/cogs/admin_cog.py @@ -289,7 +289,7 @@ async def leaderboard_create_impl( # noqa: C901 except discord.Forbidden: await send_discord_message( interaction, - "Error: Bot doesn't have permission to create forum threads." " Leaderboard was not created.", + "Error: Bot doesn't have permission to create forum threads. Leaderboard was not created.", ephemeral=True, ) except discord.HTTPException: @@ -512,10 +512,11 @@ async def set_leaderboard_gpu_rate_limit( with self.bot.leaderboard_db as db: db.set_leaderboard_gpu_rate_limit(leaderboard_name, gpu_type, rate_limit_seconds) + rate_limit_str = f"{rate_limit_seconds} seconds" if rate_limit_seconds is not None else "no rate limit" await send_discord_message( interaction, - f"Leaderboard GPU rate limit set to {rate_limit_seconds} seconds for {gpu_type} on {leaderboard_name}", + f"Leaderboard GPU rate limit set to {rate_limit_str} for {gpu_type} on {leaderboard_name}", ephemeral=True, ) diff --git a/src/libkernelbot/leaderboard_db.py b/src/libkernelbot/leaderboard_db.py index b4bfb95d..0106b545 100644 --- a/src/libkernelbot/leaderboard_db.py +++ b/src/libkernelbot/leaderboard_db.py @@ -1205,7 +1205,7 @@ def validate_cli_id(self, cli_id: str) -> Optional[dict[str, str]]: logger.exception("Error validating CLI ID %s", cli_id, exc_info=e) raise KernelBotError("Error validating CLI ID") from e - def is_user_rate_limited(self, user_id: int, leaderboard_id: int, gpu_type: str) -> bool: + def is_user_rate_limited(self, user_id: int, leaderboard_id: int, gpu_type: str) -> tuple[bool, str | None]: try: self.cursor.execute( """ @@ -1220,23 +1220,26 @@ def is_user_rate_limited(self, user_id: int, leaderboard_id: int, gpu_type: str) return False, None rate_limit_seconds = row[0] + print(f"GPU type: {gpu_type}, Rate limit: {rate_limit_seconds}") self.cursor.execute( """ - SELECT submission_time - FROM leaderboard.submission - WHERE user_id = %s - AND leaderboard_id = %s - AND submission_time > NOW() - make_interval(secs => %s) - ORDER BY submission_time DESC + SELECT s.submission_time + FROM leaderboard.submission s + JOIN leaderboard.runs r ON r.submission_id = s.id + WHERE s.user_id = %s + AND s.leaderboard_id = %s + AND s.submission_time > NOW() - make_interval(secs => %s) + AND r.runner = %s + ORDER BY s.submission_time DESC LIMIT 1 """, - (str(user_id), leaderboard_id, rate_limit_seconds), + (str(user_id), leaderboard_id, rate_limit_seconds, gpu_type), ) last = self.cursor.fetchone() if last is not None: last_time = last[0] return True, ( - f"Rate limit exceeded for {gpu_type}`. " + f"Rate limit exceeded for `{gpu_type}`. " f"You can submit once every {rate_limit_seconds} seconds. " f"Last submission: {last_time.strftime('%Y-%m-%d %H:%M:%S UTC')}." ) @@ -1265,7 +1268,7 @@ def is_user_allowed_to_submit(self, user_id: int, leaderboard: str, gpus: list) logger.exception("Error checking rate limit for user %s", user_id, exc_info=e) raise KernelBotError("Error checking submission rate limit") from e - def set_leaderboard_gpu_rate_limit(self, leaderboard_name: str, gpu_type: str, rate_limit_seconds: int): + def set_leaderboard_gpu_rate_limit(self, leaderboard_name: str, gpu_type: str, rate_limit_seconds: int | None): try: leaderboard_id = self.get_leaderboard_id(leaderboard_name) self.cursor.execute( @@ -1300,6 +1303,7 @@ def get_leaderboard_rate_limits(self, leaderboard_name: str) -> dict[str, int]: except psycopg2.Error as e: logger.exception("Error getting leaderboard rate limits for %s", leaderboard_name, exc_info=e) self.connection.rollback() + raise KernelBotError("Error getting leaderboard rate limits") from e class LeaderboardDoesNotExist(KernelBotError): diff --git a/tests/test_admin_api.py b/tests/test_admin_api.py index ada13823..0976c3ff 100644 --- a/tests/test_admin_api.py +++ b/tests/test_admin_api.py @@ -1,6 +1,5 @@ """Tests for admin API endpoints.""" -import datetime from unittest.mock import MagicMock, patch import pytest @@ -20,8 +19,9 @@ def mock_backend(): def test_client(mock_backend): """Create a test client with mocked backend.""" # Patch env before importing the app - with patch.dict('os.environ', {'ADMIN_TOKEN': 'test_token'}): + with patch.dict("os.environ", {"ADMIN_TOKEN": "test_token"}): from kernelbot.api.main import app, init_api + init_api(mock_backend) yield TestClient(app) @@ -37,19 +37,13 @@ def test_admin_requires_auth_header(self, test_client): def test_admin_rejects_invalid_token(self, test_client): """Admin endpoints reject invalid tokens.""" - response = test_client.post( - "/admin/start", - headers={"Authorization": "Bearer wrong_token"} - ) + response = test_client.post("/admin/start", headers={"Authorization": "Bearer wrong_token"}) assert response.status_code == 401 assert response.json()["detail"] == "Invalid admin token" def test_admin_accepts_valid_token(self, test_client, mock_backend): """Admin endpoints accept valid tokens.""" - response = test_client.post( - "/admin/start", - headers={"Authorization": "Bearer test_token"} - ) + response = test_client.post("/admin/start", headers={"Authorization": "Bearer test_token"}) assert response.status_code == 200 assert response.json()["status"] == "ok" assert mock_backend.accepts_jobs is True @@ -61,10 +55,7 @@ class TestAdminStartStop: def test_admin_start(self, test_client, mock_backend): """POST /admin/start enables job acceptance.""" mock_backend.accepts_jobs = False - response = test_client.post( - "/admin/start", - headers={"Authorization": "Bearer test_token"} - ) + response = test_client.post("/admin/start", headers={"Authorization": "Bearer test_token"}) assert response.status_code == 200 assert response.json() == {"status": "ok", "accepts_jobs": True} assert mock_backend.accepts_jobs is True @@ -72,10 +63,7 @@ def test_admin_start(self, test_client, mock_backend): def test_admin_stop(self, test_client, mock_backend): """POST /admin/stop disables job acceptance.""" mock_backend.accepts_jobs = True - response = test_client.post( - "/admin/stop", - headers={"Authorization": "Bearer test_token"} - ) + response = test_client.post("/admin/stop", headers={"Authorization": "Bearer test_token"}) assert response.status_code == 200 assert response.json() == {"status": "ok", "accepts_jobs": False} assert mock_backend.accepts_jobs is False @@ -88,15 +76,14 @@ def test_admin_stats(self, test_client, mock_backend): """GET /admin/stats returns statistics.""" mock_backend.db.__enter__ = MagicMock(return_value=mock_backend.db) mock_backend.db.__exit__ = MagicMock(return_value=None) - mock_backend.db.generate_stats = MagicMock(return_value={ - "num_submissions": 10, - "num_users": 5, - }) - - response = test_client.get( - "/admin/stats", - headers={"Authorization": "Bearer test_token"} + mock_backend.db.generate_stats = MagicMock( + return_value={ + "num_submissions": 10, + "num_users": 5, + } ) + + response = test_client.get("/admin/stats", headers={"Authorization": "Bearer test_token"}) assert response.status_code == 200 data = response.json() assert data["status"] == "ok" @@ -106,15 +93,14 @@ def test_admin_stats_last_day_only(self, test_client, mock_backend): """GET /admin/stats with last_day_only parameter.""" mock_backend.db.__enter__ = MagicMock(return_value=mock_backend.db) mock_backend.db.__exit__ = MagicMock(return_value=None) - mock_backend.db.generate_stats = MagicMock(return_value={ - "num_submissions": 3, - "num_users": 2, - }) - - response = test_client.get( - "/admin/stats?last_day_only=true", - headers={"Authorization": "Bearer test_token"} + mock_backend.db.generate_stats = MagicMock( + return_value={ + "num_submissions": 3, + "num_users": 2, + } ) + + response = test_client.get("/admin/stats?last_day_only=true", headers={"Authorization": "Bearer test_token"}) assert response.status_code == 200 mock_backend.db.generate_stats.assert_called_once_with(True) @@ -126,15 +112,14 @@ def test_get_submission(self, test_client, mock_backend): """GET /admin/submissions/{id} returns submission.""" mock_backend.db.__enter__ = MagicMock(return_value=mock_backend.db) mock_backend.db.__exit__ = MagicMock(return_value=None) - mock_backend.db.get_submission_by_id = MagicMock(return_value={ - "id": 123, - "code": "test code", - }) - - response = test_client.get( - "/admin/submissions/123", - headers={"Authorization": "Bearer test_token"} + mock_backend.db.get_submission_by_id = MagicMock( + return_value={ + "id": 123, + "code": "test code", + } ) + + response = test_client.get("/admin/submissions/123", headers={"Authorization": "Bearer test_token"}) assert response.status_code == 200 data = response.json() assert data["status"] == "ok" @@ -146,10 +131,7 @@ def test_get_submission_not_found(self, test_client, mock_backend): mock_backend.db.__exit__ = MagicMock(return_value=None) mock_backend.db.get_submission_by_id = MagicMock(return_value=None) - response = test_client.get( - "/admin/submissions/999", - headers={"Authorization": "Bearer test_token"} - ) + response = test_client.get("/admin/submissions/999", headers={"Authorization": "Bearer test_token"}) assert response.status_code == 404 def test_delete_submission(self, test_client, mock_backend): @@ -158,10 +140,7 @@ def test_delete_submission(self, test_client, mock_backend): mock_backend.db.__exit__ = MagicMock(return_value=None) mock_backend.db.delete_submission = MagicMock() - response = test_client.delete( - "/admin/submissions/123", - headers={"Authorization": "Bearer test_token"} - ) + response = test_client.delete("/admin/submissions/123", headers={"Authorization": "Bearer test_token"}) assert response.status_code == 200 mock_backend.db.delete_submission.assert_called_once_with(123) @@ -174,7 +153,7 @@ def test_create_leaderboard_missing_directory(self, test_client, mock_backend): response = test_client.post( "/admin/leaderboards", headers={"Authorization": "Bearer test_token"}, - json={} # missing directory + json={}, # missing directory ) assert response.status_code == 400 assert "Missing required field: directory" in response.json()["detail"] @@ -186,7 +165,7 @@ def test_create_leaderboard_invalid_directory(self, test_client, mock_backend): headers={"Authorization": "Bearer test_token"}, json={ "directory": "../../../etc/passwd", # path traversal attempt - } + }, ) assert response.status_code == 400 @@ -201,12 +180,12 @@ def test_create_leaderboard_with_gpu_list(self, test_client, mock_backend): mock_definition = MagicMock() mock_definition.gpus = ["H100", "A100"] - with patch('kernelbot.api.main.resolve_problem_directory', return_value="/valid/path"): - with patch('kernelbot.api.main.make_task_definition', return_value=mock_definition): + with patch("kernelbot.api.main.resolve_problem_directory", return_value="/valid/path"): + with patch("kernelbot.api.main.make_task_definition", return_value=mock_definition): response = test_client.post( "/admin/leaderboards", headers={"Authorization": "Bearer test_token"}, - json={"directory": "identity_py"} + json={"directory": "identity_py"}, ) assert response.status_code == 200 assert response.json()["leaderboard"] == "identity_py-dev" @@ -223,12 +202,12 @@ def test_create_leaderboard_without_gpu(self, test_client, mock_backend): mock_definition = MagicMock() mock_definition.gpus = [] - with patch('kernelbot.api.main.resolve_problem_directory', return_value="/valid/path"): - with patch('kernelbot.api.main.make_task_definition', return_value=mock_definition): + with patch("kernelbot.api.main.resolve_problem_directory", return_value="/valid/path"): + with patch("kernelbot.api.main.make_task_definition", return_value=mock_definition): response = test_client.post( "/admin/leaderboards", headers={"Authorization": "Bearer test_token"}, - json={"directory": "identity_py"} + json={"directory": "identity_py"}, ) assert response.status_code == 400 assert "No gpus specified in task.yml" in response.json()["detail"] @@ -240,8 +219,7 @@ def test_delete_leaderboard(self, test_client, mock_backend): mock_backend.db.delete_leaderboard = MagicMock() response = test_client.delete( - "/admin/leaderboards/test-leaderboard", - headers={"Authorization": "Bearer test_token"} + "/admin/leaderboards/test-leaderboard", headers={"Authorization": "Bearer test_token"} ) assert response.status_code == 200 assert response.json()["leaderboard"] == "test-leaderboard" @@ -254,8 +232,7 @@ def test_delete_leaderboard_force(self, test_client, mock_backend): mock_backend.db.delete_leaderboard = MagicMock() response = test_client.delete( - "/admin/leaderboards/test-leaderboard?force=true", - headers={"Authorization": "Bearer test_token"} + "/admin/leaderboards/test-leaderboard?force=true", headers={"Authorization": "Bearer test_token"} ) assert response.status_code == 200 assert response.json()["force"] is True @@ -281,11 +258,9 @@ def test_update_problems_success(self, test_client, mock_backend): mock_result.skipped = [{"name": "problem4", "reason": "no changes"}] mock_result.errors = [] - with patch('kernelbot.api.main.sync_problems', return_value=mock_result) as mock_sync: + with patch("kernelbot.api.main.sync_problems", return_value=mock_result) as mock_sync: response = test_client.post( - "/admin/update-problems", - headers={"Authorization": "Bearer test_token"}, - json={} + "/admin/update-problems", headers={"Authorization": "Bearer test_token"}, json={} ) assert response.status_code == 200 data = response.json() @@ -314,11 +289,9 @@ def test_update_problems_with_problem_set(self, test_client, mock_backend): mock_result.skipped = [] mock_result.errors = [] - with patch('kernelbot.api.main.sync_problems', return_value=mock_result) as mock_sync: + with patch("kernelbot.api.main.sync_problems", return_value=mock_result) as mock_sync: response = test_client.post( - "/admin/update-problems", - headers={"Authorization": "Bearer test_token"}, - json={"problem_set": "nvidia"} + "/admin/update-problems", headers={"Authorization": "Bearer test_token"}, json={"problem_set": "nvidia"} ) assert response.status_code == 200 call_kwargs = mock_sync.call_args[1] @@ -335,11 +308,9 @@ def test_update_problems_with_force(self, test_client, mock_backend): mock_result.skipped = [] mock_result.errors = [] - with patch('kernelbot.api.main.sync_problems', return_value=mock_result) as mock_sync: + with patch("kernelbot.api.main.sync_problems", return_value=mock_result) as mock_sync: response = test_client.post( - "/admin/update-problems", - headers={"Authorization": "Bearer test_token"}, - json={"force": True} + "/admin/update-problems", headers={"Authorization": "Bearer test_token"}, json={"force": True} ) assert response.status_code == 200 call_kwargs = mock_sync.call_args[1] @@ -356,14 +327,11 @@ def test_update_problems_with_custom_repo_and_branch(self, test_client, mock_bac mock_result.skipped = [] mock_result.errors = [] - with patch('kernelbot.api.main.sync_problems', return_value=mock_result) as mock_sync: + with patch("kernelbot.api.main.sync_problems", return_value=mock_result) as mock_sync: response = test_client.post( "/admin/update-problems", headers={"Authorization": "Bearer test_token"}, - json={ - "repository": "other-org/other-repo", - "branch": "develop" - } + json={"repository": "other-org/other-repo", "branch": "develop"}, ) assert response.status_code == 200 call_kwargs = mock_sync.call_args[1] @@ -375,11 +343,11 @@ def test_update_problems_value_error(self, test_client, mock_backend): mock_backend.db.__enter__ = MagicMock(return_value=mock_backend.db) mock_backend.db.__exit__ = MagicMock(return_value=None) - with patch('kernelbot.api.main.sync_problems', side_effect=ValueError("Invalid branch name")): + with patch("kernelbot.api.main.sync_problems", side_effect=ValueError("Invalid branch name")): response = test_client.post( "/admin/update-problems", headers={"Authorization": "Bearer test_token"}, - json={"branch": "invalid/branch"} + json={"branch": "invalid/branch"}, ) assert response.status_code == 400 assert "Invalid branch name" in response.json()["detail"] @@ -395,11 +363,9 @@ def test_update_problems_with_errors(self, test_client, mock_backend): mock_result.skipped = [] mock_result.errors = [{"name": "bad-problem", "error": "create failed: DB error"}] - with patch('kernelbot.api.main.sync_problems', return_value=mock_result): + with patch("kernelbot.api.main.sync_problems", return_value=mock_result): response = test_client.post( - "/admin/update-problems", - headers={"Authorization": "Bearer test_token"}, - json={} + "/admin/update-problems", headers={"Authorization": "Bearer test_token"}, json={} ) assert response.status_code == 200 data = response.json() @@ -409,38 +375,15 @@ def test_update_problems_with_errors(self, test_client, mock_backend): class TestSubmissionRateLimit: - """Test per-user submission rate limiting on Modal B200 for leaderboard 730.""" + """Test per-user submission rate limiting (generic rate limiting, no hardcoded B200/730).""" - def test_rate_limit_blocks_b200_leaderboard_730(self, test_client, mock_backend): - """Second B200 submission to leaderboard 730 within 1 hour is rejected with 429.""" - mock_backend.db.__enter__ = MagicMock(return_value=mock_backend.db) - mock_backend.db.__exit__ = MagicMock(return_value=None) - - recent_time = datetime.datetime.now(tz=datetime.timezone.utc) - mock_backend.db.check_user_rate_limit = MagicMock(return_value=recent_time) - mock_backend.db.get_leaderboard_id = MagicMock(return_value=730) - mock_backend.db.validate_cli_id = MagicMock( - return_value={"user_id": "123", "user_name": "testuser"} - ) - - response = test_client.post( - "/test-lb/B200/test", - headers={"X-Popcorn-Cli-Id": "test-cli-id"}, - files={"file": ("solution.py", b"print('hello')", "text/plain")}, - ) - assert response.status_code == 429 - assert "Rate limit exceeded" in response.json()["detail"] - assert "NVIDIA runner" in response.json()["detail"] - - def test_rate_limit_skipped_for_non_b200(self, test_client, mock_backend): - """Rate limit is not enforced for non-B200 GPUs even on leaderboard 730.""" + def test_submission_endpoint_does_not_reject_without_rate_limit(self, test_client, mock_backend): + """Submissions are not rate limited when no rate limit is configured.""" mock_backend.db.__enter__ = MagicMock(return_value=mock_backend.db) mock_backend.db.__exit__ = MagicMock(return_value=None) mock_backend.accepts_jobs = True - mock_backend.db.validate_cli_id = MagicMock( - return_value={"user_id": "123", "user_name": "testuser"} - ) + mock_backend.db.validate_cli_id = MagicMock(return_value={"user_id": "123", "user_name": "testuser"}) mock_lb = MagicMock() mock_lb.__getitem__ = lambda self, key: {"gpu_types": ["H100"]}[key] @@ -452,55 +395,83 @@ def test_rate_limit_skipped_for_non_b200(self, test_client, mock_backend): headers={"X-Popcorn-Cli-Id": "test-cli-id"}, files={"file": ("solution.py", b"print('hello')", "text/plain")}, ) - # Should not hit rate limit at all — check_user_rate_limit should not be called + # Should not be rate limited (rate limits are now configurable per-leaderboard/GPU) assert response.status_code != 429 - def test_rate_limit_skipped_for_other_leaderboard(self, test_client, mock_backend): - """Rate limit is not enforced for B200 on a leaderboard other than 730.""" + +class TestLeaderboardRateLimitsAPI: + """Test leaderboard rate limit API endpoints.""" + + def test_get_rate_limits(self, test_client, mock_backend): + """GET /leaderboard/rate-limits/{name} returns rate limits.""" mock_backend.db.__enter__ = MagicMock(return_value=mock_backend.db) mock_backend.db.__exit__ = MagicMock(return_value=None) - mock_backend.accepts_jobs = True + mock_backend.db.get_leaderboard_rate_limits = MagicMock(return_value={"A100": 3600, "H100": None}) + + response = test_client.get("/leaderboard/rate-limits/test-leaderboard") + assert response.status_code == 200 + data = response.json() + assert data["status"] == "ok" + assert data["rate_limits"]["A100"] == 3600 + assert data["rate_limits"]["H100"] is None - recent_time = datetime.datetime.now(tz=datetime.timezone.utc) - mock_backend.db.check_user_rate_limit = MagicMock(return_value=recent_time) - mock_backend.db.get_leaderboard_id = MagicMock(return_value=999) - mock_backend.db.validate_cli_id = MagicMock( - return_value={"user_id": "123", "user_name": "testuser"} + def test_set_rate_limit_requires_auth(self, test_client): + """POST /leaderboard/rate-limits/{name}/{gpu} requires authorization.""" + response = test_client.post("/leaderboard/rate-limits/test-leaderboard/A100?rate_limit_seconds=3600") + assert response.status_code == 401 + + def test_set_rate_limit_rejects_invalid_token(self, test_client): + """POST /leaderboard/rate-limits/{name}/{gpu} rejects invalid tokens.""" + response = test_client.post( + "/leaderboard/rate-limits/test-leaderboard/A100?rate_limit_seconds=3600", + headers={"Authorization": "Bearer wrong_token"}, ) + assert response.status_code == 401 - mock_lb = MagicMock() - mock_lb.__getitem__ = lambda self, key: {"gpu_types": ["B200"]}[key] - mock_lb.get = lambda key, default=None: {"gpu_types": ["B200"]}.get(key, default) - mock_backend.db.get_leaderboard = MagicMock(return_value=mock_lb) + def test_set_rate_limit(self, test_client, mock_backend): + """POST /leaderboard/rate-limits/{name}/{gpu} sets rate limit.""" + mock_backend.db.__enter__ = MagicMock(return_value=mock_backend.db) + mock_backend.db.__exit__ = MagicMock(return_value=None) + mock_backend.db.set_leaderboard_gpu_rate_limit = MagicMock() response = test_client.post( - "/other-lb/B200/test", - headers={"X-Popcorn-Cli-Id": "test-cli-id"}, - files={"file": ("solution.py", b"print('hello')", "text/plain")}, + "/leaderboard/rate-limits/test-leaderboard/A100?rate_limit_seconds=3600", + headers={"Authorization": "Bearer test_token"}, ) - # Should not be rate limited since leaderboard ID is not 730 - assert response.status_code != 429 + assert response.status_code == 200 + data = response.json() + assert data["status"] == "ok" + assert data["leaderboard_name"] == "test-leaderboard" + assert data["gpu_type"] == "A100" + assert data["rate_limit_seconds"] == 3600 + mock_backend.db.set_leaderboard_gpu_rate_limit.assert_called_once_with("test-leaderboard", "A100", 3600) - def test_rate_limit_allows_first_b200_submission(self, test_client, mock_backend): - """First B200 submission to leaderboard 730 passes the rate limit check.""" + def test_set_rate_limit_zero_clears_limit(self, test_client, mock_backend): + """POST /leaderboard/rate-limits/{name}/{gpu} with 0 clears the rate limit.""" mock_backend.db.__enter__ = MagicMock(return_value=mock_backend.db) mock_backend.db.__exit__ = MagicMock(return_value=None) - mock_backend.accepts_jobs = True + mock_backend.db.set_leaderboard_gpu_rate_limit = MagicMock() - mock_backend.db.check_user_rate_limit = MagicMock(return_value=None) - mock_backend.db.get_leaderboard_id = MagicMock(return_value=730) - mock_backend.db.validate_cli_id = MagicMock( - return_value={"user_id": "123", "user_name": "testuser"} + response = test_client.post( + "/leaderboard/rate-limits/test-leaderboard/A100?rate_limit_seconds=0", + headers={"Authorization": "Bearer test_token"}, ) + assert response.status_code == 200 + data = response.json() + assert data["rate_limit_seconds"] is None + mock_backend.db.set_leaderboard_gpu_rate_limit.assert_called_once_with("test-leaderboard", "A100", None) - mock_lb = MagicMock() - mock_lb.__getitem__ = lambda self, key: {"gpu_types": ["B200"]}[key] - mock_lb.get = lambda key, default=None: {"gpu_types": ["B200"]}.get(key, default) - mock_backend.db.get_leaderboard = MagicMock(return_value=mock_lb) + def test_set_rate_limit_negative_clears_limit(self, test_client, mock_backend): + """POST /leaderboard/rate-limits/{name}/{gpu} with negative value clears the rate limit.""" + mock_backend.db.__enter__ = MagicMock(return_value=mock_backend.db) + mock_backend.db.__exit__ = MagicMock(return_value=None) + mock_backend.db.set_leaderboard_gpu_rate_limit = MagicMock() response = test_client.post( - "/test-lb/B200/test", - headers={"X-Popcorn-Cli-Id": "test-cli-id"}, - files={"file": ("solution.py", b"print('hello')", "text/plain")}, + "/leaderboard/rate-limits/test-leaderboard/H100?rate_limit_seconds=-1", + headers={"Authorization": "Bearer test_token"}, ) - assert response.status_code != 429 + assert response.status_code == 200 + data = response.json() + assert data["rate_limit_seconds"] is None + mock_backend.db.set_leaderboard_gpu_rate_limit.assert_called_once_with("test-leaderboard", "H100", None) diff --git a/tests/test_leaderboard_db.py b/tests/test_leaderboard_db.py index 3d034cbf..f7ef2015 100644 --- a/tests/test_leaderboard_db.py +++ b/tests/test_leaderboard_db.py @@ -53,8 +53,7 @@ def _create_submission_run( db.create_submission_run( submission, start=start or datetime.datetime.now(tz=datetime.timezone.utc), - end=end - or (datetime.datetime.now(tz=datetime.timezone.utc) + datetime.timedelta(seconds=10)), + end=end or (datetime.datetime.now(tz=datetime.timezone.utc) + datetime.timedelta(seconds=10)), mode=mode, secret=secret, runner=runner, @@ -133,14 +132,10 @@ def test_leaderboard_basics(database, task_directory): assert db.get_leaderboard_submissions("test-leaderboard", "A100", "5", 100) == [] assert db.get_leaderboard_submission_count("test-leaderboard", "A100", "5") == 0 - with pytest.raises( - KernelBotError, match="Invalid GPU type 'A99' for leaderboard 'test-leaderboard'" - ): + with pytest.raises(KernelBotError, match="Invalid GPU type 'A99' for leaderboard 'test-leaderboard'"): assert db.get_leaderboard_submissions("test-leaderboard", "A99", "5", 100) == [] - with pytest.raises( - KernelBotError, match="Invalid GPU type 'A99' for leaderboard 'test-leaderboard'" - ): + with pytest.raises(KernelBotError, match="Invalid GPU type 'A99' for leaderboard 'test-leaderboard'"): assert db.get_leaderboard_submission_count("test-leaderboard", "A99", "5") == 0 @@ -232,8 +227,7 @@ def test_leaderboard_submission_basic(database, submit_leaderboard): ) expected_meta = { - k: getattr(run_result, k) - for k in ("stdout", "stderr", "success", "exit_code", "command", "duration") + k: getattr(run_result, k) for k in ("stdout", "stderr", "success", "exit_code", "command", "duration") } submission = db.get_submission_by_id(sub_id) @@ -281,12 +275,8 @@ def test_leaderboard_submission_count(database, submit_leaderboard): "submit-leaderboard", "submission.py", 5, dangerous_code, submit_time, user_name="user" ) _create_submission_run(db, sub_id, mode="test", secret=False, runner="A100") - _create_submission_run( - db, sub_id, mode="leaderboard", secret=True, runner="H100", score=5.5 - ) - _create_submission_run( - db, sub_id, mode="leaderboard", secret=False, runner="A100", score=1.5 - ) + _create_submission_run(db, sub_id, mode="leaderboard", secret=True, runner="H100", score=5.5) + _create_submission_run(db, sub_id, mode="leaderboard", secret=False, runner="A100", score=1.5) submission = db.get_submission_by_id(sub_id) assert len(submission["runs"]) == 3 @@ -368,42 +358,45 @@ def test_leaderboard_submission_ranked(database, submit_leaderboard): }, ] + def test_validate_identity_web_auth_happy_path(database, submit_leaderboard): with database as db: db.cursor.execute( - """ + """ INSERT INTO leaderboard.user_info (id, user_name, web_auth_id) VALUES (%s, %s, %s) """, - ("1234", "sara_jojo","2345" ), - ) - user_info = db.validate_identity("2345",IdentityType.WEB) - assert user_info["user_id"] =="1234" - assert user_info["user_name"] =="sara_jojo" - assert user_info["id_type"] ==IdentityType.WEB.value + ("1234", "sara_jojo", "2345"), + ) + user_info = db.validate_identity("2345", IdentityType.WEB) + assert user_info["user_id"] == "1234" + assert user_info["user_name"] == "sara_jojo" + assert user_info["id_type"] == IdentityType.WEB.value + -def test_validate_identity_web_auth_not_found(database, submit_leaderboard): +def test_validate_identity_web_auth_not_found(database, submit_leaderboard): with database as db: db.cursor.execute( - """ + """ INSERT INTO leaderboard.user_info (id, user_name) VALUES (%s, %s) """, - ("1234", "sara_jojo"), - ) - user_info = db.validate_identity("2345",IdentityType.WEB) + ("1234", "sara_jojo"), + ) + user_info = db.validate_identity("2345", IdentityType.WEB) assert user_info is None + def test_validate_identity_web_auth_missing(database, submit_leaderboard): with database as db: db.cursor.execute( - """ + """ INSERT INTO leaderboard.user_info (id, user_name) VALUES (%s, %s) """, - ("1234", "sara_jojo"), - ) - res = db.validate_identity("2345",IdentityType.WEB) + ("1234", "sara_jojo"), + ) + res = db.validate_identity("2345", IdentityType.WEB) assert res is None @@ -418,9 +411,7 @@ def test_leaderboard_submission_deduplication(database, submit_leaderboard): datetime.datetime.now(), user_name="user", ) - db.create_submission( - "submit-leaderboard", "other.py", 6, "pass", datetime.datetime.now(), user_name="other" - ) + db.create_submission("submit-leaderboard", "other.py", 6, "pass", datetime.datetime.now(), user_name="other") db.cursor.execute("SELECT COUNT(*) FROM leaderboard.code_files") assert db.cursor.fetchone()[0] == 1 @@ -447,9 +438,7 @@ def test_leaderboard_submission_delete(database, submit_leaderboard): _create_submission_run(db, sub_id, mode="leaderboard", secret=False, runner="A100", score=5) _create_submission_run(db, sub_id, mode="leaderboard", secret=True, runner="A100", score=5) - _create_submission_run( - db, other_sub, mode="leaderboard", secret=False, runner="A100", score=5 - ) + _create_submission_run(db, other_sub, mode="leaderboard", secret=False, runner="A100", score=5) db.mark_submission_done(sub_id) db.cursor.execute("SELECT COUNT(*) FROM leaderboard.runs") @@ -555,9 +544,7 @@ def test_leaderboard_update(database, task_directory): def test_generate_stats(database, submit_leaderboard): with database as db: start = datetime.datetime.now(tz=datetime.timezone.utc) - sub_id = db.create_submission( - "submit-leaderboard", "submission.py", 5, "pass", start, user_name="user" - ) + sub_id = db.create_submission("submit-leaderboard", "submission.py", 5, "pass", start, user_name="user") _create_submission_run( db, sub_id, @@ -616,9 +603,7 @@ def test_check_user_rate_limit_recent_submission(database, submit_leaderboard): """Test rate limit returns submission_time when user submitted recently""" submit_time = datetime.datetime.now(tz=datetime.timezone.utc) with database as db: - db.create_submission( - "submit-leaderboard", "file.py", 5, "code", submit_time, user_name="user" - ) + db.create_submission("submit-leaderboard", "file.py", 5, "code", submit_time, user_name="user") result = db.check_user_rate_limit("5") assert result is not None assert abs((result - submit_time).total_seconds()) < 2 @@ -628,9 +613,7 @@ def test_check_user_rate_limit_old_submission(database, submit_leaderboard): """Test rate limit returns None when submission is older than the window""" old_time = datetime.datetime.now(tz=datetime.timezone.utc) - datetime.timedelta(hours=2) with database as db: - db.create_submission( - "submit-leaderboard", "file.py", 5, "code", old_time, user_name="user" - ) + db.create_submission("submit-leaderboard", "file.py", 5, "code", old_time, user_name="user") result = db.check_user_rate_limit("5") assert result is None @@ -639,9 +622,7 @@ def test_check_user_rate_limit_different_user(database, submit_leaderboard): """Test rate limit only applies to the specific user""" submit_time = datetime.datetime.now(tz=datetime.timezone.utc) with database as db: - db.create_submission( - "submit-leaderboard", "file.py", 5, "code", submit_time, user_name="user5" - ) + db.create_submission("submit-leaderboard", "file.py", 5, "code", submit_time, user_name="user5") # User 6 should not be rate limited result = db.check_user_rate_limit("6") assert result is None @@ -820,3 +801,181 @@ def test_get_user_submissions_with_multiple_runs(database, submit_leaderboard): assert 1.5 in scores assert 2.0 in scores + +# --- Leaderboard GPU Rate Limit Tests --- + + +def test_get_leaderboard_rate_limits_empty(database, submit_leaderboard): + """Test get_leaderboard_rate_limits returns None values when no limits set""" + with database as db: + result = db.get_leaderboard_rate_limits("submit-leaderboard") + # Rate limits should be None by default + assert "A100" in result + assert "H100" in result + assert result["A100"] is None + assert result["H100"] is None + + +def test_set_leaderboard_gpu_rate_limit(database, submit_leaderboard): + """Test setting a rate limit for a specific GPU on a leaderboard""" + with database as db: + db.set_leaderboard_gpu_rate_limit("submit-leaderboard", "A100", 3600) + result = db.get_leaderboard_rate_limits("submit-leaderboard") + assert result["A100"] == 3600 + assert result["H100"] is None + + +def test_set_leaderboard_gpu_rate_limit_multiple(database, submit_leaderboard): + """Test setting different rate limits for different GPUs""" + with database as db: + db.set_leaderboard_gpu_rate_limit("submit-leaderboard", "A100", 3600) + db.set_leaderboard_gpu_rate_limit("submit-leaderboard", "H100", 7200) + result = db.get_leaderboard_rate_limits("submit-leaderboard") + assert result["A100"] == 3600 + assert result["H100"] == 7200 + + +def test_set_leaderboard_gpu_rate_limit_update(database, submit_leaderboard): + """Test updating an existing rate limit""" + with database as db: + db.set_leaderboard_gpu_rate_limit("submit-leaderboard", "A100", 3600) + db.set_leaderboard_gpu_rate_limit("submit-leaderboard", "A100", 1800) + result = db.get_leaderboard_rate_limits("submit-leaderboard") + assert result["A100"] == 1800 + + +def test_is_user_rate_limited_no_limit_set(database, submit_leaderboard): + """Test is_user_rate_limited returns False when no rate limit is configured""" + with database as db: + lb_id = db.get_leaderboard_id("submit-leaderboard") + is_limited, reason = db.is_user_rate_limited(5, lb_id, "A100") + assert is_limited is False + assert reason is None + + +def test_is_user_rate_limited_no_submissions(database, submit_leaderboard): + """Test is_user_rate_limited returns False when user has no submissions""" + with database as db: + db.set_leaderboard_gpu_rate_limit("submit-leaderboard", "A100", 3600) + lb_id = db.get_leaderboard_id("submit-leaderboard") + is_limited, reason = db.is_user_rate_limited(999, lb_id, "A100") + assert is_limited is False + assert reason is None + + +def test_is_user_rate_limited_recent_submission(database, submit_leaderboard): + """Test is_user_rate_limited returns True when user submitted recently""" + with database as db: + db.set_leaderboard_gpu_rate_limit("submit-leaderboard", "A100", 3600) + db.create_submission( + "submit-leaderboard", + "file.py", + 5, + "code", + datetime.datetime.now(tz=datetime.timezone.utc), + user_name="user", + ) + lb_id = db.get_leaderboard_id("submit-leaderboard") + is_limited, reason = db.is_user_rate_limited(5, lb_id, "A100") + assert is_limited is True + assert reason is not None + assert "Rate limit exceeded" in reason + assert "3600 seconds" in reason + + +def test_is_user_rate_limited_old_submission(database, submit_leaderboard): + """Test is_user_rate_limited returns False when submission is older than rate limit window""" + with database as db: + db.set_leaderboard_gpu_rate_limit("submit-leaderboard", "A100", 3600) + old_time = datetime.datetime.now(tz=datetime.timezone.utc) - datetime.timedelta(hours=2) + db.create_submission("submit-leaderboard", "file.py", 5, "code", old_time, user_name="user") + lb_id = db.get_leaderboard_id("submit-leaderboard") + is_limited, reason = db.is_user_rate_limited(5, lb_id, "A100") + assert is_limited is False + assert reason is None + + +def test_is_user_rate_limited_different_user(database, submit_leaderboard): + """Test is_user_rate_limited only applies to the specific user""" + with database as db: + db.set_leaderboard_gpu_rate_limit("submit-leaderboard", "A100", 3600) + db.create_submission( + "submit-leaderboard", + "file.py", + 5, + "code", + datetime.datetime.now(tz=datetime.timezone.utc), + user_name="user5", + ) + lb_id = db.get_leaderboard_id("submit-leaderboard") + # User 6 should not be rate limited + is_limited, reason = db.is_user_rate_limited(6, lb_id, "A100") + assert is_limited is False + # User 5 should be rate limited + is_limited, reason = db.is_user_rate_limited(5, lb_id, "A100") + assert is_limited is True + + +def test_is_user_allowed_to_submit_no_limit(database, submit_leaderboard): + """Test is_user_allowed_to_submit returns True when no rate limit is configured""" + with database as db: + allowed, reason = db.is_user_allowed_to_submit(5, "submit-leaderboard", ["A100"]) + assert allowed is True + assert reason is None + + +def test_is_user_allowed_to_submit_allowed(database, submit_leaderboard): + """Test is_user_allowed_to_submit returns True when user hasn't submitted recently""" + with database as db: + db.set_leaderboard_gpu_rate_limit("submit-leaderboard", "A100", 3600) + allowed, reason = db.is_user_allowed_to_submit(999, "submit-leaderboard", ["A100"]) + assert allowed is True + assert reason is None + + +def test_is_user_allowed_to_submit_blocked(database, submit_leaderboard): + """Test is_user_allowed_to_submit returns False when user submitted recently""" + with database as db: + db.set_leaderboard_gpu_rate_limit("submit-leaderboard", "A100", 3600) + db.create_submission( + "submit-leaderboard", + "file.py", + 5, + "code", + datetime.datetime.now(tz=datetime.timezone.utc), + user_name="user", + ) + allowed, reason = db.is_user_allowed_to_submit(5, "submit-leaderboard", ["A100"]) + assert allowed is False + assert reason is not None + assert "Rate limit exceeded" in reason + + +def test_is_user_allowed_to_submit_multiple_gpus_one_blocked(database, submit_leaderboard): + """Test is_user_allowed_to_submit returns False if any GPU has rate limit exceeded""" + with database as db: + # Set rate limit only on A100 + db.set_leaderboard_gpu_rate_limit("submit-leaderboard", "A100", 3600) + db.create_submission( + "submit-leaderboard", + "file.py", + 5, + "code", + datetime.datetime.now(tz=datetime.timezone.utc), + user_name="user", + ) + # Both GPUs requested, A100 is rate limited + allowed, reason = db.is_user_allowed_to_submit(5, "submit-leaderboard", ["A100", "H100"]) + assert allowed is False + assert "A100" in reason + + +def test_is_user_allowed_to_submit_multiple_gpus_all_allowed(database, submit_leaderboard): + """Test is_user_allowed_to_submit returns True when all GPUs are within limit""" + with database as db: + db.set_leaderboard_gpu_rate_limit("submit-leaderboard", "A100", 3600) + db.set_leaderboard_gpu_rate_limit("submit-leaderboard", "H100", 3600) + # No submissions, so all GPUs should be allowed + allowed, reason = db.is_user_allowed_to_submit(5, "submit-leaderboard", ["A100", "H100"]) + assert allowed is True + assert reason is None diff --git a/tests/test_submission.py b/tests/test_submission.py index e22fcb8e..b5ac76f7 100644 --- a/tests/test_submission.py +++ b/tests/test_submission.py @@ -31,6 +31,8 @@ def mock_backend(): "name": "test_board", } db_context.get_leaderboard_gpu_types.return_value = ["A100", "V100"] + # Default: user is allowed to submit (not rate-limited) + db_context.is_user_allowed_to_submit.return_value = (True, None) return backend @@ -49,9 +51,7 @@ def test_check_deadline(): "name": "test", } - with pytest.raises( - KernelBotError, match=r"The deadline to submit to test has passed\.\nIt was.*and today is.*" - ): + with pytest.raises(KernelBotError, match=r"The deadline to submit to test has passed\.\nIt was.*and today is.*"): submission.check_deadline(past_deadline) @@ -125,9 +125,7 @@ def test_get_popcorn_directives_invalid(): # Empty leaderboard but valid GPU code_empty_leaderboard = """#!POPCORN gpu A100 #!POPCORN leaderboard""" - with pytest.raises( - KernelBotError, match="!POPCORN directive missing argument: #!POPCORN leaderboard" - ): + with pytest.raises(KernelBotError, match="!POPCORN directive missing argument: #!POPCORN leaderboard"): submission._get_popcorn_directives(code_empty_leaderboard) # Invalid directive @@ -137,9 +135,7 @@ def test_get_popcorn_directives_invalid(): # Multiple leaderboard directives (last one wins or first one wins?) code_multiple_leaderboard = """#!POPCORN leaderboard first_board #!POPCORN leaderboard second_board""" - with pytest.raises( - KernelBotError, match="Found multiple values for !POPCORN directive leaderboard" - ): + with pytest.raises(KernelBotError, match="Found multiple values for !POPCORN directive leaderboard"): submission._get_popcorn_directives(code_multiple_leaderboard) @@ -280,22 +276,100 @@ def test_prepare_submission_checks(mock_backend): # invalid file extension req.file_name = "test.txt" - with pytest.raises( - KernelBotError, match=r"Please provide a Python \(.py\) or CUDA \(.cu / .cuh / .cpp\) file" - ): + with pytest.raises(KernelBotError, match=r"Please provide a Python \(.py\) or CUDA \(.cu / .cuh / .cpp\) file"): submission.prepare_submission(req, mock_backend) req.file_name = "test.py" req.gpus = ["A99"] with pytest.raises( KernelBotError, - match=re.escape( - "GPU A99 not available for `test_board`\nChoose one of: * A100\n * V100\n" - ), + match=re.escape("GPU A99 not available for `test_board`\nChoose one of: * A100\n * V100\n"), ): submission.prepare_submission(req, mock_backend) +def test_prepare_submission_rate_limited(mock_backend): + """Test that prepare_submission raises KernelBotError when user is rate-limited.""" + req = submission.SubmissionRequest( + code="#!POPCORN leaderboard test_board\nprint('hello world')", + file_name="test.py", + user_id=1, + user_name="test_user", + gpus=["A100"], + leaderboard=None, + ) + + # Mock rate limit exceeded + rate_limit_reason = "Rate limit exceeded for A100. You can submit once every 3600 seconds." + mock_backend.db.is_user_allowed_to_submit.return_value = (False, rate_limit_reason) + + with pytest.raises(KernelBotError, match="Rate limit exceeded"): + submission.prepare_submission(req, mock_backend) + + # Verify is_user_allowed_to_submit was called with correct arguments + mock_backend.db.is_user_allowed_to_submit.assert_called_once_with(1, "test_board", ["A100"]) + + +def test_prepare_submission_rate_limit_allowed(mock_backend): + """Test that prepare_submission succeeds when user is not rate-limited.""" + req = submission.SubmissionRequest( + code="#!POPCORN leaderboard test_board\nprint('hello world')", + file_name="test.py", + user_id=1, + user_name="test_user", + gpus=["A100"], + leaderboard=None, + ) + + # Explicitly set user is allowed (default, but being explicit for clarity) + mock_backend.db.is_user_allowed_to_submit.return_value = (True, None) + + result = submission.prepare_submission(req, mock_backend) + + assert isinstance(result, submission.ProcessedSubmissionRequest) + assert result.leaderboard == "test_board" + mock_backend.db.is_user_allowed_to_submit.assert_called_once_with(1, "test_board", ["A100"]) + + +def test_prepare_submission_rate_limit_called_with_correct_gpus(mock_backend): + """Test that is_user_allowed_to_submit is called with the correct GPU list.""" + # Test with multiple GPUs from POPCORN directive + mock_backend.db.get_leaderboard_gpu_types.return_value = ["A100", "V100", "H100"] + + req = submission.SubmissionRequest( + code="#!POPCORN gpu V100 H100\n#!POPCORN leaderboard test_board\nprint('test')", + file_name="test.py", + user_id=42, + user_name="test_user", + gpus=None, + leaderboard=None, + ) + + result = submission.prepare_submission(req, mock_backend) + + assert result.gpus == ["V100", "H100"] + mock_backend.db.is_user_allowed_to_submit.assert_called_once_with(42, "test_board", ["V100", "H100"]) + + +def test_prepare_submission_rate_limit_single_gpu_auto_assign(mock_backend): + """Test that is_user_allowed_to_submit is called with auto-assigned single GPU.""" + mock_backend.db.get_leaderboard_gpu_types.return_value = ["H100"] + + req = submission.SubmissionRequest( + code="#!POPCORN leaderboard test_board\nprint('test')", + file_name="test.py", + user_id=5, + user_name="test_user", + gpus=None, + leaderboard=None, + ) + + result = submission.prepare_submission(req, mock_backend) + + assert result.gpus == ["H100"] + mock_backend.db.is_user_allowed_to_submit.assert_called_once_with(5, "test_board", ["H100"]) + + def test_compute_score(): mock_task = mock.Mock() mock_result = mock.Mock() @@ -376,9 +450,7 @@ def test_generate_run_verdict(mock_backend): assert result == "> 5th place on A100: 1500 ms" # Test personal best (rank > 10) - mock_backend.db.get_leaderboard_submissions.return_value = [ - {"submission_id": 123, "user_id": 42, "rank": 15} - ] + mock_backend.db.get_leaderboard_submissions.return_value = [{"submission_id": 123, "user_id": 42, "rank": 15}] result = generate_run_verdict(mock_backend, run_item, sub_data) assert result == "> Personal best on A100: 1500 ms" From e08a4b6ebcdc33b1a636e6ce272c3360a4841091 Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Sat, 7 Feb 2026 22:25:47 +0100 Subject: [PATCH 3/3] stuff more --- src/libkernelbot/leaderboard_db.py | 1 - tests/test_leaderboard_db.py | 16 ++++++++++++---- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/src/libkernelbot/leaderboard_db.py b/src/libkernelbot/leaderboard_db.py index 0106b545..af6de1d0 100644 --- a/src/libkernelbot/leaderboard_db.py +++ b/src/libkernelbot/leaderboard_db.py @@ -1220,7 +1220,6 @@ def is_user_rate_limited(self, user_id: int, leaderboard_id: int, gpu_type: str) return False, None rate_limit_seconds = row[0] - print(f"GPU type: {gpu_type}, Rate limit: {rate_limit_seconds}") self.cursor.execute( """ SELECT s.submission_time diff --git a/tests/test_leaderboard_db.py b/tests/test_leaderboard_db.py index f7ef2015..ec16a97c 100644 --- a/tests/test_leaderboard_db.py +++ b/tests/test_leaderboard_db.py @@ -867,7 +867,7 @@ def test_is_user_rate_limited_recent_submission(database, submit_leaderboard): """Test is_user_rate_limited returns True when user submitted recently""" with database as db: db.set_leaderboard_gpu_rate_limit("submit-leaderboard", "A100", 3600) - db.create_submission( + sub_id = db.create_submission( "submit-leaderboard", "file.py", 5, @@ -875,6 +875,8 @@ def test_is_user_rate_limited_recent_submission(database, submit_leaderboard): datetime.datetime.now(tz=datetime.timezone.utc), user_name="user", ) + # Create a run with A100 GPU - required since rate limit query joins with runs table + _create_submission_run(db, sub_id, runner="A100", score=1.0) lb_id = db.get_leaderboard_id("submit-leaderboard") is_limited, reason = db.is_user_rate_limited(5, lb_id, "A100") assert is_limited is True @@ -899,7 +901,7 @@ def test_is_user_rate_limited_different_user(database, submit_leaderboard): """Test is_user_rate_limited only applies to the specific user""" with database as db: db.set_leaderboard_gpu_rate_limit("submit-leaderboard", "A100", 3600) - db.create_submission( + sub_id = db.create_submission( "submit-leaderboard", "file.py", 5, @@ -907,6 +909,8 @@ def test_is_user_rate_limited_different_user(database, submit_leaderboard): datetime.datetime.now(tz=datetime.timezone.utc), user_name="user5", ) + # Create a run with A100 GPU - required since rate limit query joins with runs table + _create_submission_run(db, sub_id, runner="A100", score=1.0) lb_id = db.get_leaderboard_id("submit-leaderboard") # User 6 should not be rate limited is_limited, reason = db.is_user_rate_limited(6, lb_id, "A100") @@ -937,7 +941,7 @@ def test_is_user_allowed_to_submit_blocked(database, submit_leaderboard): """Test is_user_allowed_to_submit returns False when user submitted recently""" with database as db: db.set_leaderboard_gpu_rate_limit("submit-leaderboard", "A100", 3600) - db.create_submission( + sub_id = db.create_submission( "submit-leaderboard", "file.py", 5, @@ -945,6 +949,8 @@ def test_is_user_allowed_to_submit_blocked(database, submit_leaderboard): datetime.datetime.now(tz=datetime.timezone.utc), user_name="user", ) + # Create a run with A100 GPU - required since rate limit query joins with runs table + _create_submission_run(db, sub_id, runner="A100", score=1.0) allowed, reason = db.is_user_allowed_to_submit(5, "submit-leaderboard", ["A100"]) assert allowed is False assert reason is not None @@ -956,7 +962,7 @@ def test_is_user_allowed_to_submit_multiple_gpus_one_blocked(database, submit_le with database as db: # Set rate limit only on A100 db.set_leaderboard_gpu_rate_limit("submit-leaderboard", "A100", 3600) - db.create_submission( + sub_id = db.create_submission( "submit-leaderboard", "file.py", 5, @@ -964,6 +970,8 @@ def test_is_user_allowed_to_submit_multiple_gpus_one_blocked(database, submit_le datetime.datetime.now(tz=datetime.timezone.utc), user_name="user", ) + # Create a run with A100 GPU - required since rate limit query joins with runs table + _create_submission_run(db, sub_id, runner="A100", score=1.0) # Both GPUs requested, A100 is rate limited allowed, reason = db.is_user_allowed_to_submit(5, "submit-leaderboard", ["A100", "H100"]) assert allowed is False