Fix Timeout Bug (#280)

* use workflows timeout

* remove magic number

* lint

* Update nvidia_workflow.yml

* Update nvidia_workflow.yml

* fix verify tests

* make reviewable

* add buffer

Author: PaliC

.github/workflows/nvidia_workflow.yml

@@ -29,6 +29,7 @@ jobs:
       shell: bash
       run: |
         # Extract the payload content without printing it
+        apt-get update && apt-get install -y jq
         PAYLOAD=$(jq -r '.inputs.payload' $GITHUB_EVENT_PATH)
         
         # Apply mask to the extracted content

src/discord-cluster-manager/consts.py

@@ -145,6 +145,8 @@ class RankCriterion(Enum):
 ]
 MODAL_CUDA_INCLUDE_DIRS = ["/ThunderKittens/include"]
 
+DEFAULT_GITHUB_TIMEOUT_MINUTES = 10  # Default timeout for GitHub launcher in minutes
+
 NVIDIA_REQUIREMENTS = """
 numpy
 torch
@@ -157,3 +159,6 @@ class RankCriterion(Enum):
 --index-url https://download.pytorch.org/whl/rocm6.2.4
 torch
 """
+
+# A buffer for timeouts to account for github setup time
+TIMEOUT_BUFFER_MINUTES = 2

src/discord-cluster-manager/launchers/github.py

@@ -1,13 +1,22 @@
 import asyncio
 import datetime
 import json
+import math
 import pprint
 import tempfile
 import zipfile
 from typing import Awaitable, Callable, Optional
 
 import requests
-from consts import AMD_REQUIREMENTS, GPU, NVIDIA_REQUIREMENTS, GitHubGPU
+from consts import (
+    AMD_REQUIREMENTS,
+    DEFAULT_GITHUB_TIMEOUT_MINUTES,
+    GPU,
+    NVIDIA_REQUIREMENTS,
+    TIMEOUT_BUFFER_MINUTES,
+    GitHubGPU,
+    SubmissionMode,
+)
 from github import Github, UnknownObjectException, WorkflowRun
 from report import RunProgressReporter
 from run_eval import CompileResult, EvalResult, FullResult, RunResult, SystemInfo
@@ -17,6 +26,15 @@
 
 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):
@@ -70,7 +88,13 @@ async def run_submission(
 
         await status.push("⏳ Waiting for workflow to start...")
         logger.info("Waiting for workflow to start...")
-        await run.wait_for_completion(lambda x: self.wait_callback(x, status))
+
+        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...")