Merge pull request #41 from gpu-mode/alzhang/leaderboards

docker-compose.yml

@@ -0,0 +1,18 @@
+version: '3.8'
+
+services:
+  postgres:
+    image: postgres:latest
+    container_name: postgres_db
+    environment:
+      - POSTGRES_USER=postgres
+      - POSTGRES_PASSWORD=postgres
+      - POSTGRES_DB=mydb
+    ports:
+      - "5432:5432"
+    volumes:
+      - postgres_data:/var/lib/postgresql/data
+    restart: unless-stopped
+
+volumes:
+  postgres_data:

requirements.txt

@@ -4,4 +4,5 @@ discord.py
 audioop-lts # discord.py imports using * syntax 
 python-dotenv
 requests
-modal
+modal
+psycopg2-binary

scripts/flush_db.py

@@ -0,0 +1,53 @@
+#!/usr/bin/env python3
+
+import psycopg2
+from psycopg2 import Error
+from dotenv import load_dotenv
+import os
+
+def flush_database():
+    # Load environment variables
+    load_dotenv()
+
+    # Get database connection parameters from environment
+    connection_params = {
+        "host": os.getenv("POSTGRES_HOST"),
+        "database": os.getenv("POSTGRES_DATABASE"),
+        "user": os.getenv("POSTGRES_USER"),
+        "password": os.getenv("POSTGRES_PASSWORD"),
+        "port": os.getenv("POSTGRES_PORT", "5432")
+    }
+
+    # Verify all parameters are present
+    missing_params = [k for k, v in connection_params.items() if not v]
+    if missing_params:
+        print(f"❌ Missing environment variables: {', '.join(missing_params)}")
+        return
+
+    try:
+        # Connect to database
+        print("📡 Connecting to database...")
+        connection = psycopg2.connect(**connection_params)
+        cursor = connection.cursor()
+
+        # Drop existing tables
+        print("🗑️  Dropping existing tables...")
+        drop_tables_query = """
+        DROP TABLE IF EXISTS submissions CASCADE;
+        DROP TABLE IF EXISTS leaderboard CASCADE;
+        """
+        cursor.execute(drop_tables_query)
+        # Commit changes
+        connection.commit()
+        print("✅ Database flushed and recreated successfully!")
+
+    except Error as e:
+        print(f"❌ Database error: {e}")
+    finally:
+        if 'connection' in locals():
+            cursor.close()
+            connection.close()
+            print("🔌 Database connection closed")
+
+if __name__ == "__main__":
+    flush_database() 

src/discord-cluster-manager/bot.py

@@ -11,9 +11,16 @@
     DISCORD_DEBUG_TOKEN,
     DISCORD_CLUSTER_STAGING_ID,
     DISCORD_DEBUG_CLUSTER_STAGING_ID,
+    POSTGRES_USER,
+    POSTGRES_PASSWORD,
+    POSTGRES_HOST,
+    POSTGRES_PORT,
+    POSTGRES_DATABASE,
 )
 from cogs.modal_cog import ModalCog
 from cogs.github_cog import GitHubCog
+from cogs.leaderboard_cog import LeaderboardCog
+from leaderboard_db import LeaderboardDB
 
 logger = setup_logging()
 
@@ -31,13 +38,27 @@ def __init__(self, debug_mode=False):
         )
         self.tree.add_command(self.run_group)
 
+        self.leaderboard_group = app_commands.Group(
+            name="leaderboard", description="Leaderboard commands"
+        )
+        self.tree.add_command(self.leaderboard_group)
+
+        self.leaderboard_db = LeaderboardDB(
+            POSTGRES_HOST,
+            POSTGRES_DATABASE,
+            POSTGRES_USER,
+            POSTGRES_PASSWORD,
+            POSTGRES_PORT,
+        )
+
     async def setup_hook(self):
         logger.info(f"Syncing commands for staging guild {DISCORD_CLUSTER_STAGING_ID}")
         try:
             # Load cogs
             await self.add_cog(ModalCog(self))
             await self.add_cog(GitHubCog(self))
             await self.add_cog(BotManagerCog(self))
+            await self.add_cog(LeaderboardCog(self))
 
             guild_id = (
                 DISCORD_CLUSTER_STAGING_ID
@@ -96,7 +117,7 @@ async def send_chunked_message(
         for i, chunk in enumerate(chunks):
             if code_block:
                 await channel.send(
-                    f"```\nOutput (part {i+1}/{len(chunks)}):\n{chunk}\n```"
+                    f"```\nOutput (part {i + 1}/{len(chunks)}):\n{chunk}\n```"
                 )
             else:
                 await channel.send(chunk)

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

@@ -0,0 +1,257 @@
+import discord
+from discord import app_commands
+from discord.ext import commands
+from datetime import datetime
+
+from typing import TYPE_CHECKING
+from consts import GitHubGPU, ModalGPU
+
+import random
+
+if TYPE_CHECKING:
+    from bot import ClusterBot
+
+
+class LeaderboardSubmitCog(app_commands.Group):
+    def __init__(self, bot):
+        self.bot: ClusterBot = bot
+        super().__init__(name="submit", description="Submit leaderboard data")
+
+    # Parent command that defines global options
+    @app_commands.describe(
+        leaderboard_name="Name of the competition / kernel to optimize",
+        script="The Python / CUDA script file to run",
+        dtype="dtype (e.g. FP32, BF16, FP4) that the input and output expects.",
+        shape="Data input shape as a tuple",
+    )
+    # TODO: Modularize this so all the write functionality is in here. Haven't figured
+    # a good way to do this yet.
+    async def submit(
+        self,
+        interaction: discord.Interaction,
+        leaderboard_name: str,
+        script: discord.Attachment,
+        dtype: app_commands.Choice[str] = None,
+        shape: app_commands.Choice[str] = None,
+    ):
+        pass
+
+    @app_commands.command(name="modal", description="Submit leaderboard data for modal")
+    @app_commands.describe(
+        gpu_type="Choose the GPU type for Modal",
+    )
+    @app_commands.choices(
+        gpu_type=[
+            app_commands.Choice(name=gpu.value, value=gpu.value) for gpu in ModalGPU
+        ]
+    )
+    async def modal(
+        self,
+        interaction: discord.Interaction,
+        leaderboard_name: str,
+        script: discord.Attachment,
+        gpu_type: app_commands.Choice[str],
+        dtype: app_commands.Choice[str] = "fp32",
+        shape: app_commands.Choice[str] = None,
+    ):
+        try:
+            # Read the template file
+            submission_content = await script.read()
+
+            # Compute eval or submission score, call runner here.
+            score = random.random()
+
+            with self.bot.leaderboard_db as db:
+                db.create_submission({
+                    "submission_name": script.filename,
+                    "submission_time": datetime.now(),
+                    "leaderboard_name": leaderboard_name,
+                    "code": submission_content,
+                    "user_id": interaction.user.id,
+                    "submission_score": score,
+                })
+
+            await interaction.response.send_message(
+                f"Ran on Modal. Leaderboard '{leaderboard_name}'. Submission title: {script.filename}. Submission user: {interaction.user.id}. Runtime: {score} ms",
+                ephemeral=True,
+            )
+        except ValueError:
+            await interaction.response.send_message(
+                "Invalid date format. Please use YYYY-MM-DD or YYYY-MM-DD HH:MM",
+                ephemeral=True,
+            )
+
+    ### GITHUB SUBCOMMAND
+    @app_commands.command(
+        name="github", description="Submit leaderboard data for GitHub"
+    )
+    @app_commands.describe(
+        gpu_type="Choose the GPU type for Github Runners",
+    )
+    @app_commands.choices(
+        gpu_type=[
+            app_commands.Choice(name=gpu.value, value=gpu.value) for gpu in GitHubGPU
+        ]
+    )
+    async def github(
+        self,
+        interaction: discord.Interaction,
+        leaderboard_name: str,
+        script: discord.Attachment,
+        gpu_type: app_commands.Choice[str],
+        dtype: app_commands.Choice[str] = "fp32",
+        shape: app_commands.Choice[str] = None,
+    ):
+        try:
+            # Read the template file
+            submission_content = await script.read()
+
+            # Compute eval or submission score, call runner here.
+            score = random.random()
+
+            with self.bot.leaderboard_db as db:
+                db.create_submission({
+                    "submission_name": script.filename,
+                    "submission_time": datetime.now(),
+                    "leaderboard_name": leaderboard_name,
+                    "code": submission_content,
+                    "user_id": interaction.user.id,
+                    "submission_score": score,
+                })
+
+            await interaction.response.send_message(
+                f"Ran on GH. Leaderboard '{leaderboard_name}'. Submission title: {script.filename}. Submission user: {interaction.user.id}. Runtime: {score} ms",
+                ephemeral=True,
+            )
+        except ValueError:
+            await interaction.response.send_message(
+                "Invalid date format. Please use YYYY-MM-DD or YYYY-MM-DD HH:MM",
+                ephemeral=True,
+            )
+
+
+class LeaderboardCog(commands.Cog):
+    def __init__(self, bot):
+        self.bot: ClusterBot = bot
+        self.get_leaderboards = bot.leaderboard_group.command(name="get")(
+            self.get_leaderboards
+        )
+        self.leaderboard_create = bot.leaderboard_group.command(
+            name="create", description="Create a new leaderboard"
+        )(self.leaderboard_create)
+
+        # self.leaderboard_submit = bot.leaderboard_group.command(
+        #     name="submit", description="Submit a file to the leaderboard"
+        # )(self.leaderboard_submit)
+
+        bot.leaderboard_group.add_command(LeaderboardSubmitCog(bot))
+
+        self.get_leaderboard_submissions = bot.leaderboard_group.command(
+            name="show", description="Get all submissions for a leaderboard"
+        )(self.get_leaderboard_submissions)
+
+    async def get_leaderboards(self, interaction: discord.Interaction):
+        """Display all leaderboards in a table format"""
+        await interaction.response.defer()
+
+        with self.bot.leaderboard_db as db:
+            leaderboards = db.get_leaderboards()
+
+        if not leaderboards:
+            await interaction.followup.send("No leaderboards found.", ephemeral=True)
+            return
+
+        # Create embed
+        embed = discord.Embed(title="Active Leaderboards", color=discord.Color.blue())
+
+        # Add fields for each leaderboard
+        for lb in leaderboards:
+            deadline_str = lb["deadline"].strftime("%Y-%m-%d %H:%M")
+            embed.add_field(
+                name=lb["name"], value=f"Deadline: {deadline_str}", inline=False
+            )
+
+        await interaction.followup.send(embed=embed)
+
+    @discord.app_commands.describe(
+        leaderboard_name="Name of the leaderboard",
+        deadline="Competition deadline in the form: 'Y-m-d'",
+        reference_code="Reference implementation of kernel. Also includes eval code.",
+    )
+    async def leaderboard_create(
+        self,
+        interaction: discord.Interaction,
+        leaderboard_name: str,
+        deadline: str,
+        reference_code: discord.Attachment,
+    ):
+        try:
+            # Try parsing with time first
+            try:
+                date_value = datetime.strptime(deadline, "%Y-%m-%d %H:%M")
+            except ValueError:
+                # If that fails, try parsing just the date (will set time to 00:00)
+                date_value = datetime.strptime(deadline, "%Y-%m-%d")
+
+            # Read the template file
+            template_content = await reference_code.read()
+
+            with self.bot.leaderboard_db as db:
+                print(
+                    leaderboard_name,
+                    type(date_value),
+                    type(template_content.decode("utf-8")),
+                )
+                db.create_leaderboard({
+                    "name": leaderboard_name,
+                    "deadline": date_value,
+                    "reference_code": template_content.decode("utf-8"),
+                })
+
+            await interaction.response.send_message(
+                f"Leaderboard '{leaderboard_name}'. Submission deadline: {date_value}",
+                ephemeral=True,
+            )
+        except ValueError:
+            await interaction.response.send_message(
+                "Invalid date format. Please use YYYY-MM-DD or YYYY-MM-DD HH:MM",
+                ephemeral=True,
+            )
+
+    @discord.app_commands.describe(leaderboard_name="Name of the leaderboard")
+    async def get_leaderboard_submissions(
+        self,
+        interaction: discord.Interaction,
+        leaderboard_name: str,
+        dtype: app_commands.Choice[str] = "fp32",
+    ):
+        with self.bot.leaderboard_db as db:
+            leaderboard_id = db.get_leaderboard_id(leaderboard_name)
+            if not leaderboard_id:
+                await interaction.response.send_message(
+                    "Leaderboard not found.", ephemeral=True
+                )
+                return
+
+            # submissions = db.get_leaderboard_submissions(leaderboard_id)  # Add dtype
+            submissions = db.get_leaderboard_submissions(leaderboard_name)  # Add dtype
+
+        if not submissions:
+            await interaction.response.send_message(
+                "No submissions found.", ephemeral=True
+            )
+            return
+
+        # Create embed
+        embed = discord.Embed(
+            title="Leaderboard Submissions", color=discord.Color.blue()
+        )
+
+        for submission in submissions:
+            embed.add_field(
+                name=f"{submission['user_id']}: {submission['submission_name']}",
+                value=f"Submission speed: {submission['submission_score']}",
+                inline=False,
+            )
+
+        await interaction.response.send_message(embed=embed)