import asyncio import base64 import dataclasses import datetime import io import json import math import pprint import uuid import zipfile import zlib from typing import Awaitable, Callable, Optional import github import requests from github import Github, UnknownObjectException from github.GithubObject import NotSet, Opt from github.Workflow import Workflow from github.WorkflowRun import WorkflowRun from libkernelbot.consts import ( AMD_REQUIREMENTS, DEFAULT_GITHUB_TIMEOUT_MINUTES, GPU, NVIDIA_REQUIREMENTS, TIMEOUT_BUFFER_MINUTES, GitHubGPU, SubmissionMode, ) from libkernelbot.report import RunProgressReporter from libkernelbot.run_eval import ( CompileResult, EvalResult, FullResult, ProfileResult, RunResult, SystemInfo, ) from libkernelbot.utils import KernelBotError, setup_logging from .launcher import Launcher logger = setup_logging() def get_timeout(config: dict) -> int: mode = config.get("mode") sec_map = { SubmissionMode.TEST.value: config.get("test_timeout"), SubmissionMode.BENCHMARK.value: config.get("benchmark_timeout"), SubmissionMode.LEADERBOARD.value: config.get("ranked_timeout"), } seconds = sec_map.get(mode) or DEFAULT_GITHUB_TIMEOUT_MINUTES * 60 return math.ceil(seconds / 60) class GitHubLauncher(Launcher): def __init__(self, repo: str, token: str, branch: str): super().__init__(name="GitHub", gpus=GitHubGPU) self.repo = repo self.token = token self.branch = branch async def run_submission( # noqa: C901 self, config: dict, gpu_type: GPU, status: RunProgressReporter ) -> FullResult: gpu_vendor = None if gpu_type.value in ["MI300", "MI250", "MI300x8"]: selected_workflow = "amd_workflow.yml" runner_name = { "MI300": "amdgpu-mi300-x86-64", "MI250": "amdgpu-mi250-x86-64", "MI300x8": "amdgpu-mi300-8-x86-64", }[gpu_type.value] gpu_vendor = "AMD" requirements = AMD_REQUIREMENTS elif gpu_type.value == "NVIDIA": selected_workflow = "nvidia_workflow.yml" gpu_vendor = "NVIDIA" requirements = NVIDIA_REQUIREMENTS else: raise ValueError(f"Invalid GPU type: {gpu_type.value}") lang = config["lang"] if lang == "cu" and gpu_vendor == "AMD": # TODO implement HIP raise NotImplementedError("Cannot use CUDA runs with AMD GPUs") lang_name = {"py": "Python", "cu": "CUDA"}[lang] logger.info(f"Attempting to trigger GitHub action for {lang_name} on {selected_workflow}") run = GitHubRun(self.repo, self.token, self.branch, selected_workflow) logger.info(f"Successfully created GitHub run: {run.run_id}") payload = base64.b64encode(zlib.compress(json.dumps(config).encode("utf-8"))).decode( "utf-8" ) inputs = {"payload": payload} if lang == "py": inputs["requirements"] = requirements if gpu_vendor == "AMD": inputs["runner"] = runner_name if not await run.trigger(inputs): raise RuntimeError("Failed to trigger GitHub Action. Please check the configuration.") await status.push("⏳ Waiting for workflow to start...") logger.info("Waiting for workflow to start...") timeout = get_timeout(config) + TIMEOUT_BUFFER_MINUTES logger.info(f"Waiting for workflow to complete... (timeout: {timeout} minutes)") await run.wait_for_completion( lambda x: self.wait_callback(x, status), timeout_minutes=timeout ) await status.update(f"Workflow [{run.run_id}](<{run.html_url}>) completed") logger.info(f"Workflow [{run.run_id}]({run.html_url}) completed") await status.push("Downloading artifacts...") logger.info("Downloading artifacts...") index = run.get_artifact_index() if "run-result" not in index: logger.error("Could not find `run-result` among artifacts: %s", index.keys()) await status.push("Downloading artifacts... failed") return FullResult( success=False, error="Could not download artifacts", runs={}, system=SystemInfo() ) artifact = await run.download_artifact(index["run-result"]) logs = artifact["result.json"].decode("utf-8") await status.update("Downloading artifacts... done") logger.info("Downloading artifacts... done") data = json.loads(logs) runs = {} # convert json back to EvalResult structures, which requires # special handling for datetime and our dataclasses. for k, v in data["runs"].items(): comp_res = None if v.get("compilation") is None else CompileResult(**v["compilation"]) run_res = None if v.get("run") is None else RunResult(**v["run"]) profile_res = None if v.get("profile") is None else ProfileResult(**v["profile"]) # Update profile artifact to the actual download URL. # For the GitHub launcher the profile_artifact currently just contains # the name of the artifact. if profile_res is not None and "profile-data" in index: profile_res.download_url = index["profile-data"].public_download_url res = EvalResult( start=datetime.datetime.fromisoformat(v["start"]), end=datetime.datetime.fromisoformat(v["end"]), compilation=comp_res, run=run_res, profile=profile_res, ) runs[k] = res system = SystemInfo(**data.get("system", {})) return FullResult(success=True, error="", runs=runs, system=system) 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)" ) @dataclasses.dataclass class GitHubArtifact: name: str archive_download_url: str public_download_url: str _WORKFLOW_FILE_CACHE: dict[str, Workflow] = {} def patched_create_dispatch( workflow: Workflow, ref: github.Branch.Branch | github.Tag.Tag | github.Commit.Commit | str, inputs: Opt[dict] = NotSet, ) -> bool: """ :calls: `POST /repos/{owner}/{repo}/actions/workflows/{workflow_id}/dispatches `_ """ assert ( isinstance(ref, github.Branch.Branch) or isinstance(ref, github.Tag.Tag) or isinstance(ref, github.Commit.Commit) or isinstance(ref, str) ), ref assert inputs is NotSet or isinstance(inputs, dict), inputs if isinstance(ref, github.Branch.Branch): ref = ref.name elif isinstance(ref, github.Commit.Commit): ref = ref.sha elif isinstance(ref, github.Tag.Tag): ref = ref.name if inputs is NotSet: inputs = {} status, _, _ = workflow._requester.requestJson( "POST", f"{workflow.url}/dispatches", input={"ref": ref, "inputs": inputs} ) return status == 200 or status == 204 class GitHubRun: def __init__(self, repo: str, token: str, branch: str, workflow_file: str): gh = Github(token) try: self.repo = gh.get_repo(repo) except UnknownObjectException as e: raise KernelBotError(f"Could not find GitHub repository {repo}: 404") from e self.token = token self.branch = branch 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 get_workflow(self) -> Workflow: if self.workflow_file in _WORKFLOW_FILE_CACHE: logger.info(f"Returning cached workflow {self.workflow_file}") return _WORKFLOW_FILE_CACHE[self.workflow_file] logger.info(f"Fetching workflow {self.workflow_file} from GitHub") workflow = self.repo.get_workflow(self.workflow_file) _WORKFLOW_FILE_CACHE[self.workflow_file] = workflow return workflow 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, """ run_id = str(uuid.uuid4()) inputs_with_run_id = {**inputs, "run_id": run_id} if self.workflow_file == "amd_workflow.yml": expected_run_name = f"AMD Job - {run_id}" elif self.workflow_file == "nvidia_workflow.yml": expected_run_name = f"NVIDIA Job - {run_id}" else: raise ValueError(f"Unknown workflow file: {self.workflow_file}") trigger_time = datetime.datetime.now(datetime.timezone.utc) try: workflow = await self.get_workflow() 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.info( "Dispatching workflow %s on branch %s with run_id %s", self.workflow_file, self.branch, run_id, ) logger.debug( "Dispatching workflow %s on branch %s with inputs %s", self.workflow_file, self.branch, pprint.pformat(inputs_with_run_id), ) success = await asyncio.to_thread( patched_create_dispatch, workflow, self.branch, inputs=inputs_with_run_id ) # noqa: E501 if success: wait_seconds = 10 logger.info( f"Workflow dispatch successful. Waiting {wait_seconds}s for the run to appear..." ) await asyncio.sleep(wait_seconds) recent_runs_paginated = await asyncio.to_thread( workflow.get_runs, event="workflow_dispatch" ) logger.info(f"Looking for workflow run with name: '{expected_run_name}'") found_run = None runs_checked = 0 try: run_iterator = recent_runs_paginated.__iter__() while runs_checked < 100: try: run = next(run_iterator) runs_checked += 1 logger.debug( f"Checking run {run.id} with name '{run.name}'" f" created at {run.created_at.isoformat()}" ) if run.name == expected_run_name and run.created_at.replace( tzinfo=datetime.timezone.utc ) > trigger_time - datetime.timedelta(seconds=30): found_run = run logger.info( f"Found matching workflow run: ID {found_run.id} " f"with name '{found_run.name}'" ) break except StopIteration: logger.debug("Reached end of recent runs list.") break except Exception as e: logger.error(f"Error iterating through recent runs: {e}", exc_info=True) return False if found_run: self.run = found_run return True else: logger.warning( f"Could not find a workflow run with name '{expected_run_name}' " f"created after {trigger_time.isoformat()}." ) return False else: logger.error( f"Failed to dispatch workflow {self.workflow_file} on branch {self.branch}." ) return False async def wait_for_completion( self, callback: Callable[["GitHubRun"], Awaitable[None]], timeout_minutes: int = 10 ): 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: run_update = await asyncio.to_thread(self.repo.get_workflow_run, self.run_id) self.run = run = run_update 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(30) # Yield control while waiting except TimeoutError: raise # Re-raise the specific TimeoutError from the timeout block except Exception as e: logger.error(f"Error waiting for GitHub run {self.run_id}: {e}", exc_info=e) raise # Re-raise other exceptions def get_artifact_index(self) -> dict[str, GitHubArtifact]: logger.info("Creating artifact index for run %s", self.run_id) artifacts = self.run.get_artifacts() extracted = {} for artifact in artifacts: extracted[artifact.name] = GitHubArtifact( name=artifact.name, archive_download_url=artifact.archive_download_url, # Non-machine users cannot download from the archive_download_url and # the GitHub API does not give us access to the public download url. public_download_url=f"{self.repo.html_url}/actions/runs/{self.run_id}/artifacts/{artifact.id}", ) return extracted async def download_artifact(self, artifact: GitHubArtifact) -> dict: logger.info("Attempting to download artifact '%s' for run %s", artifact.name, self.run_id) url = artifact.archive_download_url headers = {"Authorization": f"token {self.token}"} response = requests.get(url, headers=headers) if response.status_code == 200: artifact_dict = {} with zipfile.ZipFile(io.BytesIO(response.content)) as z: for file in z.namelist(): with z.open(file) as f: artifact_dict[file] = f.read() logger.info("Downloaded artifact '%s' for run %s", artifact.name, self.run_id) return artifact_dict else: raise RuntimeError( f"Failed to download artifact {artifact.name}. Status code: {response.status_code}" )