Add Buildkite backend integration and pre-built image support

msaroufim · msaroufim · commit 0bcb540eaefa · 2026-02-04T19:24:41.000-08:00
Backend integration:
- Register BuildkiteLauncher in create_backend() when BUILDKITE_API_TOKEN is set
- Add BUILDKITE_API_TOKEN, BUILDKITE_ORG, BUILDKITE_PIPELINE env vars
- Results now flow to database same as GitHub/Modal

Pre-built Docker image for fast cold starts:
- Add Dockerfile with all dependencies pre-installed
- Add build-image.sh script for local image building
- Add pipeline-fast.yml for using pre-built image (~5s vs ~40s cold start)
- Update setup-node-simple.sh with BUILD_IMAGE=true option

Update skills doc with operational model for both approaches
diff --git a/SKILLS/buildkite.md b/SKILLS/buildkite.md
@@ -554,16 +554,15 @@ BUILDKITE_API_TOKEN=xxx uv run python scripts/submit_buildkite_job.py --eval ide
 
 ## Operational Model
 
-### No Pre-Built Docker Image (Current Setup)
+### Option 1: No Pre-Built Image (Current Default)
 
-The pipeline does **NOT** use a pre-built Docker image. Each job:
+The pipeline installs dependencies at runtime. Each job:
 
 1. Uses base `nvidia/cuda:12.4.0-devel-ubuntu22.04` image
-2. Installs dependencies at runtime:
+2. Installs dependencies at runtime (~30-40 seconds):
    - `uv` for Python package management
-   - Clones kernelbot repo from `buildkite-infrastructure` branch
-   - Runs `uv sync` to install project dependencies
-   - Runs `uv pip install torch triton numpy` for GPU packages
+   - Clones kernelbot repo
+   - Runs `uv sync` and `uv pip install torch triton numpy`
 3. Runs the evaluation
 
 **Advantages:**
@@ -573,17 +572,44 @@ The pipeline does **NOT** use a pre-built Docker image. Each job:
 - **No admin action needed** after code updates
 
 **Trade-off:**
-- Slightly longer job startup time (~30-40 seconds for dependency installation)
+- Slower cold starts (~40 seconds)
 
-### When Admin Action Is Needed
+### Option 2: Pre-Built Image (Fast Cold Starts)
+
+For faster cold starts (~5 seconds), build the Docker image on each node:
 
-The only time the machine admin needs to run anything is:
+```bash
+# During initial setup:
+sudo BUILDKITE_AGENT_TOKEN=xxx GPU_TYPE=test BUILD_IMAGE=true ./deployment/buildkite/setup-node-simple.sh
 
-1. **Initial setup**: Run `setup-node-simple.sh` once when onboarding a new node
-2. **Buildkite agent updates**: If Buildkite releases a new agent version (rare)
-3. **System-level changes**: NVIDIA driver updates, OS updates, etc.
+# Or build separately:
+./deployment/buildkite/build-image.sh
+```
+
+Then update the Buildkite pipeline config to use the local image:
+```yaml
+image: "kernelbot:latest"
+```
+
+**When to rebuild the image:**
+- When dependencies change (new PyTorch version, new packages)
+- When you want the latest kernelbot code baked in
+- NOT needed for problem/task changes (those come via config)
+
+**Rebuild command:**
+```bash
+./deployment/buildkite/build-image.sh
+```
+
+### When Admin Action Is Needed
 
-Code changes to kernelbot require **no admin action** - the pipeline clones fresh code each run.
+| Scenario | Action Required |
+|----------|-----------------|
+| Code changes (no deps) | None - pipeline clones fresh code |
+| Dependency changes | Rebuild image: `./build-image.sh` |
+| Initial node setup | Run `setup-node-simple.sh` once |
+| NVIDIA driver updates | May need to rebuild image |
+| Buildkite agent updates | Rare - Buildkite handles this |
 
 ### Shared Evaluation Logic
 
diff --git a/deployment/buildkite/Dockerfile b/deployment/buildkite/Dockerfile
@@ -1,42 +1,42 @@
 # Kernelbot evaluation image
+# Pre-built with all dependencies for fast cold starts
 FROM nvidia/cuda:12.4.0-devel-ubuntu22.04
 
 ENV DEBIAN_FRONTEND=noninteractive
 ENV PYTHONUNBUFFERED=1
 
 # System packages
 RUN apt-get update && apt-get install -y --no-install-recommends \
-    python3.11 \
-    python3.11-dev \
-    python3.11-venv \
-    python3-pip \
-    git \
-    wget \
     curl \
+    ca-certificates \
+    git \
     build-essential \
     ninja-build \
     cmake \
     && rm -rf /var/lib/apt/lists/*
 
-# Set Python 3.11 as default
-RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.11 1 && \
-    update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 1
+# Install uv
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh
+ENV PATH="/root/.local/bin:$PATH"
+
+# Clone and install kernelbot
+WORKDIR /opt/kernelbot
+RUN git clone --depth 1 --branch buildkite-infrastructure https://github.com/gpu-mode/kernelbot.git .
+
+# Install dependencies with uv
+RUN uv sync
 
-# Upgrade pip
-RUN python -m pip install --no-cache-dir --upgrade pip setuptools wheel
+# Install PyTorch and GPU packages
+RUN uv pip install torch triton numpy --index-url https://download.pytorch.org/whl/cu124
 
-# PyTorch + CUDA
-RUN pip install --no-cache-dir \
-    torch==2.4.0 \
-    triton \
-    numpy \
-    scipy
+# Ensure venv is activated for any commands
+ENV VIRTUAL_ENV=/opt/kernelbot/.venv
+ENV PATH="$VIRTUAL_ENV/bin:$PATH"
 
-# Copy kernelbot
-WORKDIR /app
-COPY pyproject.toml .
-COPY src/ src/
-RUN pip install --no-cache-dir -e .
+# Verify installation
+RUN python -c "import torch; print(f'PyTorch {torch.__version__}')" && \
+    python -c "import triton; print(f'Triton installed')" && \
+    python -c "from libkernelbot.run_eval import run_config; print('kernelbot installed')"
 
 # Default command
-CMD ["python", "/app/src/runners/buildkite-runner.py"]
+CMD ["python", "/opt/kernelbot/src/runners/buildkite-runner.py"]
diff --git a/deployment/buildkite/build-image.sh b/deployment/buildkite/build-image.sh
@@ -0,0 +1,43 @@
+#!/bin/bash
+# Build the kernelbot Docker image locally on a GPU node
+# Usage: ./build-image.sh [--push]
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
+
+IMAGE_NAME="${KERNELBOT_IMAGE:-kernelbot:latest}"
+BRANCH="${KERNELBOT_BRANCH:-buildkite-infrastructure}"
+
+echo "=== Building Kernelbot Image ==="
+echo "Image: $IMAGE_NAME"
+echo "Branch: $BRANCH"
+echo ""
+
+# Update Dockerfile to use correct branch
+sed -i "s|--branch [a-zA-Z0-9_-]*|--branch $BRANCH|g" "$SCRIPT_DIR/Dockerfile" 2>/dev/null || \
+    sed -i '' "s|--branch [a-zA-Z0-9_-]*|--branch $BRANCH|g" "$SCRIPT_DIR/Dockerfile"
+
+echo "Building image..."
+docker build -t "$IMAGE_NAME" -f "$SCRIPT_DIR/Dockerfile" "$REPO_ROOT"
+
+echo ""
+echo "=== Build Complete ==="
+echo "Image: $IMAGE_NAME"
+docker images "$IMAGE_NAME"
+
+# Optional: push to registry
+if [[ "${1:-}" == "--push" ]]; then
+    REGISTRY="${KERNELBOT_REGISTRY:-ghcr.io/gpu-mode}"
+    REMOTE_IMAGE="$REGISTRY/kernelbot:latest"
+    echo ""
+    echo "Pushing to $REMOTE_IMAGE..."
+    docker tag "$IMAGE_NAME" "$REMOTE_IMAGE"
+    docker push "$REMOTE_IMAGE"
+    echo "Pushed: $REMOTE_IMAGE"
+fi
+
+echo ""
+echo "To use this image, update your pipeline config:"
+echo "  image: \"$IMAGE_NAME\""
diff --git a/deployment/buildkite/pipeline-fast.yml b/deployment/buildkite/pipeline-fast.yml
@@ -0,0 +1,53 @@
+# Kernelbot Fast Evaluation Pipeline
+# Uses pre-built image for fast cold starts (~5s vs ~40s)
+#
+# Prerequisites:
+# 1. Build image on node: ./deployment/buildkite/build-image.sh
+# 2. Or pull from registry: docker pull ghcr.io/gpu-mode/kernelbot:latest
+
+steps:
+  - label: ":rocket: Kernel Evaluation"
+    agents:
+      queue: "${KERNELBOT_QUEUE:-test}"
+
+    plugins:
+      - docker#v5.11.0:
+          image: "${KERNELBOT_IMAGE:-kernelbot:latest}"
+          always-pull: false
+          gpus: "all"
+          propagate-environment: true
+          shell: ["/bin/bash", "-e", "-c"]
+          environment:
+            - NVIDIA_VISIBLE_DEVICES
+            - CUDA_VISIBLE_DEVICES
+            - KERNELBOT_PAYLOAD
+            - KERNELBOT_RUN_ID
+          cpus: "${KERNELBOT_CPUS:-8}"
+          memory: "${KERNELBOT_MEMORY:-64g}"
+          workdir: /workdir
+
+    command: |
+      set -e
+
+      echo "=== Environment ==="
+      echo "NVIDIA_VISIBLE_DEVICES=$NVIDIA_VISIBLE_DEVICES"
+      echo "KERNELBOT_RUN_ID=$KERNELBOT_RUN_ID"
+      nvidia-smi -L
+
+      echo ""
+      echo "=== Running Evaluation ==="
+      cd /opt/kernelbot
+      python src/runners/buildkite-runner.py
+
+      echo ""
+      echo "=== Copying Artifacts ==="
+      cp result.json /workdir/result.json
+      cp -r profile_data /workdir/profile_data 2>/dev/null || true
+
+      echo "=== Done ==="
+
+    artifact_paths:
+      - "result.json"
+      - "profile_data/*"
+
+    timeout_in_minutes: 15
diff --git a/deployment/buildkite/setup-node-simple.sh b/deployment/buildkite/setup-node-simple.sh
@@ -203,3 +203,27 @@ echo '    - label: "GPU Test"'
 echo '      command: "echo NVIDIA_VISIBLE_DEVICES=$$NVIDIA_VISIBLE_DEVICES && nvidia-smi -L"'
 echo '      agents:'
 echo "        queue: \"${GPU_TYPE}\""
+
+# === BUILD DOCKER IMAGE (optional) ===
+if [[ "${BUILD_IMAGE:-}" == "true" ]]; then
+    echo ""
+    echo "=== Building Docker Image ==="
+    SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+    if [[ -f "$SCRIPT_DIR/Dockerfile" ]]; then
+        docker build -t kernelbot:latest -f "$SCRIPT_DIR/Dockerfile" "$SCRIPT_DIR/../.."
+        echo "Docker image built: kernelbot:latest"
+        echo ""
+        echo "To use the fast pipeline, update Buildkite config to use:"
+        echo "  image: \"kernelbot:latest\""
+    else
+        echo "WARNING: Dockerfile not found at $SCRIPT_DIR/Dockerfile"
+        echo "Clone the repo first: git clone https://github.com/gpu-mode/kernelbot.git"
+    fi
+fi
+
+echo ""
+echo "For faster cold starts, build the Docker image:"
+echo "  BUILD_IMAGE=true $0"
+echo "Or manually:"
+echo "  ./deployment/buildkite/build-image.sh"
diff --git a/src/kernelbot/env.py b/src/kernelbot/env.py
@@ -33,6 +33,11 @@
 env.GITHUB_WORKFLOW_BRANCH = os.getenv("GITHUB_WORKFLOW_BRANCH", get_github_branch_name())
 env.PROBLEMS_REPO = os.getenv("PROBLEMS_REPO")
 
+# Buildkite-specific constants
+env.BUILDKITE_API_TOKEN = os.getenv("BUILDKITE_API_TOKEN")
+env.BUILDKITE_ORG = os.getenv("BUILDKITE_ORG", "gpu-mode")
+env.BUILDKITE_PIPELINE = os.getenv("BUILDKITE_PIPELINE", "kernelbot")
+
 # Directory that will be used for local problem development.
 env.PROBLEM_DEV_DIR = os.getenv("PROBLEM_DEV_DIR", "examples")
 
diff --git a/src/kernelbot/main.py b/src/kernelbot/main.py
@@ -16,7 +16,8 @@
 from libkernelbot import consts
 from libkernelbot.backend import KernelBackend
 from libkernelbot.background_submission_manager import BackgroundSubmissionManager
-from libkernelbot.launchers import GitHubLauncher, ModalLauncher
+from libkernelbot.launchers import BuildkiteLauncher, GitHubLauncher, ModalLauncher
+from libkernelbot.launchers.buildkite import BuildkiteConfig
 from libkernelbot.utils import setup_logging
 
 logger = setup_logging(__name__)
@@ -29,6 +30,17 @@ def create_backend(debug_mode: bool = False) -> KernelBackend:
     backend.register_launcher(
         GitHubLauncher(env.GITHUB_REPO, env.GITHUB_TOKEN, env.GITHUB_WORKFLOW_BRANCH)
     )
+
+    # Register Buildkite launcher if API token is configured
+    if env.BUILDKITE_API_TOKEN:
+        buildkite_config = BuildkiteConfig(
+            org_slug=env.BUILDKITE_ORG,
+            pipeline_slug=env.BUILDKITE_PIPELINE,
+            api_token=env.BUILDKITE_API_TOKEN,
+        )
+        backend.register_launcher(BuildkiteLauncher(buildkite_config))
+        logger.info("Buildkite launcher registered")
+
     return backend
 
 
