fold GitHubRun into github launcher

Author: ngc92

src/discord-cluster-manager/bot.py

@@ -4,6 +4,7 @@
 
 import consts
 import discord
+import env
 import uvicorn
 from api.main import app, init_api
 from cogs.admin_cog import AdminCog
@@ -78,7 +79,7 @@ async def setup_hook(self):
             # Load cogs
             submit_cog = SubmitCog(self)
             submit_cog.register_launcher(ModalLauncher(consts.MODAL_CUDA_INCLUDE_DIRS))
-            submit_cog.register_launcher(GitHubLauncher())
+            submit_cog.register_launcher(GitHubLauncher(env.GITHUB_REPO, env.GITHUB_TOKEN))
             await self.add_cog(submit_cog)
             await self.add_cog(BotManagerCog(self))
             await self.add_cog(LeaderboardCog(self))

src/discord-cluster-manager/github_runner.py

@@ -1,20 +1,28 @@
+import asyncio
 import datetime
 import json
+import pprint
+import tempfile
+import zipfile
+from typing import Awaitable, Callable, Optional
 
+import requests
 from consts import AMD_REQUIREMENTS, GPU, NVIDIA_REQUIREMENTS, GitHubGPU, GPUType
-from github_runner import GitHubRun
+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 utils import get_github_branch_name, setup_logging
 
 from .launcher import Launcher
 
 logger = setup_logging()
 
 
 class GitHubLauncher(Launcher):
-    def __init__(self):
+    def __init__(self, repo: str, token: str):
         super().__init__(name="GitHub", gpus=GitHubGPU)
+        self.repo = repo
+        self.token = token
 
     async def run_submission(
         self, config: dict, gpu_type: GPU, status: RunProgressReporter
@@ -37,7 +45,7 @@ async def run_submission(
             logger.info(f"Running on {gpu_name} amd gpu")
 
         workflow_file = selected_gpu.value
-        run = GitHubRun(workflow_file)
+        run = GitHubRun(self.repo, self.token, workflow_file)
 
         payload = json.dumps(config)
 
@@ -89,8 +97,152 @@ async def run_submission(
         system = SystemInfo(**data.get("system", {}))
         return FullResult(success=True, error="", runs=runs, system=system)
 
-    async def wait_callback(self, run: GitHubRun, status: RunProgressReporter):
+    async def wait_callback(self, run: "GitHubRun", status: RunProgressReporter):
         await status.update(
             f"⏳ Workflow [{run.run_id}]({run.html_url}): {run.status} "
             f"({run.elapsed_time.total_seconds():.1f}s)"
         )
+
+
+class GitHubRun:
+    def __init__(self, repo: str, token: str, workflow_file: str):
+        gh = Github(token)
+        self.repo = gh.get_repo(repo)
+        self.token = token
+        self.workflow_file = workflow_file
+        self.run: Optional[WorkflowRun.WorkflowRun] = None
+        self.start_time = None
+
+    @property
+    def run_id(self):
+        if self.run is None:
+            return None
+        return self.run.id
+
+    @property
+    def html_url(self):
+        if self.run is None:
+            return None
+        return self.run.html_url
+
+    @property
+    def status(self):
+        if self.run is None:
+            return None
+        return self.run.status
+
+    @property
+    def elapsed_time(self):
+        if self.start_time is None:
+            return None
+        return datetime.datetime.now(datetime.timezone.utc) - self.start_time
+
+    async def trigger(self, inputs: dict) -> bool:
+        """
+        Trigger this run with the provided inputs.
+        Sets `self.run` to the new WorkflowRun on success.
+
+        Returns: Whether the run was successfully triggered,
+        """
+        trigger_time = datetime.datetime.now(datetime.timezone.utc)
+        try:
+            workflow = self.repo.get_workflow(self.workflow_file)
+        except UnknownObjectException as e:
+            logger.error(f"Could not find workflow {self.workflow_file}", exc_info=e)
+            raise ValueError(f"Could not find workflow {self.workflow_file}") from e
+
+        logger.debug(
+            "Dispatching workflow %s on branch %s with inputs %s",
+            self.workflow_file,
+            get_github_branch_name(),
+            pprint.pformat(inputs),
+        )
+        success = workflow.create_dispatch(get_github_branch_name(), inputs=inputs)
+        if success:
+            await asyncio.sleep(2)
+            runs = list(workflow.get_runs())
+
+            for run in runs:
+                if run.created_at.replace(tzinfo=datetime.timezone.utc) > trigger_time:
+                    self.run = run
+                    return True
+        return False
+
+    async def wait_for_completion(
+        self, callback: Callable[["GitHubRun"], Awaitable[None]], timeout_minutes: int = 5
+    ):
+        if self.run is None:
+            raise ValueError("Run needs to be triggered before a status check!")
+
+        self.start_time = datetime.datetime.now(datetime.timezone.utc)
+        timeout = datetime.timedelta(minutes=timeout_minutes)
+
+        while True:
+            try:
+                # update run status
+                self.run = run = self.repo.get_workflow_run(self.run_id)
+
+                if self.elapsed_time > timeout:
+                    try:
+                        self.run.cancel()
+                        # Wait briefly to ensure cancellation is processed
+                        # And Verify the run was actually cancelled
+                        await asyncio.sleep(5)
+                        run = self.repo.get_workflow_run(self.run_id)
+                        if run.status != "completed":
+                            logger.warning(f"Failed to cancel workflow run {self.run_id}")
+                    except Exception as e:
+                        logger.error(f"Error cancelling workflow: {str(e)}", exc_info=e)
+                        raise
+
+                    logger.warning(
+                        f"Workflow {self.run_id} cancelled - "
+                        f"exceeded {timeout_minutes} minute timeout"
+                    )
+                    raise TimeoutError(
+                        f"Workflow {self.run_id} cancelled - "
+                        f"exceeded {timeout_minutes} minute timeout"
+                    )
+
+                if run.status == "completed":
+                    return
+
+                await callback(self)
+                await asyncio.sleep(3)
+            except TimeoutError:
+                raise
+            except Exception as e:
+                logger.error(f"Error waiting for GitHub run {self.run_id}: {e}", exc_info=e)
+                raise
+
+    async def download_artifacts(self) -> dict:
+        logger.info("Attempting to download artifacts for run %s", self.run_id)
+        artifacts = self.run.get_artifacts()
+
+        extracted = {}
+
+        for artifact in artifacts:
+            url = artifact.archive_download_url
+            headers = {"Authorization": f"token {self.token}"}
+            response = requests.get(url, headers=headers)
+
+            if response.status_code == 200:
+                with tempfile.NamedTemporaryFile("w+b") as temp:
+                    temp.write(response.content)
+                    temp.flush()
+
+                    with zipfile.ZipFile(temp.name) as z:
+                        artifact_dict = {}
+                        for file in z.namelist():
+                            with z.open(file) as f:
+                                artifact_dict[file] = f.read()
+
+                extracted[artifact.name] = artifact_dict
+            else:
+                raise RuntimeError(
+                    f"Failed to download artifact {artifact.name}. "
+                    f"Status code: {response.status_code}"
+                )
+
+        logger.info("Download artifacts for run %s: %s", self.run_id, list(extracted.keys()))
+        return extracted