moved all discord-specific utilities into their own file

Author: ngc92

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

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

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

@@ -10,6 +10,12 @@
 )
 from discord import app_commands
 from discord.ext import commands
+from discord_utils import (
+    get_user_from_id,
+    leaderboard_name_autocomplete,
+    send_discord_message,
+    with_error_handling,
+)
 from leaderboard_db import (
     LeaderboardItem,
     LeaderboardRankedEntry,
@@ -22,11 +28,7 @@
 from ui.table import create_table
 from utils import (
     format_time,
-    get_user_from_id,
-    leaderboard_name_autocomplete,
-    send_discord_message,
     setup_logging,
-    with_error_handling,
 )
 
 if TYPE_CHECKING:

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

@@ -5,8 +5,9 @@
 import psycopg2
 from discord import app_commands
 from discord.ext import commands
+from discord_utils import send_discord_message
 from env import DATABASE_URL, DISABLE_SSL
-from utils import send_discord_message, setup_logging
+from utils import setup_logging
 
 if TYPE_CHECKING:
     from ..bot import ClusterBot

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

@@ -11,6 +11,7 @@
 from consts import GPU, GPU_TO_SM, RankCriterion, SubmissionMode, get_gpu_by_name
 from discord import app_commands
 from discord.ext import commands
+from discord_utils import send_discord_message, with_error_handling
 from report import (
     MultiProgressReporter,
     RunProgressReporter,
@@ -21,9 +22,7 @@
 from task import LeaderboardTask, build_task_config
 from utils import (
     KernelBotError,
-    send_discord_message,
     setup_logging,
-    with_error_handling,
 )
 
 logger = setup_logging()

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

@@ -14,10 +14,11 @@
 from discord import app_commands
 from discord.app_commands import Choice
 from discord.ext import commands
+from discord_utils import send_discord_message, with_error_handling
 from leaderboard_db import RunItem, SubmissionItem
 from report import MultiProgressReporter
 from task import make_task
-from utils import send_discord_message, setup_logging, with_error_handling
+from utils import setup_logging
 
 logger = setup_logging()
 

src/discord-cluster-manager/discord_utils.py

@@ -0,0 +1,105 @@
+import functools
+import logging
+
+import discord
+from utils import KernelBotError, setup_logging
+
+logger = setup_logging(__name__)
+
+
+def with_error_handling(f: callable):
+    @functools.wraps(f)
+    async def wrap(self, interaction: discord.Interaction, *args, **kwargs):
+        try:
+            await f(self, interaction, *args, **kwargs)
+        except KernelBotError as e:
+            await send_discord_message(
+                interaction,
+                str(e),
+                ephemeral=True,
+            )
+        except Exception as e:
+            logging.exception("Unhandled exception %s", e, exc_info=e)
+            await send_discord_message(
+                interaction,
+                "An unexpected error occurred. Please report this to the developers.",
+                ephemeral=True,
+            )
+
+    return wrap
+
+
+async def get_user_from_id(bot, id) -> str:
+    with bot.leaderboard_db as db:
+        return db.get_user_from_id(id) or id
+
+
+async def send_discord_message(
+    interaction: discord.Interaction, msg: str, *, ephemeral=False, **kwargs
+) -> None:
+    """
+    To get around response messages in slash commands that are
+    called externally, send a message using the followup.
+    """
+    if interaction.response.is_done():
+        await interaction.followup.send(msg, ephemeral=ephemeral, **kwargs)
+    else:
+        await interaction.response.send_message(msg, ephemeral=ephemeral, **kwargs)
+
+
+async def send_logs(thread: discord.Thread, logs: str) -> None:
+    """Send logs to a Discord thread, splitting by lines and respecting Discord's character limit.
+
+    Args:
+        thread: The Discord thread to send logs to
+        logs: The log string to send
+    """
+    # Split logs into lines
+    log_lines = logs.splitlines()
+
+    current_chunk = []
+    current_length = 0
+
+    for line in log_lines:
+        # Add 1 for the newline character
+        line_length = len(line) + 1
+
+        # If adding this line would exceed Discord's limit, send current chunk
+        if current_length + line_length > 1990:  # Leave room for code block markers
+            if current_chunk:
+                chunk_text = "\n".join(current_chunk)
+                await thread.send(f"```\n{chunk_text}\n```")
+                current_chunk = []
+                current_length = 0
+
+        current_chunk.append(line)
+        current_length += line_length
+
+    # Send any remaining lines
+    if current_chunk:
+        chunk_text = "\n".join(current_chunk)
+        await thread.send(f"```\n{chunk_text}\n```")
+
+
+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 []

src/discord-cluster-manager/ui/misc.py

@@ -2,7 +2,8 @@
 
 import discord
 from discord import Interaction, SelectOption, ui
-from utils import KernelBotError, send_discord_message
+from discord_utils import send_discord_message
+from utils import KernelBotError
 
 
 class GPUSelectionView(ui.View):

src/discord-cluster-manager/utils.py

@@ -1,10 +1,7 @@
-import functools
 import logging
 import subprocess
 from typing import Any, Optional
 
-import discord
-
 
 def setup_logging(name: Optional[str] = None):
     """Configure and setup logging for the application"""
@@ -26,31 +23,6 @@ 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):
-        try:
-            await f(self, interaction, *args, **kwargs)
-        except KernelBotError as e:
-            await send_discord_message(
-                interaction,
-                str(e),
-                ephemeral=True,
-            )
-        except Exception as e:
-            logging.exception("Unhandled exception %s", e, exc_info=e)
-            await send_discord_message(
-                interaction,
-                "An unexpected error occurred. Please report this to the developers.",
-                ephemeral=True,
-            )
-
-    return wrap
-
-
 class KernelBotError(Exception):
     """
     This class represents an Exception that has been sanitized,
@@ -75,58 +47,6 @@ def get_github_branch_name():
         return "main"
 
 
-async def get_user_from_id(bot, id) -> str:
-    with bot.leaderboard_db as db:
-        return db.get_user_from_id(id) or id
-
-
-async def send_discord_message(
-    interaction: discord.Interaction, msg: str, *, ephemeral=False, **kwargs
-) -> None:
-    """
-    To get around response messages in slash commands that are
-    called externally, send a message using the followup.
-    """
-    if interaction.response.is_done():
-        await interaction.followup.send(msg, ephemeral=ephemeral, **kwargs)
-    else:
-        await interaction.response.send_message(msg, ephemeral=ephemeral, **kwargs)
-
-
-async def send_logs(thread: discord.Thread, logs: str) -> None:
-    """Send logs to a Discord thread, splitting by lines and respecting Discord's character limit.
-
-    Args:
-        thread: The Discord thread to send logs to
-        logs: The log string to send
-    """
-    # Split logs into lines
-    log_lines = logs.splitlines()
-
-    current_chunk = []
-    current_length = 0
-
-    for line in log_lines:
-        # Add 1 for the newline character
-        line_length = len(line) + 1
-
-        # If adding this line would exceed Discord's limit, send current chunk
-        if current_length + line_length > 1990:  # Leave room for code block markers
-            if current_chunk:
-                chunk_text = "\n".join(current_chunk)
-                await thread.send(f"```\n{chunk_text}\n```")
-                current_chunk = []
-                current_length = 0
-
-        current_chunk.append(line)
-        current_length += line_length
-
-    # Send any remaining lines
-    if current_chunk:
-        chunk_text = "\n".join(current_chunk)
-        await thread.send(f"```\n{chunk_text}\n```")
-
-
 class LRUCache:
     def __init__(self, max_size: int):
         """LRU Cache implementation, as functools.lru doesn't work in async code
@@ -215,27 +135,3 @@ 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 []