unified API and discord submit paths

Author: ngc92

src/discord-cluster-manager/api/main.py

@@ -227,7 +227,6 @@ async def cli_auth(auth_provider: str, code: str, state: str, db_context=Depends
 
 async def _stream_submission_response(
     submission_request: SubmissionRequest,
-    user_info: dict,
     submission_mode_enum: SubmissionMode,
     backend: KernelBackend,
 ):
@@ -237,7 +236,6 @@ async def _stream_submission_response(
         task = asyncio.create_task(
             _run_submission(
                 submission_request,
-                user_info,
                 submission_mode_enum,
                 backend,
             )
@@ -398,6 +396,7 @@ async def run_submission(  # noqa: C901
             code=submission_code,
             file_name=file.filename or "submission.py",
             user_id=user_id,
+            user_name=user_name,
             gpus=[gpu_type],
             leaderboard=leaderboard_name,
         )
@@ -412,7 +411,6 @@ async def run_submission(  # noqa: C901
 
     generator = _stream_submission_response(
         submission_request=submission_request,
-        user_info={"user_id": user_id, "user_name": user_name},
         submission_mode_enum=submission_mode_enum,
         backend=backend_instance,
     )

src/discord-cluster-manager/api/utils.py

@@ -1,10 +1,6 @@
-import asyncio
-from datetime import datetime
-from typing import List
-
 import requests
 from backend import KernelBackend
-from consts import SubmissionMode, get_gpu_by_name
+from consts import SubmissionMode
 from env import (
     CLI_DISCORD_CLIENT_ID,
     CLI_DISCORD_CLIENT_SECRET,
@@ -13,7 +9,7 @@
     CLI_TOKEN_URL,
 )
 from fastapi import HTTPException
-from report import Log, RunProgressReporter, RunResultReport, Text
+from report import Log, MultiProgressReporter, RunProgressReporter, RunResultReport, Text
 from submission import SubmissionRequest, prepare_submission
 
 
@@ -140,73 +136,35 @@ async def _handle_github_oauth(code: str, redirect_uri: str) -> tuple[str, str]:
 
 
 async def _run_submission(
-    submission: SubmissionRequest, user_info: dict, mode: SubmissionMode, backend: KernelBackend
+    submission: SubmissionRequest, mode: SubmissionMode, backend: KernelBackend
 ):
     try:
         req = prepare_submission(submission, backend.db)
     except Exception as e:
         raise HTTPException(status_code=400, detail=str(e)) from e
 
-    selected_gpus = [get_gpu_by_name(gpu) for gpu in req.gpus]
-    if len(selected_gpus) > 1 or selected_gpus[0] is None:
+    if len(req.gpus) != 1:
         raise HTTPException(status_code=400, detail="Invalid GPU type")
 
-    user_name = user_info["user_name"]
-    user_id = user_info["user_id"]
-
-    with backend.db as db:
-        sub_id = db.create_submission(
-            leaderboard=req.leaderboard,
-            file_name=submission.file_name,
-            code=submission.code,
-            user_id=user_id,
-            time=datetime.now(),
-            user_name=user_name,
-        )
+    reporter = MultiProgressReporterAPI()
+    sub_id, results = await backend.submit_full(req, mode, reporter)
+    return results, [rep.get_message() + "\n" + rep.long_report for rep in reporter.runs]
 
-    gpu = selected_gpus[0]
 
-    reporters: List[RunProgressReporterAPI] = []
+class MultiProgressReporterAPI(MultiProgressReporter):
+    def __init__(self):
+        self.runs = []
 
-    def add_reporter(title: str):
+    async def show(self, title: str):
+        return
+
+    def add_run(self, title: str) -> "RunProgressReporterAPI":
         rep = RunProgressReporterAPI(title)
-        reporters.append(rep)
+        self.runs.append(rep)
         return rep
 
-    try:
-        tasks = [
-            backend.submit_leaderboard(
-                sub_id,
-                submission.code,
-                submission.file_name,
-                gpu,
-                add_reporter(f"{gpu.name} on {gpu.runner}"),
-                req.task,
-                mode,
-                None,
-            )
-        ]
-
-        if mode == SubmissionMode.LEADERBOARD:
-            tasks += [
-                backend.submit_leaderboard(
-                    sub_id,
-                    submission.code,
-                    submission.file_name,
-                    gpu,
-                    add_reporter(f"{gpu.name} on {gpu.runner} (secret)"),
-                    req.task,
-                    SubmissionMode.PRIVATE,
-                    req.secret_seed,
-                )
-            ]
-
-        results = await asyncio.gather(*tasks)
-    finally:
-        with backend.db as db:
-            db.mark_submission_done(sub_id)
-
-    return results, [rep.get_message() + "\n" + rep.long_report for rep in reporters]
+    def make_message(self):
+        return
 
 
 class RunProgressReporterAPI(RunProgressReporter):

src/discord-cluster-manager/backend.py

@@ -1,13 +1,16 @@
+import asyncio
 import copy
 import math
+from datetime import datetime
 from typing import Optional
 
 import env
-from consts import GPU, GPU_TO_SM, RankCriterion, SubmissionMode
+from consts import GPU, GPU_TO_SM, RankCriterion, SubmissionMode, get_gpu_by_name
 from launchers import Launcher
 from leaderboard_db import LeaderboardDB
-from report import RunProgressReporter, generate_report, make_short_report
+from report import MultiProgressReporter, RunProgressReporter, generate_report, make_short_report
 from run_eval import FullResult
+from submission import ProcessedSubmissionRequest
 from task import LeaderboardTask, build_task_config
 from utils import KernelBotError, setup_logging
 
@@ -44,6 +47,60 @@ def register_launcher(self, launcher: Launcher):
         for gpu in launcher.gpus:
             self.launcher_map[gpu.value] = launcher
 
+    async def submit_full(
+        self, req: ProcessedSubmissionRequest, mode: SubmissionMode, reporter: MultiProgressReporter
+    ):
+        with self.db as db:
+            sub_id = db.create_submission(
+                leaderboard=req.leaderboard,
+                file_name=req.file_name,
+                code=req.code,
+                user_id=req.user_id,
+                time=datetime.now(),
+                user_name=req.user_name,
+            )
+
+        selected_gpus = [get_gpu_by_name(gpu) for gpu in req.gpus]
+
+        try:
+            tasks = [
+                self.submit_leaderboard(
+                    sub_id,
+                    req.code,
+                    req.file_name,
+                    gpu,
+                    reporter.add_run(f"{gpu.name} on {gpu.runner}"),
+                    req.task,
+                    mode,
+                    None,
+                )
+                for gpu in selected_gpus
+            ]
+
+            if mode == SubmissionMode.LEADERBOARD:
+                tasks += [
+                    self.submit_leaderboard(
+                        sub_id,
+                        req.code,
+                        req.file_name,
+                        gpu,
+                        reporter.add_run(f"{gpu.name} on {gpu.runner} (secret)"),
+                        req.task,
+                        SubmissionMode.PRIVATE,
+                        req.secret_seed,
+                    )
+                    for gpu in selected_gpus
+                ]
+            await reporter.show(
+                f"Submission **{sub_id}**: `{req.file_name}` for `{req.leaderboard}`"
+            )
+            results = await asyncio.gather(*tasks)
+        finally:
+            with self.db as db:
+                db.mark_submission_done(sub_id)
+
+        return sub_id, results
+
     async def submit_leaderboard(  # noqa: C901
         self,
         submission_id: int,

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

@@ -1,16 +1,14 @@
-import asyncio
 from datetime import datetime, timedelta
 from io import StringIO
 from typing import TYPE_CHECKING, List, Optional
 
 import discord
 from consts import (
     SubmissionMode,
-    get_gpu_by_name,
 )
 from discord import app_commands
 from discord.ext import commands
-from discord_reporter import MultiProgressReporter
+from discord_reporter import MultiProgressReporterDiscord
 from discord_utils import (
     get_user_from_id,
     leaderboard_name_autocomplete,
@@ -85,79 +83,25 @@ async def on_submit_hook(  # noqa: C901
             )
             return -1
 
+        if not interaction.response.is_done():
+            await interaction.response.defer(ephemeral=True)
+
         req = SubmissionRequest(
             code=submission_content,
             file_name=script.filename,
             user_id=interaction.user.id,
+            user_name=interaction.user.global_name or interaction.user.name,
             gpus=cmd_gpus,
             leaderboard=leaderboard_name,
         )
         req = prepare_submission(req, self.bot.leaderboard_db)
 
-        # if there is more than one candidate GPU, display UI to let user select,
-        # otherwise just run on that GPU
-        if not interaction.response.is_done():
-            await interaction.response.defer(ephemeral=True)
-
         if req.gpus is None:
             view = await self.select_gpu_view(interaction, leaderboard_name, req.task_gpus)
-            selected_gpus = view.selected_gpus
-        else:
-            selected_gpus = req.gpus
-
-        selected_gpus = [get_gpu_by_name(gpu) for gpu in selected_gpus]
+            req.gpus = view.selected_gpus
 
-        command = self.bot.backend.submit_leaderboard
-
-        user_name = interaction.user.global_name or interaction.user.name
-        # Create a submission entry in the database
-        with self.bot.leaderboard_db as db:
-            sub_id = db.create_submission(
-                leaderboard=req.leaderboard,
-                file_name=script.filename,
-                code=submission_content,
-                user_id=interaction.user.id,
-                time=datetime.now(),
-                user_name=user_name,
-            )
-
-        run_msg = f"Submission **{sub_id}**: `{script.filename}` for `{req.leaderboard}`"
-        reporter = MultiProgressReporter(interaction, run_msg)
-        try:
-            tasks = [
-                command(
-                    sub_id,
-                    submission_content,
-                    script.filename,
-                    gpu,
-                    reporter.add_run(f"{gpu.name} on {gpu.runner}"),
-                    req.task,
-                    mode,
-                    None,
-                )
-                for gpu in selected_gpus
-            ]
-
-            # also schedule secret run
-            if mode == SubmissionMode.LEADERBOARD:
-                tasks += [
-                    command(
-                        sub_id,
-                        submission_content,
-                        script.filename,
-                        gpu,
-                        reporter.add_run(f"{gpu.name} on {gpu.runner} (secret)"),
-                        req.task,
-                        SubmissionMode.PRIVATE,
-                        req.secret_seed,
-                    )
-                    for gpu in selected_gpus
-                ]
-            await reporter.show()
-            await asyncio.gather(*tasks)
-        finally:
-            with self.bot.leaderboard_db as db:
-                db.mark_submission_done(sub_id)
+        reporter = MultiProgressReporterDiscord(interaction)
+        sub_id, results = await self.bot.backend.submit_full(req, mode, reporter)
 
         if mode == SubmissionMode.LEADERBOARD:
             await self.post_submit_hook(interaction, sub_id)

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

@@ -9,7 +9,7 @@
 from consts import SubmissionMode, get_gpu_by_name
 from discord import app_commands
 from discord.ext import commands
-from discord_reporter import MultiProgressReporter
+from discord_reporter import MultiProgressReporterDiscord
 from discord_utils import send_discord_message, with_error_handling
 from utils import (
     setup_logging,
@@ -70,7 +70,7 @@ async def run_script(
         """
         Function invoked by the `run` command to run a single script.
         """
-        reporter = MultiProgressReporter(interaction, "Script run")
+        reporter = MultiProgressReporterDiscord(interaction, "Script run")
         rep = reporter.add_run(f"{gpu_type.name}")
         await reporter.show()
         gpu_type = get_gpu_by_name(gpu_type.name)

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

@@ -13,7 +13,7 @@
 from discord import app_commands
 from discord.app_commands import Choice
 from discord.ext import commands
-from discord_reporter import MultiProgressReporter
+from discord_reporter import MultiProgressReporterDiscord
 from discord_utils import send_discord_message, with_error_handling
 from leaderboard_db import RunItem, SubmissionItem
 from task import make_task
@@ -291,7 +291,7 @@ async def verify_runs(self, interaction: discord.Interaction):
             amd = get_gpu_by_name("mi300")
             t4 = get_gpu_by_name("T4")
 
-            reporter = MultiProgressReporter("Verifying")
+            reporter = MultiProgressReporterDiscord("Verifying")
             await reporter.show(interaction)
 
             results = await asyncio.gather(

src/discord-cluster-manager/discord_reporter.py

@@ -1,18 +1,19 @@
 import discord
 from discord_utils import _send_split_log
-from report import Log, RunProgressReporter, RunResultReport, Text
+from report import Log, MultiProgressReporter, RunProgressReporter, RunResultReport, Text
 
 
-class MultiProgressReporter:
-    def __init__(self, interaction: discord.Interaction, header: str):
-        self.header = header
+class MultiProgressReporterDiscord(MultiProgressReporter):
+    def __init__(self, interaction: discord.Interaction):
+        self.header = ""
         self.runs = []
         self.interaction = interaction
 
-    async def show(self):
+    async def show(self, title: str):
+        self.header = title
         await self._update_message()
 
-    def add_run(self, title: str) -> "RunProgressReporter":
+    def add_run(self, title: str) -> "RunProgressReporterDiscord":
         rpr = RunProgressReporterDiscord(self, self.interaction, title)
         self.runs.append(rpr)
         return rpr
@@ -34,7 +35,7 @@ async def _update_message(self):
 class RunProgressReporterDiscord(RunProgressReporter):
     def __init__(
         self,
-        root: MultiProgressReporter,
+        root: MultiProgressReporterDiscord,
         interaction: discord.Interaction,
         title: str,
     ):