Merge branch 'main' of https://github.com/gpu-mode/discord-cluster-manager

Author: msaroufim

README.md

@@ -141,7 +141,10 @@ Create a `.env` file with the following environment variables:
 - `DISCORD_DEBUG_CLUSTER_STAGING_ID` : The ID of the "staging" server you want to connect to
 - `DISCORD_CLUSTER_STAGING_ID` : The ID of the "production" server you want to connect to
 - `GITHUB_TOKEN` : A Github token with permissions to trigger workflows, for now only new branches from [discord-cluster-manager](https://github.com/gpu-mode/discord-cluster-manager) are tested, since the bot triggers workflows on your behalf
+- `GITHUB_REPO` : The repository where the cluster manager is hosted.
+- `GITHUB_WORKFLOW_BRANCH` : The branch to start the GitHub Actions jobs from when submitting a task.
 - `DATABASE_URL` : The URL you use to connect to Postgres.
+- `DISABLE_SSL` : (Optional) set if you want to disable SSL when connecting to Postgres.
 
 Below is where to find these environment variables:
 
@@ -169,14 +172,22 @@ Below is where to find these environment variables:
       <img width="1440" alt="Screenshot 2024-12-30 at 8 51 59 AM" src="https://github.com/user-attachments/assets/e3467871-bd2c-4f94-b0c5-c8a6ef5ce89e">
   </details>
 
+- `GITHUB_REPO`: This should be set to this repository, which is usually `gpu-mode/discord-cluster-manager`.
+
+- `GITHUB_WORKFLOW_BRANCH`: Usually `main` or the branch you are working from.
+
 - `DATABASE_URL`: This contains the connection details for your local database, and has the form `postgresql://user:password@localhost/clusterdev`.
 
+- `DISABLE_SSL`: Set to `1` when developing.
+
 ### Verify Setup
 
+Install the kernel bot as editable using `pip install -e .`
+
 Run the following command to run the bot:
 
 ```
-python src/discord-cluster-manager/bot.py --debug
+python src/kernelbot/main.py --debug
 ```
 
 Then in your staging server, use the `/verifyruns` command to test basic functionalities of the bot and the `/verifydb` command to check database connectivity.
@@ -232,7 +243,7 @@ specify the available GPUs that the leaderboard evaluates on.
 The Discord bot internally contains an `eval.py` script that handles the correctness and timing
 analysis for the leaderboard. The `reference_code` that the leaderboard creator submits must have
 the following function signatures with their implementations filled out. `InputType` and
-`OutputType` are generics that could be a `torch.Tensor`, `List[torch.Tensor]`, etc. 
+`OutputType` are generics that could be a `torch.Tensor`, `List[torch.Tensor]`, etc.
 depending on the reference code specifications. We leave this flexibility to the leaderboard creator.
 
 ```python
@@ -257,8 +268,8 @@ handle the typing system for tensors. The `reference.cu` that the leaderboard cr
 the following function signatures with their implementations filled out:
 
 The main difference is we now need to define an alias for the type that the input / outputs are. A
-simple and common example is a list of FP32 tensors, which can be defined using a pre-defined array of 
-`const int`s called `N_SIZES`, then define an array of containers, e.g. 
+simple and common example is a list of FP32 tensors, which can be defined using a pre-defined array of
+`const int`s called `N_SIZES`, then define an array of containers, e.g.
 `std::array<std::vector<float>, N_SIZES>`.
 
 ```cuda
@@ -293,7 +304,7 @@ bool check_implementation(output_t out, output_t ref) {
 ```
 
 The leaderboard submission for _Python code_ requires the following function signatures, where
-`InputType` and `OutputType` are generics that could be a `torch.Tensor`, `List[torch.Tensor]`, 
+`InputType` and `OutputType` are generics that could be a `torch.Tensor`, `List[torch.Tensor]`,
 etc. depending on the reference code specifications.
 
 ```python
@@ -354,7 +365,7 @@ If you'd like to donate a GPU to our efforts, we can make you a CI admin in Gith
 
 ## Citation
 
-If you used our software please cite it as 
+If you used our software please cite it as
 
 ```
 @misc{kernelbot2025,

docker-compose.test.yml

@@ -0,0 +1,32 @@
+services:
+  postgres-test:
+    image: postgres:15
+    container_name: postgres_db_test
+    environment:
+      POSTGRES_USER: postgres
+      POSTGRES_PASSWORD: postgres
+      POSTGRES_DB: clusterdev
+    ports:
+      - "5433:5432"
+    tmpfs:
+      - /var/lib/postgresql/data
+    healthcheck:
+      test: ["CMD-SHELL", "pg_isready -U postgres -d clusterdev"]
+      interval: 5s
+      timeout: 3s
+      retries: 10
+
+  migrate-test:
+    image: python:3.11-slim
+    depends_on:
+      postgres-test:
+        condition: service_healthy
+    volumes:
+      - ./src/migrations:/migrations
+    command: >
+      sh -c "
+        pip install yoyo-migrations psycopg2-binary &&
+        yoyo apply -b -d 'postgresql://postgres:postgres@postgres-test:5432/clusterdev' -v /migrations/ &&
+        yoyo list -d 'postgresql://postgres:postgres@postgres-test:5432/clusterdev' /migrations/
+      "
+    restart: "no"

pyproject.toml

@@ -44,6 +44,8 @@ relative_files = true
 exclude_lines = [
    "pragma: no cover",
    "raise NotImplementedError",
+    # For now, don't require coverage of db errors
+   "except psycopg2.Error"
 ]
 
 [tool.pytest.ini_options]

src/kernelbot/cogs/admin_cog.py

@@ -184,7 +184,7 @@ async def leaderboard_create_local(
             try:
                 old_lb = db.get_leaderboard(leaderboard_name)
             except LeaderboardDoesNotExist:
-                pass
+                old_lb = None
             db.delete_leaderboard(leaderboard_name, force=True)
 
         # get existing forum thread or create new one

src/kernelbot/cogs/verify_run_cog.py

@@ -53,20 +53,20 @@ async def trigger_run(self, interaction: discord.Interaction, gpu: GPU, reporter
             sub_code = create_mock_attachment(
                 "submission.py", Path("examples/identity_py/submission.py").read_text()
             )
-            task = make_task_definition("examples/identity_py")
+            leaderboard = make_task_definition("examples/identity_py")
         else:
             sub_code = create_mock_attachment(
                 "test.cu", Path("examples/identity_cuda/submission.cu").read_text()
             )
-            task = make_task_definition("examples/identity_cuda")
+            leaderboard = make_task_definition("examples/identity_cuda")
 
         return await submit_leaderboard(
             interaction,
             -1,
             sub_code,
             gpu,
             reporter=reporter,
-            task=task,
+            task=leaderboard.task,
             mode=SubmissionMode.TEST,
             seed=None,
         )
@@ -292,8 +292,8 @@ async def verify_runs(self, interaction: discord.Interaction):
             amd = get_gpu_by_name("mi300")
             t4 = get_gpu_by_name("T4")
 
-            reporter = MultiProgressReporterDiscord("Verifying")
-            await reporter.show(interaction)
+            reporter = MultiProgressReporterDiscord(interaction)
+            await reporter.show("Verifying")
 
             results = await asyncio.gather(
                 self.verify_github_run(interaction, nvidia, reporter.add_run("NVIDIA-PY"), "py"),

src/libkernelbot/leaderboard_db.py

@@ -87,6 +87,8 @@ def create_leaderboard(
         forum_id: int,
         gpu_types: list | str,
     ) -> int:
+        # to prevent surprises, ensure we have specified a timezone
+        assert deadline.tzinfo is not None
         try:
             task = definition.task
             self.cursor.execute(
@@ -127,10 +129,10 @@ def create_leaderboard(
             self.name_cache.invalidate()  # Invalidate autocomplete cache
             return leaderboard_id
         except psycopg2.Error as e:
-            logger.exception("Error in leaderboard creation.", e)
+            logger.exception("Error in leaderboard creation.", exc_info=e)
             if isinstance(e, psycopg2.errors.UniqueViolation):
                 raise KernelBotError(
-                    "Error: Tried to create a leaderboard " f'"{name}" that already exists.'
+                    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
@@ -140,31 +142,21 @@ def update_leaderboard(
     ):
         task = definition.task
         try:
+            lb_id = self.get_leaderboard_id(name)
             self.cursor.execute(
                 """
                 UPDATE leaderboard.leaderboard
                 SET deadline = %s, task = %s, description = %s
-                WHERE name = %s;
+                WHERE id = %s;
                 """,
                 (
                     deadline.astimezone(datetime.timezone.utc),
                     task.to_str(),
                     definition.description,
-                    name,
+                    lb_id,
                 ),
             )
 
-            self.cursor.execute(
-                """
-                SELECT id
-                FROM leaderboard.leaderboard
-                WHERE name = %s
-                """,
-                (name,),
-            )
-
-            lb_id = self.cursor.fetchone()[0]
-
             # replace templates
             self.cursor.execute(
                 """
@@ -235,8 +227,13 @@ def delete_leaderboard(self, leaderboard_name: str, force: bool = False):
             self.name_cache.invalidate()  # Invalidate autocomplete cache
         except psycopg2.Error as e:
             self.connection.rollback()
+            if isinstance(e, psycopg2.errors.ForeignKeyViolation):
+                raise KernelBotError(
+                    f"Could not delete leaderboard `{leaderboard_name}` with existing submissions."
+                ) from e
+
             logger.exception("Could not delete leaderboard %s.", leaderboard_name, exc_info=e)
-            raise KernelBotError(f"Could not delete leaderboard {leaderboard_name}.") from e
+            raise KernelBotError(f"Could not delete leaderboard `{leaderboard_name}`.") from e
 
     def create_submission(
         self,
@@ -260,7 +257,7 @@ def create_submission(
 
             code_id = None
             for candidate in self.cursor.fetchall():
-                if candidate[1] == code:
+                if bytes(candidate[1]).decode("utf-8") == code:
                     code_id = candidate[0]
                     break
 
@@ -357,6 +354,24 @@ def create_submission_run(
             if compilation is not None:
                 compilation = json.dumps(dataclasses.asdict(compilation))
 
+            # check validity
+            self.cursor.execute(
+                """
+            SELECT done FROM leaderboard.submission WHERE id = %s
+            """,
+                (submission,),
+            )
+            if self.cursor.fetchone()[0]:
+                logger.error(
+                    "Submission '%s' is already marked as done when trying to add %s run.",
+                    submission,
+                    mode,
+                )
+                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"]
@@ -408,7 +423,7 @@ def get_leaderboard_names(self, active_only: bool = False) -> list[str]:
     def get_leaderboards(self) -> list["LeaderboardItem"]:
         self.cursor.execute(
             """
-            SELECT id, name, deadline, task, creator_id, forum_id, description
+            SELECT id, name, deadline, task, creator_id, forum_id, description, secret_seed
             FROM leaderboard.leaderboard
             """
         )
@@ -432,6 +447,7 @@ def get_leaderboards(self) -> list["LeaderboardItem"]:
                     creator_id=lb[4],
                     forum_id=lb[5],
                     description=lb[6],
+                    secret_seed=lb[7],
                 )
             )
 
@@ -461,7 +477,7 @@ def get_leaderboard_gpu_types(self, leaderboard_name: str) -> List[str]:
 
         return [x[0] for x in self.cursor.fetchall()]
 
-    def get_leaderboard_templates(self, leaderboard_name: str) -> Dict[str, str]:
+    def get_leaderboard_id(self, leaderboard_name: str) -> int:
         self.cursor.execute(
             """
             SELECT id
@@ -473,14 +489,18 @@ def get_leaderboard_templates(self, leaderboard_name: str) -> Dict[str, str]:
         lb_id = self.cursor.fetchone()
         if lb_id is None:
             raise LeaderboardDoesNotExist(leaderboard_name)
+        return lb_id[0]
+
+    def get_leaderboard_templates(self, leaderboard_name: str) -> Dict[str, str]:
+        lb_id = self.get_leaderboard_id(leaderboard_name)
 
         self.cursor.execute(
             """
             SELECT lang, code
             FROM leaderboard.templates
             WHERE leaderboard_id = %s
             """,
-            (lb_id[0],),
+            (lb_id,),
         )
 
         return {x[0]: x[1] for x in self.cursor.fetchall()}
@@ -585,7 +605,7 @@ def get_leaderboard_submissions(
 
         self.cursor.execute(query, args)
 
-        return [
+        result = [
             LeaderboardRankedEntry(
                 submission_name=submission[0],
                 submission_id=submission[1],
@@ -599,6 +619,19 @@ def get_leaderboard_submissions(
             )
             for submission in self.cursor.fetchall()
         ]
+        if len(result) == 0:
+            # try to diagnose why we didn't get anything
+            # this will raise if the LB does not exist at all.
+            self.get_leaderboard_id(leaderboard_name)
+
+            # 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}'"
+                )
+
+        return result
 
     def generate_stats(self, last_day: bool):
         try:
@@ -784,7 +817,7 @@ def get_submission_by_id(self, submission_id: int) -> Optional["SubmissionItem"]
             user_id=submission[3],
             submission_time=submission[4],
             done=submission[5],
-            code=submission[6],
+            code=bytes(submission[6]).decode("utf-8"),
             runs=runs,
         )
 
@@ -824,7 +857,20 @@ def get_leaderboard_submission_count(
             args = (leaderboard_name, gpu_name)
 
         self.cursor.execute(query, args)
-        return self.cursor.fetchone()[0]
+        count = self.cursor.fetchone()[0]
+        if count == 0:
+            # try to diagnose why we didn't get anything
+            # this will raise if the LB does not exist at all.
+            self.get_leaderboard_id(leaderboard_name)
+
+            # 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}'"
+                )
+
+        return count
 
     def init_user_from_cli(self, cli_id: str, auth_provider: str):
         """