update

Author: Mark Saroufim

.buildkite/pipeline.yml

@@ -9,6 +9,10 @@
 # Environment variables passed from BuildkiteLauncher:
 #   - SUBMISSION_PAYLOAD: Base64-encoded, zlib-compressed submission config
 #   - GPU_QUEUE: Queue name for agent routing
+#
+# Note: timeout_in_minutes is set high (60) as a safety net.
+# The BuildkiteLauncher handles dynamic timeouts based on submission mode
+# and will cancel jobs that exceed their configured timeout.
 
 steps:
   - label: ":gpu: Run Kernel Submission"
@@ -19,7 +23,8 @@ steps:
     agents:
       # Route to agent with matching queue tag
       queue: "${GPU_QUEUE}"
-    timeout_in_minutes: 15
+    # Safety timeout - BuildkiteLauncher handles actual timeout enforcement
+    timeout_in_minutes: 60
     artifact_paths:
       - "result.json"
       - "profile_data/**/*"

src/libkernelbot/consts.py

@@ -135,6 +135,14 @@ class RankCriterion(Enum):
     "MI300": None,
     "MI300x8": None,
     "MI250": None,
+    # Buildkite GPU types (vendor-prefixed queue names)
+    "nvidia-h100": "90a",
+    "nvidia-b200": "100",
+    "nvidia-a100": "80",
+    "nebius-h100": "90a",
+    "amd-mi300": None,
+    "amd-mi250": None,
+    "google-tpu": None,
 }
 
 

src/libkernelbot/launchers/buildkite.py

@@ -221,7 +221,8 @@ async def _wait_for_completion(
                 if state in ("passed", "failed", "canceled", "blocked"):
                     if state != "passed":
                         logger.warning(f"Build #{build_number} finished with state: {state}")
-                    return
+                        raise RuntimeError(f"Build #{build_number} {state}")
+                    return state
 
                 await callback(state, elapsed)
                 await asyncio.sleep(10)  # Poll every 10 seconds

src/runners/buildkite-runner.py

@@ -13,6 +13,8 @@
 import base64
 import json
 import os
+import sys
+import traceback
 import zlib
 from dataclasses import asdict
 from datetime import datetime
@@ -28,28 +30,52 @@ def serialize(obj: object):
     raise TypeError(f"Type {type(obj)} not serializable")
 
 
+def write_error_result(error_message: str):
+    """Write an error result to result.json when execution fails."""
+    error_result = {
+        "success": False,
+        "error": error_message,
+        "runs": {},
+        "system": {},
+    }
+    Path("result.json").write_text(json.dumps(error_result, default=serialize))
+
+
 def main():
-    # Get payload from environment variable
-    payload_b64 = os.environ.get("SUBMISSION_PAYLOAD")
-    if not payload_b64:
-        raise RuntimeError("SUBMISSION_PAYLOAD environment variable not set")
-
-    # Decompress and parse config
-    payload = zlib.decompress(base64.b64decode(payload_b64)).decode("utf-8")
-    config = json.loads(payload)
-
-    # Run the submission
-    result = run_config(config)
-
-    # Write result to file for artifact upload
-    result_dict = asdict(result)
-    Path("result.json").write_text(json.dumps(result_dict, default=serialize))
-
-    # Create profile_data directory if profiling was enabled
-    # (profile artifacts will be written there by run_config)
-    profile_dir = Path("profile_data")
-    if profile_dir.exists():
-        print(f"Profile data available in {profile_dir}")
+    try:
+        # Get payload from environment variable
+        payload_b64 = os.environ.get("SUBMISSION_PAYLOAD")
+        if not payload_b64:
+            write_error_result("SUBMISSION_PAYLOAD environment variable not set")
+            sys.exit(1)
+
+        # Decompress and parse config
+        try:
+            payload = zlib.decompress(base64.b64decode(payload_b64)).decode("utf-8")
+            config = json.loads(payload)
+        except Exception as e:
+            write_error_result(f"Failed to decompress/parse payload: {e}")
+            sys.exit(1)
+
+        # Run the submission
+        result = run_config(config)
+
+        # Write result to file for artifact upload
+        result_dict = asdict(result)
+        Path("result.json").write_text(json.dumps(result_dict, default=serialize))
+
+        # Create profile_data directory if profiling was enabled
+        # (profile artifacts will be written there by run_config)
+        profile_dir = Path("profile_data")
+        if profile_dir.exists():
+            print(f"Profile data available in {profile_dir}")
+
+    except Exception as e:
+        # Catch any unexpected errors and write them to result.json
+        error_msg = f"Runner error: {e}\n{traceback.format_exc()}"
+        print(error_msg, file=sys.stderr)
+        write_error_result(error_msg)
+        sys.exit(1)
 
 
 if __name__ == "__main__":

tests/test_buildkite.py

@@ -70,7 +70,8 @@ def mock_config(self):
         return {
             "lang": "py",
             "mode": "test",
-            "files": {"main.py": "print('hello')"},
+            "main": "main.py",
+            "sources": {"main.py": "print('hello')"},
             "tests": [],
             "benchmarks": [],
             "test_timeout": 180,
@@ -141,6 +142,9 @@ async def test_run_submission_creates_build(self, launcher, mock_config, gpu_typ
                     assert "SUBMISSION_PAYLOAD" in body["env"]
                     assert body["env"]["GPU_QUEUE"] == "nvidia-h100"
 
+                    # Verify result
+                    assert result.success is True
+
     @pytest.mark.asyncio
     async def test_run_submission_handles_api_error(self, launcher, mock_config, gpu_type):
         """Test that API errors are handled gracefully."""