moved name_cache into db

Author: ngc92

src/discord-cluster-manager/cogs/admin_cog.py

@@ -13,13 +13,13 @@
 from consts import GitHubGPU, ModalGPU
 from discord import app_commands
 from discord.ext import commands, tasks
-from leaderboard_db import leaderboard_name_autocomplete
 from task import LeaderboardTask, make_task
 from ui.misc import ConfirmationView, DeleteConfirmationModal, GPUSelectionView
 from utils import (
     KernelBotError,
     LeaderboardItem,
     SubmissionItem,
+    leaderboard_name_autocomplete,
     send_discord_message,
     setup_logging,
     with_error_handling,

src/discord-cluster-manager/cogs/leaderboard_cog.py

@@ -10,7 +10,6 @@
 )
 from discord import app_commands
 from discord.ext import commands
-from leaderboard_db import leaderboard_name_autocomplete
 from report import MultiProgressReporter
 from submission import SubmissionRequest, prepare_submission
 from ui.misc import GPUSelectionView
@@ -22,6 +21,7 @@
     SubmissionItem,
     format_time,
     get_user_from_id,
+    leaderboard_name_autocomplete,
     send_discord_message,
     setup_logging,
     with_error_handling,

src/discord-cluster-manager/leaderboard_db.py

@@ -4,16 +4,10 @@
 import logging
 from typing import List, Optional
 
-import discord
 import psycopg2
 from env import (
     DATABASE_URL,
     DISABLE_SSL,
-    POSTGRES_DATABASE,
-    POSTGRES_HOST,
-    POSTGRES_PASSWORD,
-    POSTGRES_PORT,
-    POSTGRES_USER,
 )
 from run_eval import CompileResult, RunResult, SystemInfo
 from task import LeaderboardTask
@@ -27,34 +21,9 @@
     setup_logging,
 )
 
-leaderboard_name_cache = LRUCache(max_size=512)
-
 logger = setup_logging(__name__)
 
 
-async def leaderboard_name_autocomplete(
-    interaction: discord.Interaction,
-    current: str,
-) -> list[discord.app_commands.Choice[str]]:
-    """Return leaderboard names that match the current typed name"""
-    try:
-        cached_value = leaderboard_name_cache[current]
-        if cached_value is not None:
-            return cached_value
-
-        bot = interaction.client
-        with bot.leaderboard_db as db:
-            leaderboards = db.get_leaderboard_names()
-        filtered = [lb for lb in leaderboards if current.lower() in lb.lower()]
-        leaderboard_name_cache[current] = [
-            discord.app_commands.Choice(name=name, value=name) for name in filtered[:25]
-        ]
-        return leaderboard_name_cache[current]
-    except Exception as e:
-        logger.exception("Error in leaderboard autocomplete", exc_info=e)
-        return []
-
-
 class LeaderboardDB:
     def __init__(self, host: str, database: str, user: str, password: str, port: str = "5432"):
         """Initialize database connection parameters"""
@@ -68,6 +37,7 @@ def __init__(self, host: str, database: str, user: str, password: str, port: str
         self.connection: Optional[psycopg2.extensions.connection] = None
         self.refcount: int = 0
         self.cursor: Optional[psycopg2.extensions.cursor] = None
+        self.name_cache = LRUCache(max_size=512)
 
     def connect(self) -> bool:
         """Establish connection to the database"""
@@ -143,7 +113,7 @@ def create_leaderboard(self, leaderboard: LeaderboardItem) -> int:
                 )
 
             self.connection.commit()
-            leaderboard_name_cache.invalidate()  # Invalidate autocomplete cache
+            self.name_cache.invalidate()  # Invalidate autocomplete cache
             return leaderboard_id
         except psycopg2.Error as e:
             logger.exception("Error in leaderboard creation.", e)
@@ -205,7 +175,7 @@ def delete_leaderboard(self, leaderboard_name: str, force: bool = False):
                 (leaderboard_name,),
             )
             self.connection.commit()
-            leaderboard_name_cache.invalidate()  # Invalidate autocomplete cache
+            self.name_cache.invalidate()  # Invalidate autocomplete cache
         except psycopg2.Error as e:
             self.connection.rollback()
             logger.exception("Could not delete leaderboard %s.", leaderboard_name, exc_info=e)
@@ -923,23 +893,3 @@ def validate_cli_id(self, cli_id: str) -> Optional[dict[str, str]]:
             self.connection.rollback()
             logger.exception("Error validating CLI ID %s", cli_id, exc_info=e)
             raise KernelBotError("Error validating CLI ID") from e
-
-
-if __name__ == "__main__":
-    print(
-        POSTGRES_HOST,
-        POSTGRES_DATABASE,
-        POSTGRES_USER,
-        POSTGRES_PASSWORD,
-        POSTGRES_PORT,
-    )
-
-    leaderboard_db = LeaderboardDB(
-        POSTGRES_HOST,
-        POSTGRES_DATABASE,
-        POSTGRES_USER,
-        POSTGRES_PASSWORD,
-        POSTGRES_PORT,
-    )
-    leaderboard_db.connect()
-    leaderboard_db.disconnect()

src/discord-cluster-manager/utils.py

@@ -31,6 +31,9 @@ def setup_logging(name: Optional[str] = None):
     return logger
 
 
+logger = setup_logging(__name__)
+
+
 def with_error_handling(f: callable):
     @functools.wraps(f)
     async def wrap(self, interaction: discord.Interaction, *args, **kwargs):
@@ -336,3 +339,27 @@ def format_time(value: float | str, err: Optional[float | str] = None, scale=Non
             return f"{value:.0f} ± {err:.1f} {unit}"
         else:
             return f"{value:.0f} {unit}"
+
+
+async def leaderboard_name_autocomplete(
+    interaction: discord.Interaction,
+    current: str,
+) -> list[discord.app_commands.Choice[str]]:
+    """Return leaderboard names that match the current typed name"""
+    try:
+        bot = interaction.client
+        name_cache = bot.leaderboard_db.name_cache
+        cached_value = name_cache[current]
+        if cached_value is not None:
+            return cached_value
+
+        with bot.leaderboard_db as db:
+            leaderboards = db.get_leaderboard_names()
+        filtered = [lb for lb in leaderboards if current.lower() in lb.lower()]
+        name_cache[current] = [
+            discord.app_commands.Choice(name=name, value=name) for name in filtered[:25]
+        ]
+        return name_cache[current]
+    except Exception as e:
+        logger.exception("Error in leaderboard autocomplete", exc_info=e)
+        return []