diff --git a/.github/workflows/amd_workflow.yml b/.github/workflows/amd_workflow.yml index c264c5ff..2ac899db 100644 --- a/.github/workflows/amd_workflow.yml +++ b/.github/workflows/amd_workflow.yml @@ -69,7 +69,7 @@ jobs: if [[ "${{ github.event.inputs.runner }}" == "amdgpu-mi250-x86-64" ]]; then source ${VENV_DIR}/bin/activate fi - python3 .github/workflows/runner.py + PYTHONPATH=src python3 src/runners/github-runner.py - name: Upload training artifacts uses: actions/upload-artifact@v4 diff --git a/.github/workflows/nvidia_workflow.yml b/.github/workflows/nvidia_workflow.yml index 17663ff9..ab4b527e 100644 --- a/.github/workflows/nvidia_workflow.yml +++ b/.github/workflows/nvidia_workflow.yml @@ -66,7 +66,7 @@ jobs: - name: Run script shell: bash run: | - python .github/workflows/runner.py + PYTHONPATH=src python src/runners/github-runner.py cat result.json # Debug: show output - name: Upload training artifacts diff --git a/.github/workflows/runner_ci.yml b/.github/workflows/runner_ci.yml index 594647aa..e2c4b47d 100644 --- a/.github/workflows/runner_ci.yml +++ b/.github/workflows/runner_ci.yml @@ -30,6 +30,8 @@ jobs: - name: Run script shell: bash + env: + PYTHONPATH: src run: pytest scripts/ci_test_cuda.py env: @@ -62,6 +64,8 @@ jobs: - name: Run script shell: bash + env: + PYTHONPATH: src run: pytest scripts/ci_test_python.py env: diff --git a/examples/identity_py/submission.py b/examples/identity_py/submission.py index 0ef8b529..1a7a7e24 100644 --- a/examples/identity_py/submission.py +++ b/examples/identity_py/submission.py @@ -1,4 +1,4 @@ -#!POPCORN leaderboard identity_py +#!POPCORN leaderboard identity_py-dev from task import input_t, output_t diff --git a/requirements.txt b/requirements.txt index c9b9796c..c2160978 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,6 +10,8 @@ yoyo-migrations ruff pre-commit better_profanity +pytest +PyYAML # api fastapi[all] # install all to avoid random bugs diff --git a/scripts/ci_test_cuda.py b/scripts/ci_test_cuda.py index c3661925..ae1a6cd7 100644 --- a/scripts/ci_test_cuda.py +++ b/scripts/ci_test_cuda.py @@ -1,17 +1,10 @@ -import os -import sys import tempfile from pathlib import Path import pytest -if Path().resolve().name == "scripts": - os.chdir("..") - -sys.path.append("src/discord-cluster-manager") - -from consts import ExitCode, SubmissionMode -from run_eval import compile_cuda_script, run_cuda_script +from libkernelbot.consts import ExitCode, SubmissionMode +from libkernelbot.run_eval import compile_cuda_script, run_cuda_script ref = Path("examples/identity_cuda/reference.cuh").read_text() task_h = Path("examples/identity_cuda/task.h").read_text() diff --git a/scripts/ci_test_python.py b/scripts/ci_test_python.py index fc43592b..41ac92bd 100644 --- a/scripts/ci_test_python.py +++ b/scripts/ci_test_python.py @@ -1,14 +1,7 @@ -import os -import sys from pathlib import Path -if Path().resolve().name == "scripts": - os.chdir("..") - -sys.path.append("src/discord-cluster-manager") - -from consts import ExitCode, SubmissionMode -from run_eval import run_pytorch_script +from libkernelbot.consts import ExitCode, SubmissionMode +from libkernelbot.run_eval import run_pytorch_script ref = Path("examples/identity_py/reference.py").read_text() task = Path("examples/identity_py/task.py").read_text() diff --git a/scripts/local-test.py b/scripts/local-test.py deleted file mode 100644 index ea42e772..00000000 --- a/scripts/local-test.py +++ /dev/null @@ -1,25 +0,0 @@ -import pprint -import sys -from pathlib import Path - -sys.path.append("src/discord-cluster-manager") - -from run_eval import run_cuda_script - -ref = Path("examples/identity_cuda/reference.cuh") -sub = Path("examples/identity_cuda/submission.cu") -util = Path("examples/identity_cuda/utils.h") -task = Path("examples/identity_cuda/task.h") - -result = run_cuda_script( - { - "eval.cu": Path("examples/identity_cuda/eval.cu").read_text(), - "submission.cu": sub.read_text(), - }, - {"reference.cuh": ref.read_text(), "utils.h": util.read_text(), "task.h": task.read_text()}, - arch=None, - tests="size: 128; seed: 45\nsize: 512; seed: 123", - mode="test", -) - -pprint.pprint(result) diff --git a/scripts/modal-test.py b/scripts/modal-test.py deleted file mode 100644 index 7695ae88..00000000 --- a/scripts/modal-test.py +++ /dev/null @@ -1,47 +0,0 @@ -# Run modal token new -# Then run this script as a simple test - -# Script executed successfully -# Stopping app - local entrypoint completed. -# tensor([1., 2., 3., 4., 5.], device='cuda:0') -# tensor([1., 2., 3., 4., 5.], device='cuda:0') -# tensor([ 2., 4., 6., 8., 10.], device='cuda:0') - -import modal - -# Modal app setup -modal_app = modal.App("discord-bot-runner") - - -@modal_app.function( - gpu="T4", image=modal.Image.debian_slim(python_version="3.12").pip_install(["torch"]) -) -async def run_pytorch_script_on_modal(): - """ - Runs a Python script on Modal with GPU - """ - try: - # Define main script content - main_script = """ -import sys -import torch - -# Your actual script content -a = torch.Tensor([1, 2, 3, 4, 5]).cuda() -b = torch.Tensor([1, 2, 3, 4, 5]).cuda() -print(a) -print(b) -print(a + b) -""" - # Execute the script content directly - exec(main_script, {"__name__": "__main__"}) - return "Script executed successfully" - except Exception as e: - return f"Error executing script: {str(e)}" - - -# Run the function -if __name__ == "__main__": - with modal_app.run(): - result = run_pytorch_script_on_modal.remote() - print(result) diff --git a/src/discord-cluster-manager/env.py b/src/discord-cluster-manager/env.py deleted file mode 100644 index 99e2b18e..00000000 --- a/src/discord-cluster-manager/env.py +++ /dev/null @@ -1,50 +0,0 @@ -import os - -from dotenv import load_dotenv -from utils import get_github_branch_name - - -def init_environment(): - load_dotenv() - - # Validate environment - required_env_vars = ["DISCORD_TOKEN", "GITHUB_TOKEN", "GITHUB_REPO"] - for var in required_env_vars: - if not os.getenv(var): - raise ValueError(f"{var} not found") - - -init_environment() - -# Discord-specific constants -DISCORD_TOKEN = os.getenv("DISCORD_TOKEN") -DISCORD_DEBUG_TOKEN = os.getenv("DISCORD_DEBUG_TOKEN") -DISCORD_CLUSTER_STAGING_ID = os.getenv("DISCORD_CLUSTER_STAGING_ID") -DISCORD_DEBUG_CLUSTER_STAGING_ID = os.getenv("DISCORD_DEBUG_CLUSTER_STAGING_ID") - -# Only required to run the CLI against this instance -# setting these is required only to run the CLI against local instance -CLI_DISCORD_CLIENT_ID = os.getenv("CLI_DISCORD_CLIENT_ID", "") -CLI_DISCORD_CLIENT_SECRET = os.getenv("CLI_DISCORD_CLIENT_SECRET", "") -CLI_TOKEN_URL = os.getenv("CLI_TOKEN_URL", "") -CLI_GITHUB_CLIENT_ID = os.getenv("CLI_GITHUB_CLIENT_ID", "") -CLI_GITHUB_CLIENT_SECRET = os.getenv("CLI_GITHUB_CLIENT_SECRET", "") - - -# GitHub-specific constants -GITHUB_TOKEN = os.getenv("GITHUB_TOKEN") -GITHUB_REPO = os.getenv("GITHUB_REPO") -GITHUB_WORKFLOW_BRANCH = os.getenv("GITHUB_WORKFLOW_BRANCH", get_github_branch_name()) -PROBLEMS_REPO = os.getenv("PROBLEMS_REPO") - -# Directory that will be used for local problem development. -PROBLEM_DEV_DIR = os.getenv("PROBLEM_DEV_DIR", "examples") - -# PostgreSQL-specific constants -POSTGRES_HOST = os.getenv("POSTGRES_HOST") -POSTGRES_DATABASE = os.getenv("POSTGRES_DATABASE") -POSTGRES_USER = os.getenv("POSTGRES_USER") -POSTGRES_PASSWORD = os.getenv("POSTGRES_PASSWORD") -POSTGRES_PORT = os.getenv("POSTGRES_PORT") -DATABASE_URL = os.getenv("DATABASE_URL") -DISABLE_SSL = os.getenv("DISABLE_SSL") diff --git a/src/kernelbot/api/__init__.py b/src/kernelbot/api/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/discord-cluster-manager/api/api_utils.py b/src/kernelbot/api/api_utils.py similarity index 91% rename from src/discord-cluster-manager/api/api_utils.py rename to src/kernelbot/api/api_utils.py index 26c269ef..b5cda480 100644 --- a/src/discord-cluster-manager/api/api_utils.py +++ b/src/kernelbot/api/api_utils.py @@ -1,23 +1,24 @@ import requests -from backend import KernelBackend -from consts import SubmissionMode -from env import ( - CLI_DISCORD_CLIENT_ID, - CLI_DISCORD_CLIENT_SECRET, - CLI_GITHUB_CLIENT_ID, - CLI_GITHUB_CLIENT_SECRET, - CLI_TOKEN_URL, -) from fastapi import HTTPException -from report import Log, MultiProgressReporter, RunProgressReporter, RunResultReport, Text -from submission import SubmissionRequest, prepare_submission + +from kernelbot.env import env +from libkernelbot.backend import KernelBackend +from libkernelbot.consts import SubmissionMode +from libkernelbot.report import ( + Log, + MultiProgressReporter, + RunProgressReporter, + RunResultReport, + Text, +) +from libkernelbot.submission import SubmissionRequest, prepare_submission async def _handle_discord_oauth(code: str, redirect_uri: str) -> tuple[str, str]: """Handles the Discord OAuth code exchange and user info retrieval.""" - client_id = CLI_DISCORD_CLIENT_ID - client_secret = CLI_DISCORD_CLIENT_SECRET - token_url = CLI_TOKEN_URL + client_id = env.CLI_DISCORD_CLIENT_ID + client_secret = env.CLI_DISCORD_CLIENT_SECRET + token_url = env.CLI_TOKEN_URL user_api_url = "https://discord.com/api/users/@me" if not client_id: @@ -76,8 +77,8 @@ async def _handle_discord_oauth(code: str, redirect_uri: str) -> tuple[str, str] async def _handle_github_oauth(code: str, redirect_uri: str) -> tuple[str, str]: """Handles the GitHub OAuth code exchange and user info retrieval.""" - client_id = CLI_GITHUB_CLIENT_ID - client_secret = CLI_GITHUB_CLIENT_SECRET + client_id = env.CLI_GITHUB_CLIENT_ID + client_secret = env.CLI_GITHUB_CLIENT_SECRET token_url = "https://github.com/login/oauth/access_token" user_api_url = "https://api.github.com/user" diff --git a/src/discord-cluster-manager/api/main.py b/src/kernelbot/api/main.py similarity index 97% rename from src/discord-cluster-manager/api/main.py rename to src/kernelbot/api/main.py index 6120c7a5..2848d383 100644 --- a/src/discord-cluster-manager/api/main.py +++ b/src/kernelbot/api/main.py @@ -7,13 +7,14 @@ from dataclasses import asdict from typing import Annotated, Optional -from backend import KernelBackend -from consts import SubmissionMode from fastapi import Depends, FastAPI, Header, HTTPException, Request, UploadFile from fastapi.responses import JSONResponse, StreamingResponse -from leaderboard_db import LeaderboardRankedEntry -from submission import SubmissionRequest -from utils import KernelBotError + +from libkernelbot.backend import KernelBackend +from libkernelbot.consts import SubmissionMode +from libkernelbot.leaderboard_db import LeaderboardRankedEntry +from libkernelbot.submission import SubmissionRequest +from libkernelbot.utils import KernelBotError from .api_utils import _handle_discord_oauth, _handle_github_oauth, _run_submission @@ -309,7 +310,7 @@ async def run_submission( # noqa: C901 user_id (str): The validated user ID obtained from the X-Popcorn-Cli-Id header. Raises: - HTTPException: If the bot is not initialized, or header/input is invalid. + HTTPException: If the kernelbot is not initialized, or header/input is invalid. Returns: StreamingResponse: A streaming response containing the status and results of the submission. diff --git a/src/kernelbot/cogs/__init__.py b/src/kernelbot/cogs/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/discord-cluster-manager/cogs/admin_cog.py b/src/kernelbot/cogs/admin_cog.py similarity index 97% rename from src/discord-cluster-manager/cogs/admin_cog.py rename to src/kernelbot/cogs/admin_cog.py index 03b3420c..c33ad847 100644 --- a/src/discord-cluster-manager/cogs/admin_cog.py +++ b/src/kernelbot/cogs/admin_cog.py @@ -8,22 +8,27 @@ from typing import TYPE_CHECKING, Optional, TypedDict import discord -import env import yaml -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 LeaderboardDoesNotExist, LeaderboardItem, SubmissionItem -from task import LeaderboardDefinition, make_task_definition -from ui.misc import ConfirmationView, DeleteConfirmationModal, GPUSelectionView -from utils import ( + +from kernelbot.discord_utils import ( + leaderboard_name_autocomplete, + send_discord_message, + with_error_handling, +) +from kernelbot.env import env +from kernelbot.ui.misc import ConfirmationView, DeleteConfirmationModal, GPUSelectionView +from libkernelbot.consts import GitHubGPU, ModalGPU +from libkernelbot.leaderboard_db import LeaderboardDoesNotExist, LeaderboardItem, SubmissionItem +from libkernelbot.task import LeaderboardDefinition, make_task_definition +from libkernelbot.utils import ( KernelBotError, setup_logging, ) if TYPE_CHECKING: - from ..bot import ClusterBot + from kernelbot.main import ClusterBot logger = setup_logging() @@ -66,7 +71,7 @@ class AdminCog(commands.Cog): def __init__(self, bot: "ClusterBot"): self.bot = bot - # create-local should only be used for the development bot + # create-local should only be used for the development kernelbot if self.bot.debug_mode: self.leaderboard_create_local = bot.admin_group.command( name="create-local", @@ -82,11 +87,11 @@ def __init__(self, bot: "ClusterBot"): )(self.delete_submission) self.accept_jobs = bot.admin_group.command( - name="start", description="Make the bot accept new submissions" + name="start", description="Make the kernelbot accept new submissions" )(self.start) self.reject_jobs = bot.admin_group.command( - name="stop", description="Make the bot stop accepting new submissions" + name="stop", description="Make the kernelbot stop accepting new submissions" )(self.stop) self.update_problems = bot.admin_group.command( @@ -94,7 +99,7 @@ def __init__(self, bot: "ClusterBot"): )(self.update_problems) self.show_bot_stats = bot.admin_group.command( - name="show-stats", description="Show stats for the bot" + name="show-stats", description="Show stats for the kernelbot" )(self.show_bot_stats) self.resync = bot.admin_group.command( diff --git a/src/discord-cluster-manager/cogs/leaderboard_cog.py b/src/kernelbot/cogs/leaderboard_cog.py similarity index 98% rename from src/discord-cluster-manager/cogs/leaderboard_cog.py rename to src/kernelbot/cogs/leaderboard_cog.py index 59fe2f0e..6cb8bb40 100644 --- a/src/discord-cluster-manager/cogs/leaderboard_cog.py +++ b/src/kernelbot/cogs/leaderboard_cog.py @@ -3,34 +3,30 @@ from typing import TYPE_CHECKING, List, Optional import discord -from consts import ( - SubmissionMode, -) from discord import app_commands from discord.ext import commands -from discord_reporter import MultiProgressReporterDiscord -from discord_utils import ( + +from kernelbot.discord_reporter import MultiProgressReporterDiscord +from kernelbot.discord_utils import ( get_user_from_id, leaderboard_name_autocomplete, send_discord_message, with_error_handling, ) -from leaderboard_db import ( +from kernelbot.ui.misc import GPUSelectionView +from kernelbot.ui.table import create_table +from libkernelbot.consts import SubmissionMode +from libkernelbot.leaderboard_db import ( LeaderboardItem, LeaderboardRankedEntry, RunItem, SubmissionItem, ) -from submission import SubmissionRequest, prepare_submission -from ui.misc import GPUSelectionView -from ui.table import create_table -from utils import ( - format_time, - setup_logging, -) +from libkernelbot.submission import SubmissionRequest, prepare_submission +from libkernelbot.utils import format_time, setup_logging if TYPE_CHECKING: - from ..bot import ClusterBot + from kernelbot.main import ClusterBot logger = setup_logging() diff --git a/src/discord-cluster-manager/cogs/misc_cog.py b/src/kernelbot/cogs/misc_cog.py similarity index 86% rename from src/discord-cluster-manager/cogs/misc_cog.py rename to src/kernelbot/cogs/misc_cog.py index 628cb557..09e78bc8 100644 --- a/src/discord-cluster-manager/cogs/misc_cog.py +++ b/src/kernelbot/cogs/misc_cog.py @@ -5,12 +5,13 @@ 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 setup_logging + +from kernelbot.discord_utils import send_discord_message +from kernelbot.env import env +from libkernelbot.utils import setup_logging if TYPE_CHECKING: - from ..bot import ClusterBot + from ..main import ClusterBot logger = setup_logging() @@ -27,15 +28,15 @@ async def ping(self, interaction: discord.Interaction): @app_commands.command(name="verifydb") async def verify_db(self, interaction: discord.Interaction): """Command to verify database connectivity""" - if not DATABASE_URL: + if not env.DATABASE_URL: message = "DATABASE_URL not set." logger.error(message) await send_discord_message(interaction, message) return try: - sslmode = "disable" if DISABLE_SSL else "require" - with psycopg2.connect(DATABASE_URL, sslmode=sslmode) as conn: + sslmode = "disable" if env.DISABLE_SSL else "require" + with psycopg2.connect(env.DATABASE_URL, sslmode=sslmode) as conn: with conn.cursor() as cursor: cursor.execute("SELECT RANDOM()") result = cursor.fetchone() diff --git a/src/discord-cluster-manager/cogs/verify_run_cog.py b/src/kernelbot/cogs/verify_run_cog.py similarity index 95% rename from src/discord-cluster-manager/cogs/verify_run_cog.py rename to src/kernelbot/cogs/verify_run_cog.py index 70fe5dfc..1b23d16d 100644 --- a/src/discord-cluster-manager/cogs/verify_run_cog.py +++ b/src/kernelbot/cogs/verify_run_cog.py @@ -6,18 +6,19 @@ from unittest.mock import AsyncMock import discord -import env -from cogs import admin_cog -from cogs.leaderboard_cog import LeaderboardSubmitCog -from consts import GPU, SubmissionMode, get_gpu_by_name from discord import app_commands from discord.app_commands import Choice from discord.ext import commands -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_definition -from utils import setup_logging + +from kernelbot.cogs import admin_cog +from kernelbot.cogs.leaderboard_cog import LeaderboardSubmitCog +from kernelbot.discord_reporter import MultiProgressReporterDiscord +from kernelbot.discord_utils import send_discord_message, with_error_handling +from kernelbot.env import env +from libkernelbot.consts import GPU, SubmissionMode, get_gpu_by_name +from libkernelbot.leaderboard_db import RunItem, SubmissionItem +from libkernelbot.task import make_task_definition +from libkernelbot.utils import setup_logging logger = setup_logging() diff --git a/src/discord-cluster-manager/discord_reporter.py b/src/kernelbot/discord_reporter.py similarity index 94% rename from src/discord-cluster-manager/discord_reporter.py rename to src/kernelbot/discord_reporter.py index ca71a606..f49e5225 100644 --- a/src/discord-cluster-manager/discord_reporter.py +++ b/src/kernelbot/discord_reporter.py @@ -1,6 +1,13 @@ import discord from discord_utils import _send_split_log -from report import Log, MultiProgressReporter, RunProgressReporter, RunResultReport, Text + +from libkernelbot.report import ( + Log, + MultiProgressReporter, + RunProgressReporter, + RunResultReport, + Text, +) class MultiProgressReporterDiscord(MultiProgressReporter): diff --git a/src/discord-cluster-manager/discord_utils.py b/src/kernelbot/discord_utils.py similarity index 98% rename from src/discord-cluster-manager/discord_utils.py rename to src/kernelbot/discord_utils.py index 7cad810d..d014f3ca 100644 --- a/src/discord-cluster-manager/discord_utils.py +++ b/src/kernelbot/discord_utils.py @@ -2,7 +2,8 @@ import logging import discord -from utils import KernelBotError, limit_length, setup_logging + +from libkernelbot.utils import KernelBotError, limit_length, setup_logging logger = setup_logging(__name__) diff --git a/src/kernelbot/env.py b/src/kernelbot/env.py new file mode 100644 index 00000000..b2e457da --- /dev/null +++ b/src/kernelbot/env.py @@ -0,0 +1,54 @@ +import os +import types + +from dotenv import load_dotenv + +from libkernelbot.utils import get_github_branch_name + + +def init_environment(): + load_dotenv() + + # Validate environment + required_env_vars = ["DISCORD_TOKEN", "GITHUB_TOKEN", "GITHUB_REPO"] + for var in required_env_vars: + if not os.getenv(var): + raise ValueError(f"{var} not found") + + +init_environment() + +env = types.SimpleNamespace() + +# Discord-specific constants +env.DISCORD_TOKEN = os.getenv("DISCORD_TOKEN") +env.DISCORD_DEBUG_TOKEN = os.getenv("DISCORD_DEBUG_TOKEN") +env.DISCORD_CLUSTER_STAGING_ID = os.getenv("DISCORD_CLUSTER_STAGING_ID") +env.DISCORD_DEBUG_CLUSTER_STAGING_ID = os.getenv("DISCORD_DEBUG_CLUSTER_STAGING_ID") + +# Only required to run the CLI against this instance +# setting these is required only to run the CLI against local instance +env.CLI_DISCORD_CLIENT_ID = os.getenv("CLI_DISCORD_CLIENT_ID", "") +env.CLI_DISCORD_CLIENT_SECRET = os.getenv("CLI_DISCORD_CLIENT_SECRET", "") +env.CLI_TOKEN_URL = os.getenv("CLI_TOKEN_URL", "") +env.CLI_GITHUB_CLIENT_ID = os.getenv("CLI_GITHUB_CLIENT_ID", "") +env.CLI_GITHUB_CLIENT_SECRET = os.getenv("CLI_GITHUB_CLIENT_SECRET", "") + + +# GitHub-specific constants +env.GITHUB_TOKEN = os.getenv("GITHUB_TOKEN") +env.GITHUB_REPO = os.getenv("GITHUB_REPO") +env.GITHUB_WORKFLOW_BRANCH = os.getenv("GITHUB_WORKFLOW_BRANCH", get_github_branch_name()) +env.PROBLEMS_REPO = os.getenv("PROBLEMS_REPO") + +# Directory that will be used for local problem development. +env.PROBLEM_DEV_DIR = os.getenv("PROBLEM_DEV_DIR", "examples") + +# PostgreSQL-specific constants +env.POSTGRES_HOST = os.getenv("POSTGRES_HOST") +env.POSTGRES_DATABASE = os.getenv("POSTGRES_DATABASE") +env.POSTGRES_USER = os.getenv("POSTGRES_USER") +env.POSTGRES_PASSWORD = os.getenv("POSTGRES_PASSWORD") +env.POSTGRES_PORT = os.getenv("POSTGRES_PORT") +env.DATABASE_URL = os.getenv("DATABASE_URL") +env.DISABLE_SSL = os.getenv("DISABLE_SSL") diff --git a/src/discord-cluster-manager/bot.py b/src/kernelbot/main.py similarity index 93% rename from src/discord-cluster-manager/bot.py rename to src/kernelbot/main.py index 6d029381..20e40248 100644 --- a/src/discord-cluster-manager/bot.py +++ b/src/kernelbot/main.py @@ -2,27 +2,21 @@ import asyncio import os -import consts import discord -import env import uvicorn from api.main import app, init_api -from backend import KernelBackend from cogs.admin_cog import AdminCog from cogs.leaderboard_cog import LeaderboardCog from cogs.misc_cog import BotManagerCog from cogs.verify_run_cog import VerifyRunCog from discord import app_commands from discord.ext import commands -from env import ( - DISCORD_CLUSTER_STAGING_ID, - DISCORD_DEBUG_CLUSTER_STAGING_ID, - DISCORD_DEBUG_TOKEN, - DISCORD_TOKEN, - init_environment, -) -from launchers import GitHubLauncher, ModalLauncher -from utils import setup_logging +from env import env, init_environment + +from libkernelbot import consts +from libkernelbot.backend import KernelBackend +from libkernelbot.launchers import GitHubLauncher, ModalLauncher +from libkernelbot.utils import setup_logging logger = setup_logging(__name__) @@ -43,7 +37,7 @@ def __init__(self, debug_mode=False): self.tree.add_command(self.admin_group) self.tree.add_command(self.leaderboard_group) - self.backend = KernelBackend(debug_mode=debug_mode) + self.backend = KernelBackend(env=env, debug_mode=debug_mode) self.backend.register_launcher(ModalLauncher(consts.MODAL_CUDA_INCLUDE_DIRS)) self.backend.register_launcher( GitHubLauncher(env.GITHUB_REPO, env.GITHUB_TOKEN, env.GITHUB_WORKFLOW_BRANCH) @@ -54,7 +48,7 @@ def leaderboard_db(self): return self.backend.db async def setup_hook(self): - logger.info(f"Syncing commands for staging guild {DISCORD_CLUSTER_STAGING_ID}") + logger.info(f"Syncing commands for staging guild {env.DISCORD_CLUSTER_STAGING_ID}") try: # Load cogs await self.add_cog(BotManagerCog(self)) @@ -63,9 +57,9 @@ async def setup_hook(self): await self.add_cog(AdminCog(self)) guild_id = ( - DISCORD_CLUSTER_STAGING_ID + env.DISCORD_CLUSTER_STAGING_ID if not self.debug_mode - else DISCORD_DEBUG_CLUSTER_STAGING_ID + else env.DISCORD_DEBUG_CLUSTER_STAGING_ID ) if guild_id: @@ -214,7 +208,7 @@ async def start_bot(self, token: str): async def start_bot_and_api(debug_mode: bool): - token = DISCORD_DEBUG_TOKEN if debug_mode else DISCORD_TOKEN + token = env.DISCORD_DEBUG_TOKEN if debug_mode else env.DISCORD_TOKEN if debug_mode and not token: raise ValueError("DISCORD_DEBUG_TOKEN not found") @@ -247,7 +241,7 @@ def main(): parser.add_argument("--debug", action="store_true", help="Run in debug/staging mode") args = parser.parse_args() - logger.info("Starting bot and API server...") + logger.info("Starting kernelbot and API server...") with asyncio.Runner(debug=args.debug) as runner: runner.get_loop().set_exception_handler(on_unhandled_exception) diff --git a/src/discord-cluster-manager/ui/misc.py b/src/kernelbot/ui/misc.py similarity index 98% rename from src/discord-cluster-manager/ui/misc.py rename to src/kernelbot/ui/misc.py index e09d0f5a..14161149 100644 --- a/src/discord-cluster-manager/ui/misc.py +++ b/src/kernelbot/ui/misc.py @@ -3,7 +3,8 @@ import discord from discord import Interaction, SelectOption, ui from discord_utils import send_discord_message -from utils import KernelBotError + +from libkernelbot.utils import KernelBotError class GPUSelectionView(ui.View): diff --git a/src/discord-cluster-manager/ui/table.py b/src/kernelbot/ui/table.py similarity index 100% rename from src/discord-cluster-manager/ui/table.py rename to src/kernelbot/ui/table.py diff --git a/src/libkernelbot/__init__.py b/src/libkernelbot/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/discord-cluster-manager/backend.py b/src/libkernelbot/backend.py similarity index 93% rename from src/discord-cluster-manager/backend.py rename to src/libkernelbot/backend.py index 086aefdc..fad3bef1 100644 --- a/src/discord-cluster-manager/backend.py +++ b/src/libkernelbot/backend.py @@ -2,17 +2,22 @@ import copy import math from datetime import datetime +from types import SimpleNamespace from typing import Optional -import env -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 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 +from libkernelbot.consts import GPU, GPU_TO_SM, RankCriterion, SubmissionMode, get_gpu_by_name +from libkernelbot.launchers import Launcher +from libkernelbot.leaderboard_db import LeaderboardDB +from libkernelbot.report import ( + MultiProgressReporter, + RunProgressReporter, + generate_report, + make_short_report, +) +from libkernelbot.run_eval import FullResult +from libkernelbot.submission import ProcessedSubmissionRequest +from libkernelbot.task import LeaderboardTask, build_task_config +from libkernelbot.utils import KernelBotError, setup_logging logger = setup_logging(__name__) @@ -20,6 +25,7 @@ class KernelBackend: def __init__( self, + env: SimpleNamespace, debug_mode=False, ): self.debug_mode = debug_mode diff --git a/src/discord-cluster-manager/consts.py b/src/libkernelbot/consts.py similarity index 100% rename from src/discord-cluster-manager/consts.py rename to src/libkernelbot/consts.py diff --git a/src/discord-cluster-manager/launchers/__init__.py b/src/libkernelbot/launchers/__init__.py similarity index 100% rename from src/discord-cluster-manager/launchers/__init__.py rename to src/libkernelbot/launchers/__init__.py diff --git a/src/discord-cluster-manager/launchers/github.py b/src/libkernelbot/launchers/github.py similarity index 98% rename from src/discord-cluster-manager/launchers/github.py rename to src/libkernelbot/launchers/github.py index d71610a8..304b7e08 100644 --- a/src/discord-cluster-manager/launchers/github.py +++ b/src/libkernelbot/launchers/github.py @@ -11,7 +11,9 @@ from typing import Awaitable, Callable, Optional import requests -from consts import ( +from github import Github, UnknownObjectException, WorkflowRun + +from libkernelbot.consts import ( AMD_REQUIREMENTS, DEFAULT_GITHUB_TIMEOUT_MINUTES, GPU, @@ -20,10 +22,9 @@ GitHubGPU, SubmissionMode, ) -from github import Github, UnknownObjectException, WorkflowRun -from report import RunProgressReporter -from run_eval import CompileResult, EvalResult, FullResult, RunResult, SystemInfo -from utils import setup_logging +from libkernelbot.report import RunProgressReporter +from libkernelbot.run_eval import CompileResult, EvalResult, FullResult, RunResult, SystemInfo +from libkernelbot.utils import setup_logging from .launcher import Launcher diff --git a/src/discord-cluster-manager/launchers/launcher.py b/src/libkernelbot/launchers/launcher.py similarity index 77% rename from src/discord-cluster-manager/launchers/launcher.py rename to src/libkernelbot/launchers/launcher.py index 5760c18c..67292bdc 100644 --- a/src/discord-cluster-manager/launchers/launcher.py +++ b/src/libkernelbot/launchers/launcher.py @@ -1,8 +1,8 @@ from enum import Enum from typing import Type -from consts import GPU -from report import RunProgressReporter +from libkernelbot.consts import GPU +from libkernelbot.report import RunProgressReporter class Launcher: diff --git a/src/discord-cluster-manager/launchers/modal.py b/src/libkernelbot/launchers/modal.py similarity index 85% rename from src/discord-cluster-manager/launchers/modal.py rename to src/libkernelbot/launchers/modal.py index dc22fe18..26adfe33 100644 --- a/src/discord-cluster-manager/launchers/modal.py +++ b/src/libkernelbot/launchers/modal.py @@ -1,10 +1,11 @@ import asyncio import modal -from consts import GPU, ModalGPU -from report import RunProgressReporter -from run_eval import FullResult -from utils import setup_logging + +from libkernelbot.consts import GPU, ModalGPU +from libkernelbot.report import RunProgressReporter +from libkernelbot.run_eval import FullResult +from libkernelbot.utils import setup_logging from .launcher import Launcher diff --git a/src/discord-cluster-manager/leaderboard_db.py b/src/libkernelbot/leaderboard_db.py similarity index 99% rename from src/discord-cluster-manager/leaderboard_db.py rename to src/libkernelbot/leaderboard_db.py index e694110a..fed6125c 100644 --- a/src/discord-cluster-manager/leaderboard_db.py +++ b/src/libkernelbot/leaderboard_db.py @@ -4,9 +4,10 @@ from typing import Dict, List, NotRequired, Optional, TypedDict import psycopg2 -from run_eval import CompileResult, RunResult, SystemInfo -from task import LeaderboardDefinition, LeaderboardTask -from utils import ( + +from libkernelbot.run_eval import CompileResult, RunResult, SystemInfo +from libkernelbot.task import LeaderboardDefinition, LeaderboardTask +from libkernelbot.utils import ( KernelBotError, LRUCache, setup_logging, diff --git a/src/discord-cluster-manager/report.py b/src/libkernelbot/report.py similarity index 98% rename from src/discord-cluster-manager/report.py rename to src/libkernelbot/report.py index bdfb6a8f..e6c8bad3 100644 --- a/src/discord-cluster-manager/report.py +++ b/src/libkernelbot/report.py @@ -3,9 +3,9 @@ import textwrap from typing import List -import consts -from run_eval import CompileResult, EvalResult, FullResult, RunResult, SystemInfo -from utils import format_time, limit_length +from libkernelbot import consts +from libkernelbot.run_eval import CompileResult, EvalResult, FullResult, RunResult, SystemInfo +from libkernelbot.utils import format_time, limit_length @dataclasses.dataclass diff --git a/src/discord-cluster-manager/run_eval.py b/src/libkernelbot/run_eval.py similarity index 99% rename from src/discord-cluster-manager/run_eval.py rename to src/libkernelbot/run_eval.py index 37eeec26..18344dd3 100644 --- a/src/discord-cluster-manager/run_eval.py +++ b/src/libkernelbot/run_eval.py @@ -10,7 +10,7 @@ from types import NoneType from typing import Optional, Protocol, Union -from consts import CUDA_FLAGS, ExitCode, Timeout +from libkernelbot.consts import CUDA_FLAGS, ExitCode, Timeout @dataclasses.dataclass diff --git a/src/discord-cluster-manager/submission.py b/src/libkernelbot/submission.py similarity index 96% rename from src/discord-cluster-manager/submission.py rename to src/libkernelbot/submission.py index 4a53d124..8bc31ef5 100644 --- a/src/discord-cluster-manager/submission.py +++ b/src/libkernelbot/submission.py @@ -5,9 +5,10 @@ from typing import Optional, Union from better_profanity import profanity -from leaderboard_db import LeaderboardDB, LeaderboardItem -from task import LeaderboardTask -from utils import KernelBotError + +from libkernelbot.leaderboard_db import LeaderboardDB, LeaderboardItem +from libkernelbot.task import LeaderboardTask +from libkernelbot.utils import KernelBotError if typing.TYPE_CHECKING: from backend import KernelBackend diff --git a/src/discord-cluster-manager/task.py b/src/libkernelbot/task.py similarity index 97% rename from src/discord-cluster-manager/task.py rename to src/libkernelbot/task.py index 0fea0e8a..7fe1b698 100644 --- a/src/discord-cluster-manager/task.py +++ b/src/libkernelbot/task.py @@ -5,8 +5,10 @@ from pathlib import Path from typing import Dict, Optional, Union -from consts import Language, RankCriterion, SubmissionMode -from utils import KernelBotError +import yaml + +from libkernelbot.consts import Language, RankCriterion, SubmissionMode +from libkernelbot.utils import KernelBotError @dataclasses.dataclass @@ -105,8 +107,6 @@ class LeaderboardDefinition: def make_task_definition(yaml_file: str | Path) -> LeaderboardDefinition: - import yaml - if Path(yaml_file).is_dir(): yaml_file = Path(yaml_file) / "task.yml" diff --git a/src/discord-cluster-manager/utils.py b/src/libkernelbot/utils.py similarity index 100% rename from src/discord-cluster-manager/utils.py rename to src/libkernelbot/utils.py diff --git a/src/discord-cluster-manager/migrations/20241208_01_p3yuR-initial-leaderboard-schema.py b/src/migrations/20241208_01_p3yuR-initial-leaderboard-schema.py similarity index 100% rename from src/discord-cluster-manager/migrations/20241208_01_p3yuR-initial-leaderboard-schema.py rename to src/migrations/20241208_01_p3yuR-initial-leaderboard-schema.py diff --git a/src/discord-cluster-manager/migrations/20241214_01_M62BX-drop-old-leaderboard-tables.py b/src/migrations/20241214_01_M62BX-drop-old-leaderboard-tables.py similarity index 100% rename from src/discord-cluster-manager/migrations/20241214_01_M62BX-drop-old-leaderboard-tables.py rename to src/migrations/20241214_01_M62BX-drop-old-leaderboard-tables.py diff --git a/src/discord-cluster-manager/migrations/20241221_01_54Oeg-rename-problem-table.py b/src/migrations/20241221_01_54Oeg-rename-problem-table.py similarity index 100% rename from src/discord-cluster-manager/migrations/20241221_01_54Oeg-rename-problem-table.py rename to src/migrations/20241221_01_54Oeg-rename-problem-table.py diff --git a/src/discord-cluster-manager/migrations/20241222_01_ELxU5-add-gpu-types.py b/src/migrations/20241222_01_ELxU5-add-gpu-types.py similarity index 100% rename from src/discord-cluster-manager/migrations/20241222_01_ELxU5-add-gpu-types.py rename to src/migrations/20241222_01_ELxU5-add-gpu-types.py diff --git a/src/discord-cluster-manager/migrations/20241224_01_Pg4FX-delete-cascade.py b/src/migrations/20241224_01_Pg4FX-delete-cascade.py similarity index 100% rename from src/discord-cluster-manager/migrations/20241224_01_Pg4FX-delete-cascade.py rename to src/migrations/20241224_01_Pg4FX-delete-cascade.py diff --git a/src/discord-cluster-manager/migrations/20241226_01_ZQSOK-add_gpu_type_to_submission.py b/src/migrations/20241226_01_ZQSOK-add_gpu_type_to_submission.py similarity index 100% rename from src/discord-cluster-manager/migrations/20241226_01_ZQSOK-add_gpu_type_to_submission.py rename to src/migrations/20241226_01_ZQSOK-add_gpu_type_to_submission.py diff --git a/src/discord-cluster-manager/migrations/20250106_01_Sgph3-add-leaderboard-creator-id.py b/src/migrations/20250106_01_Sgph3-add-leaderboard-creator-id.py similarity index 100% rename from src/discord-cluster-manager/migrations/20250106_01_Sgph3-add-leaderboard-creator-id.py rename to src/migrations/20250106_01_Sgph3-add-leaderboard-creator-id.py diff --git a/src/discord-cluster-manager/migrations/20250202_01_YYS3Q-leaderboard-rename-reference-to-task.py b/src/migrations/20250202_01_YYS3Q-leaderboard-rename-reference-to-task.py similarity index 100% rename from src/discord-cluster-manager/migrations/20250202_01_YYS3Q-leaderboard-rename-reference-to-task.py rename to src/migrations/20250202_01_YYS3Q-leaderboard-rename-reference-to-task.py diff --git a/src/discord-cluster-manager/migrations/20250221_01_GA8ro-submission-collection.py b/src/migrations/20250221_01_GA8ro-submission-collection.py similarity index 100% rename from src/discord-cluster-manager/migrations/20250221_01_GA8ro-submission-collection.py rename to src/migrations/20250221_01_GA8ro-submission-collection.py diff --git a/src/discord-cluster-manager/migrations/20250228_01_9ANYn-submission-add-user-name.py b/src/migrations/20250228_01_9ANYn-submission-add-user-name.py similarity index 100% rename from src/discord-cluster-manager/migrations/20250228_01_9ANYn-submission-add-user-name.py rename to src/migrations/20250228_01_9ANYn-submission-add-user-name.py diff --git a/src/discord-cluster-manager/migrations/20250304_01_DzORz-collect-system-information-for-each-run.py b/src/migrations/20250304_01_DzORz-collect-system-information-for-each-run.py similarity index 100% rename from src/discord-cluster-manager/migrations/20250304_01_DzORz-collect-system-information-for-each-run.py rename to src/migrations/20250304_01_DzORz-collect-system-information-for-each-run.py diff --git a/src/discord-cluster-manager/migrations/20250316_01_5oMi3-remember-forum-id.py b/src/migrations/20250316_01_5oMi3-remember-forum-id.py similarity index 100% rename from src/discord-cluster-manager/migrations/20250316_01_5oMi3-remember-forum-id.py rename to src/migrations/20250316_01_5oMi3-remember-forum-id.py diff --git a/src/discord-cluster-manager/migrations/20250329_01_7VjJJ-add-a-secret-seed-column.py b/src/migrations/20250329_01_7VjJJ-add-a-secret-seed-column.py similarity index 100% rename from src/discord-cluster-manager/migrations/20250329_01_7VjJJ-add-a-secret-seed-column.py rename to src/migrations/20250329_01_7VjJJ-add-a-secret-seed-column.py diff --git a/src/discord-cluster-manager/migrations/20250406_01_ZXjWK-user-info-add-cli-id.py b/src/migrations/20250406_01_ZXjWK-user-info-add-cli-id.py similarity index 100% rename from src/discord-cluster-manager/migrations/20250406_01_ZXjWK-user-info-add-cli-id.py rename to src/migrations/20250406_01_ZXjWK-user-info-add-cli-id.py diff --git a/src/discord-cluster-manager/migrations/20250412_01_l7Dra-user-info-fix-auth.py b/src/migrations/20250412_01_l7Dra-user-info-fix-auth.py similarity index 100% rename from src/discord-cluster-manager/migrations/20250412_01_l7Dra-user-info-fix-auth.py rename to src/migrations/20250412_01_l7Dra-user-info-fix-auth.py diff --git a/src/discord-cluster-manager/migrations/20250412_02_NN9kK-user-info-cli-drop-old.py b/src/migrations/20250412_02_NN9kK-user-info-cli-drop-old.py similarity index 100% rename from src/discord-cluster-manager/migrations/20250412_02_NN9kK-user-info-cli-drop-old.py rename to src/migrations/20250412_02_NN9kK-user-info-cli-drop-old.py diff --git a/src/discord-cluster-manager/migrations/20250506_01_38PkG-add-index-on-runs-runner-score.py b/src/migrations/20250506_01_38PkG-add-index-on-runs-runner-score.py similarity index 57% rename from src/discord-cluster-manager/migrations/20250506_01_38PkG-add-index-on-runs-runner-score.py rename to src/migrations/20250506_01_38PkG-add-index-on-runs-runner-score.py index 1c74792a..e2a347e3 100644 --- a/src/discord-cluster-manager/migrations/20250506_01_38PkG-add-index-on-runs-runner-score.py +++ b/src/migrations/20250506_01_38PkG-add-index-on-runs-runner-score.py @@ -4,11 +4,11 @@ from yoyo import step -__depends__ = {'20250412_02_NN9kK-user-info-cli-drop-old'} +__depends__ = {"20250412_02_NN9kK-user-info-cli-drop-old"} steps = [ step( "CREATE INDEX IF NOT EXISTS runs_runner_score_idx ON leaderboard.runs(runner, score)", - "DROP INDEX IF EXISTS leaderboard.runs_runner_score_idx" - ) + "DROP INDEX IF EXISTS leaderboard.runs_runner_score_idx", + ) ] diff --git a/src/discord-cluster-manager/migrations/20250617_01_c5mrF-task-split.py b/src/migrations/20250617_01_c5mrF-task-split.py similarity index 100% rename from src/discord-cluster-manager/migrations/20250617_01_c5mrF-task-split.py rename to src/migrations/20250617_01_c5mrF-task-split.py diff --git a/.github/workflows/runner.py b/src/runners/github-runner.py similarity index 86% rename from .github/workflows/runner.py rename to src/runners/github-runner.py index f5e5c42d..e408348e 100644 --- a/.github/workflows/runner.py +++ b/src/runners/github-runner.py @@ -1,14 +1,11 @@ import base64 import json -import sys import zlib from dataclasses import asdict from datetime import datetime from pathlib import Path -sys.path.append("src/discord-cluster-manager") - -from run_eval import run_config +from libkernelbot.run_eval import run_config payload = Path("payload.json").read_text() Path("payload.json").unlink() diff --git a/src/discord-cluster-manager/modal_runner.py b/src/runners/modal_runner.py similarity index 95% rename from src/discord-cluster-manager/modal_runner.py rename to src/runners/modal_runner.py index 3556a96e..eb74d173 100644 --- a/src/discord-cluster-manager/modal_runner.py +++ b/src/runners/modal_runner.py @@ -3,7 +3,8 @@ from contextlib import contextmanager from modal import App, Image -from run_eval import FullResult, SystemInfo, run_config + +from libkernelbot.run_eval import FullResult, SystemInfo, run_config # Create a stub for the Modal app # IMPORTANT: This has to stay in separate file or modal breaks @@ -28,6 +29,8 @@ "requests~=2.32.4", "packaging~=25.0", "numpy~=2.3", + "pytest", + "PyYAML", ) .pip_install( "torch~=2.7", @@ -53,10 +56,9 @@ ) cuda_image = cuda_image.add_local_python_source( - "consts", + "libkernelbot", "modal_runner", "modal_runner_archs", - "run_eval", ) diff --git a/src/discord-cluster-manager/modal_runner_archs.py b/src/runners/modal_runner_archs.py similarity index 100% rename from src/discord-cluster-manager/modal_runner_archs.py rename to src/runners/modal_runner_archs.py