Merge remote-tracking branch 'origin/main' into alexcheema/speculative-decoding

PlaceInstance was missing these fields that are accessed in placement.py
when creating MlxJaccl and MlxRing instances with speculative decoding
support, causing type errors after merging main.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

.dockerignore

@@ -1,15 +0,0 @@
-.venv/
-.direnv/
-target/
-.git/
-.idea/
-.pytest_cache/
-.ruff_cache/
-dashboard/node_modules/
-dashboard/.svelte-kit/
-dashboard/build/
-dist/
-*.pdb
-**/__pycache__
-**/.DS_Store
-.mlx_typings/

.github/workflows/e2e.yml

@@ -1,44 +0,0 @@
-name: e2e-tests
-
-on:
-  push:
-    branches:
-      - e2e-tests
-  pull_request:
-    branches:
-      - staging
-      - main
-
-jobs:
-  e2e:
-    runs-on: ubuntu-latest
-    timeout-minutes: 45
-    steps:
-      - name: Free up disk space
-        run: |
-          sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc \
-            /opt/hostedtoolcache /usr/local/share/boost /usr/share/swift \
-            /opt/microsoft /opt/az
-          docker system prune -af
-          df -h /
-
-      - name: Checkout repository
-        uses: actions/checkout@v4
-        with:
-          lfs: false
-
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3
-
-      - name: Build E2E image with cache
-        uses: docker/build-push-action@v6
-        with:
-          context: .
-          file: e2e/Dockerfile
-          tags: exo-e2e:latest
-          load: true
-          cache-from: type=gha
-          cache-to: type=gha,mode=max
-
-      - name: Run E2E tests
-        run: python3 e2e/run_all.py

.github/workflows/pipeline.yml

@@ -8,6 +8,33 @@ on:
       - main
 
 jobs:
+  typecheck:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          lfs: false
+
+      - uses: cachix/install-nix-action@v31
+        with:
+          nix_path: nixpkgs=channel:nixos-unstable
+
+      - uses: cachix/cachix-action@v14
+        name: Configure Cachix
+        with:
+          name: exo
+          authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
+
+      - name: Load nix develop environment
+        run: nix run github:nicknovitski/nix-develop/v1
+
+      - name: Sync dependencies
+        run: uv sync --all-packages
+
+      - name: Run type checker
+        run: uv run basedpyright --project pyproject.toml
+
   nix:
     name: Build and check (${{ matrix.system }})
     runs-on: ${{ matrix.runner }}

MISSED_THINGS.md

@@ -5,21 +5,21 @@
 [X] Fetching download status of all models on start
 [X] Deduplication of tasks in plan_step.
 [X] resolve_allow_patterns should just be wildcard now.
-[X] no mx_barrier in genreate.py mlx_generate at the end.
+[] no mx_barrier in genreate.py mlx_generate at the end.
 [] cache assertion not needed in auto_parallel.py PipelineLastLayer.
-[X] GPTOSS support dropped in auto_parallel.py.
-[X] sharding changed "all-to-sharded" became _all_to_sharded in auto_parallel.py.
-[X] same as above with "sharded-to-all" became _sharded_to_all in auto_parallel.py.
-[X] Dropped support for Ministral3Model, DeepseekV32Model, Glm4MoeModel, Qwen3NextModel, GptOssMode in auto_parallel.py.
+[] GPTOSS support dropped in auto_parallel.py.
+[] sharding changed "all-to-sharded" became _all_to_sharded in auto_parallel.py.
+[] same as above with "sharded-to-all" became _sharded_to_all in auto_parallel.py.
+[] Dropped support for Ministral3Model, DeepseekV32Model, Glm4MoeModel, Qwen3NextModel, GptOssMode in auto_parallel.py.
 [] Dropped prefill/decode code in auto_parallel.py and utils_mlx.py.
 [X] KV_CACHE_BITS should be None to disable quantized KV cache.
-[X] Dropped _set_nofile_limit in utils_mlx.py.
-[X] We have group optional in load_mlx_items in utils_mlx.py.
-[X] Dropped add_missing_chat_templates for GptOss in load_mlx_items in utils_mlx.py.
-[X] Dropped model.make_cache in make_kv_cache in utils_mlx.py.
+[] Dropped _set_nofile_limit in utils_mlx.py.
+[] We have group optional in load_mlx_items in utils_mlx.py.
+[] Dropped add_missing_chat_templates for GptOss in load_mlx_items in utils_mlx.py.
+[] Dropped model.make_cache in make_kv_cache in utils_mlx.py.
 [X] We put cache limit back in utils_mlx.py.
-[X] topology.py remove_node removes the connections after checking if node is is in self._node_id_to_rx_id_map. on beta_1 it checks after, so would remove stale connections I guess?
-[X] Missing Glm 4.7 model cards (this isn't ready yet but should be picked up, probably create an issue... the blocker is transforemrs version doesn't support the tokenizer for Glm 4.7. rc-1 does but we can't upgrade as it breaks other things.)
+[] topology.py remove_node removes the connections after checking if node is is in self._node_id_to_rx_id_map. on beta_1 it checks after, so would remove stale connections I guess?
+[] Missing Glm 4.7 model cards (this isn't ready yet but should be picked up, probably create an issue... the blocker is transforemrs version doesn't support the tokenizer for Glm 4.7. rc-1 does but we can't upgrade as it breaks other things.)
 [] try-except in _command_processor only excepts ValueError. This was silently failing leading to un-debuggable errors (we had a KeyError that was happening ). Changed this to catch Exception instead of ValueError. See exo-v2 89ae38405e0052e3c22405daf094b065878aa873 and fb99fea69b5a39017efc90c5dad0072e677455f0.
 [X] In placement.py, place_instance no longer looks at model_meta.supports_tensor and check if this tensor parallel number of nodes is supported by the model's tensor dimensions.
 [X] In placement.py, place_instanec, we no longer have the special case to exclude DeepSeek v3.1 pipeline parallel (it doesn't work).

bench/exo_bench.py

@@ -288,151 +288,6 @@ def resolve_model_short_id(client: ExoClient, model_arg: str) -> tuple[str, str]
     raise ValueError(f"Model not found in /models: {model_arg}")
 
 
-def run_planning_phase(
-    client: ExoClient,
-    full_model_id: str,
-    preview: dict[str, Any],
-    danger_delete: bool,
-    timeout: float,
-    settle_deadline: float | None,
-) -> None:
-    """Check disk space and ensure model is downloaded before benchmarking."""
-    # Get model size from /models
-    models = client.request_json("GET", "/models") or {}
-    model_bytes = 0
-    for m in models.get("data", []):
-        if m.get("hugging_face_id") == full_model_id:
-            model_bytes = m.get("storage_size_megabytes", 0) * 1024 * 1024
-            break
-
-    if not model_bytes:
-        logger.warning(
-            f"Could not determine size for {full_model_id}, skipping disk check"
-        )
-        return
-
-    # Get nodes from preview
-    inner = unwrap_instance(preview["instance"])
-    node_ids = list(inner["shardAssignments"]["nodeToRunner"].keys())
-    runner_to_shard = inner["shardAssignments"]["runnerToShard"]
-
-    state = client.request_json("GET", "/state")
-    downloads = state.get("downloads", {})
-    node_disk = state.get("nodeDisk", {})
-
-    for node_id in node_ids:
-        node_downloads = downloads.get(node_id, [])
-
-        # Check if model already downloaded on this node
-        already_downloaded = any(
-            "DownloadCompleted" in p
-            and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
-                "modelId"
-            ]
-            == full_model_id
-            for p in node_downloads
-        )
-        if already_downloaded:
-            continue
-
-        # Wait for disk info if settle_deadline is set
-        disk_info = node_disk.get(node_id, {})
-        backoff = _SETTLE_INITIAL_BACKOFF_S
-        while not disk_info and settle_deadline and time.monotonic() < settle_deadline:
-            remaining = settle_deadline - time.monotonic()
-            logger.info(
-                f"Waiting for disk info on {node_id} ({remaining:.0f}s remaining)..."
-            )
-            time.sleep(min(backoff, remaining))
-            backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S)
-            state = client.request_json("GET", "/state")
-            node_disk = state.get("nodeDisk", {})
-            disk_info = node_disk.get(node_id, {})
-
-        if not disk_info:
-            logger.warning(f"No disk info for {node_id}, skipping space check")
-            continue
-
-        avail = disk_info.get("available", {}).get("inBytes", 0)
-        if avail >= model_bytes:
-            continue
-
-        if not danger_delete:
-            raise RuntimeError(
-                f"Insufficient disk on {node_id}: need {model_bytes // (1024**3)}GB, "
-                f"have {avail // (1024**3)}GB. Use --danger-delete-downloads to free space."
-            )
-
-        # Delete from smallest to largest
-        completed = [
-            (
-                unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
-                    "modelId"
-                ],
-                p["DownloadCompleted"]["totalBytes"]["inBytes"],
-            )
-            for p in node_downloads
-            if "DownloadCompleted" in p
-        ]
-        for del_model, size in sorted(completed, key=lambda x: x[1]):
-            logger.info(f"Deleting {del_model} from {node_id} ({size // (1024**2)}MB)")
-            client.request_json("DELETE", f"/download/{node_id}/{del_model}")
-            avail += size
-            if avail >= model_bytes:
-                break
-
-        if avail < model_bytes:
-            raise RuntimeError(f"Could not free enough space on {node_id}")
-
-    # Start downloads (idempotent)
-    for node_id in node_ids:
-        runner_id = inner["shardAssignments"]["nodeToRunner"][node_id]
-        shard = runner_to_shard[runner_id]
-        client.request_json(
-            "POST",
-            "/download/start",
-            body={
-                "targetNodeId": node_id,
-                "shardMetadata": shard,
-            },
-        )
-        logger.info(f"Started download on {node_id}")
-
-    # Wait for downloads
-    start = time.time()
-    while time.time() - start < timeout:
-        state = client.request_json("GET", "/state")
-        downloads = state.get("downloads", {})
-        all_done = True
-        for node_id in node_ids:
-            done = any(
-                "DownloadCompleted" in p
-                and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])[
-                    "modelCard"
-                ]["modelId"]
-                == full_model_id
-                for p in downloads.get(node_id, [])
-            )
-            failed = [
-                p["DownloadFailed"]["errorMessage"]
-                for p in downloads.get(node_id, [])
-                if "DownloadFailed" in p
-                and unwrap_instance(p["DownloadFailed"]["shardMetadata"])["modelCard"][
-                    "modelId"
-                ]
-                == full_model_id
-            ]
-            if failed:
-                raise RuntimeError(f"Download failed on {node_id}: {failed[0]}")
-            if not done:
-                all_done = False
-        if all_done:
-            return
-        time.sleep(1)
-
-    raise TimeoutError("Downloads did not complete in time")
-
-
 def placement_filter(instance_meta: str, wanted: str) -> bool:
     s = (instance_meta or "").lower()
     if wanted == "both":
@@ -680,11 +535,6 @@ def main() -> int:
         default=0,
         help="Max seconds to wait for the cluster to produce valid placements (0 = try once).",
     )
-    ap.add_argument(
-        "--danger-delete-downloads",
-        action="store_true",
-        help="Delete existing models from smallest to largest to make room for benchmark model.",
-    )
     args = ap.parse_args()
 
     pp_list = parse_int_list(args.pp)
@@ -719,16 +569,13 @@ def main() -> int:
         logger.error("[exo-bench] tokenizer usable but prompt sizing failed")
         raise
 
-    settle_deadline = (
-        time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
-    )
-
     selected = fetch_and_filter_placements(client, full_model_id, args)
 
-    if not selected and settle_deadline:
+    if not selected and args.settle_timeout > 0:
         backoff = _SETTLE_INITIAL_BACKOFF_S
-        while not selected and time.monotonic() < settle_deadline:
-            remaining = settle_deadline - time.monotonic()
+        deadline = time.monotonic() + args.settle_timeout
+        while not selected and time.monotonic() < deadline:
+            remaining = deadline - time.monotonic()
             logger.warning(
                 f"No valid placements yet (cluster may still be settling). "
                 f"Retrying in {backoff:.1f}s ({remaining:.0f}s remaining)..."
@@ -760,16 +607,6 @@ def main() -> int:
     if args.dry_run:
         return 0
 
-    logger.info("Planning phase: checking downloads...")
-    run_planning_phase(
-        client,
-        full_model_id,
-        selected[0],
-        args.danger_delete_downloads,
-        args.timeout,
-        settle_deadline,
-    )
-
     all_rows: list[dict[str, Any]] = []
 
     for preview in selected:

conftest.py

@@ -1 +0,0 @@
-collect_ignore = ["tests/start_distributed_test.py"]

dashboard/src/lib/components/ChatForm.svelte

@@ -265,7 +265,6 @@
 
   function handleSubmit() {
     if ((!message.trim() && uploadedFiles.length === 0) || loading) return;
-    if (isEditOnlyWithoutImage) return;
 
     const content = message.trim();
     const files = [...uploadedFiles];
@@ -290,11 +289,7 @@
       if (imageFile.preview) {
         editImage(content, imageFile.preview);
       }
-    } else if (
-      currentModel &&
-      modelSupportsTextToImage(currentModel) &&
-      content
-    ) {
+    } else if (isImageModel() && content) {
       // Use image generation for text-to-image models
       generateImage(content);
     } else {

dashboard/src/lib/components/ChatMessages.svelte

@@ -225,7 +225,6 @@
   }
 
   function handleDeleteClick(messageId: string) {
-    if (loading) return;
     deleteConfirmId = messageId;
   }
 
@@ -256,7 +255,7 @@
 </script>
 
 <div class="flex flex-col gap-4 sm:gap-6 {className}">
-  {#each messageList as message, i (message.id)}
+  {#each messageList as message (message.id)}
     <div
       class="group flex {message.role === 'user'
         ? 'justify-end'
@@ -318,11 +317,9 @@
           <!-- Delete confirmation -->
           <div class="bg-red-500/10 border border-red-500/30 rounded-lg p-3">
             <p class="text-xs text-red-400 mb-3">
-              {#if i === messageList.length - 1}
-                Delete this message?
-              {:else}
-                Delete this message and all messages after it?
-              {/if}
+              Delete this message{message.role === "user"
+                ? " and all responses after it"
+                : ""}?
             </p>
             <div class="flex gap-2 justify-end">
               <button
@@ -754,13 +751,8 @@
             <!-- Delete button -->
             <button
               onclick={() => handleDeleteClick(message.id)}
-              disabled={loading}
-              class="p-1.5 transition-colors rounded {loading
-                ? 'text-exo-light-gray/30 cursor-not-allowed'
-                : 'text-exo-light-gray hover:text-red-400 hover:bg-red-500/10 cursor-pointer'}"
-              title={loading
-                ? "Cannot delete while generating"
-                : "Delete message"}
+              class="p-1.5 text-exo-light-gray hover:text-red-400 transition-colors rounded hover:bg-red-500/10 cursor-pointer"
+              title="Delete message"
             >
               <svg
                 class="w-3.5 h-3.5"

dashboard/src/lib/stores/app.svelte.ts

@@ -72,6 +72,8 @@ export interface Instance {
     runnerToShard?: Record<string, unknown>;
     nodeToRunner?: Record<string, string>;
   };
+  draftModel?: string;
+  numDraftTokens?: number;
 }
 
 // Granular node state types from the new state structure

e2e/Dockerfile

@@ -1,58 +0,0 @@
-# Stage 1: Build the dashboard
-FROM node:22-slim AS dashboard
-WORKDIR /app/dashboard
-COPY dashboard/package.json dashboard/package-lock.json ./
-RUN npm ci
-COPY dashboard/ .
-RUN npm run build
-
-# Stage 2: Build and run exo
-FROM python:3.13-slim
-
-# Install system dependencies
-# libblas-dev/liblapack-dev/liblapacke-dev are required by MLX CPU backend on Linux
-RUN apt-get update && apt-get install -y \
-    build-essential \
-    pkg-config \
-    libssl-dev \
-    libblas-dev \
-    liblapack-dev \
-    liblapacke-dev \
-    curl \
-    protobuf-compiler \
-    iptables \
-    && rm -rf /var/lib/apt/lists/*
-
-# Install Rust nightly
-RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain nightly
-ENV PATH="/root/.cargo/bin:${PATH}"
-
-# Wrap g++ with -fpermissive to fix MLX CPU JIT compilation with GCC 14
-# (GCC 14 treats _Float128/_Float32/_Float64 as built-in types, conflicting with MLX-generated code)
-# Must be done BEFORE uv sync so any source builds also get the fix
-RUN mv /usr/bin/g++ /usr/bin/g++.real && \
-    printf '#!/bin/sh\nexec /usr/bin/g++.real -fpermissive "$@"\n' > /usr/bin/g++ && \
-    chmod +x /usr/bin/g++
-
-# Install uv
-COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv
-
-WORKDIR /app
-
-# Copy dependency files first for better layer caching
-COPY pyproject.toml Cargo.toml uv.lock README.md ./
-COPY rust/ ./rust/
-COPY bench/pyproject.toml ./bench/pyproject.toml
-
-# Copy source and resources
-COPY src/ ./src/
-COPY resources/ ./resources/
-
-# Copy built dashboard from stage 1
-COPY --from=dashboard /app/dashboard/build ./dashboard/build/
-
-# Install Python deps and build Rust bindings, then clean up build artifacts
-# to keep the layer small (Rust target/ and cargo registry can be 1-2 GB)
-RUN uv sync && rm -rf /app/rust/target /root/.cargo/registry /root/.cargo/git
-
-CMD [".venv/bin/exo", "-v"]

e2e/conftest.py

@@ -1,195 +0,0 @@
-"""Shared E2E test infrastructure for exo cluster tests."""
-
-import asyncio
-import json
-import os
-import sys
-from pathlib import Path
-from urllib.error import URLError
-from urllib.request import Request, urlopen
-
-E2E_DIR = Path(__file__).parent.resolve()
-TIMEOUT = int(os.environ.get("E2E_TIMEOUT", "120"))
-
-
-class Cluster:
-    """Async wrapper around a docker compose exo cluster."""
-
-    def __init__(self, name: str, overrides: list[str] | None = None):
-        self.name = name
-        self.project = f"e2e-{name}"
-        compose_files = [str(E2E_DIR / "docker-compose.yml")]
-        for path in overrides or []:
-            compose_files.append(str(E2E_DIR / path))
-        self._compose_base = [
-            "docker",
-            "compose",
-            "-p",
-            self.project,
-            *[arg for f in compose_files for arg in ("-f", f)],
-        ]
-
-    async def __aenter__(self):
-        return self
-
-    async def __aexit__(self, *exc):
-        await self.stop()
-
-    async def _run(self, *args: str, check: bool = True) -> str:
-        proc = await asyncio.create_subprocess_exec(
-            *self._compose_base,
-            *args,
-            stdout=asyncio.subprocess.PIPE,
-            stderr=asyncio.subprocess.STDOUT,
-        )
-        stdout, _ = await proc.communicate()
-        output = stdout.decode()
-        if check and proc.returncode != 0:
-            print(output, file=sys.stderr)
-            raise RuntimeError(
-                f"docker compose {' '.join(args)} failed (rc={proc.returncode})"
-            )
-        return output
-
-    async def build(self):
-        # Skip build if the image was pre-built (e.g. in CI with buildx cache)
-        proc = await asyncio.create_subprocess_exec(
-            "docker",
-            "image",
-            "inspect",
-            "exo-e2e:latest",
-            stdout=asyncio.subprocess.DEVNULL,
-            stderr=asyncio.subprocess.DEVNULL,
-        )
-        await proc.wait()
-        if proc.returncode == 0:
-            print("  Using pre-built image (exo-e2e:latest)")
-            return
-        print("  Building images...")
-        await self._run("build", "--quiet")
-
-    async def start(self):
-        print("  Starting cluster...")
-        await self._run("up", "-d")
-
-    async def stop(self):
-        print("  Cleaning up...")
-        await self._run("down", "--timeout", "5", check=False)
-
-    async def logs(self) -> str:
-        return await self._run("logs", check=False)
-
-    async def exec(
-        self, service: str, *cmd: str, check: bool = True
-    ) -> tuple[int, str]:
-        """Run a command inside a running container. Returns (returncode, output)."""
-        proc = await asyncio.create_subprocess_exec(
-            *self._compose_base,
-            "exec",
-            "-T",
-            service,
-            *cmd,
-            stdout=asyncio.subprocess.PIPE,
-            stderr=asyncio.subprocess.STDOUT,
-        )
-        stdout, _ = await proc.communicate()
-        output = stdout.decode()
-        if check and proc.returncode != 0:
-            raise RuntimeError(
-                f"exec {' '.join(cmd)} in {service} failed (rc={proc.returncode})"
-            )
-        return proc.returncode, output
-
-    async def wait_for(self, description: str, check_fn, timeout: int = TIMEOUT):
-        """Poll check_fn every 2s until it returns True or timeout expires."""
-        print(f"  Waiting for {description}...")
-        deadline = asyncio.get_event_loop().time() + timeout
-        while asyncio.get_event_loop().time() < deadline:
-            if await check_fn():
-                print(f"  {description}")
-                return
-            await asyncio.sleep(2)
-        output = await self.logs()
-        print(f"--- cluster logs ---\n{output}\n---", file=sys.stderr)
-        raise TimeoutError(f"Timed out waiting for {description}")
-
-    async def assert_healthy(self):
-        """Verify the cluster formed correctly: nodes started, discovered each other, elected a master, API responds."""
-
-        async def both_nodes_started():
-            log = await self.logs()
-            return log.count("Starting node") >= 2
-
-        async def nodes_discovered():
-            log = await self.logs()
-            return log.count("ConnectionMessageType.Connected") >= 2
-
-        async def master_elected():
-            log = await self.logs()
-            return "demoting self" in log
-
-        async def api_responding():
-            try:
-                with urlopen("http://localhost:52415/v1/models", timeout=3) as resp:
-                    return resp.status == 200
-            except (URLError, OSError):
-                return False
-
-        await self.wait_for("Both nodes started", both_nodes_started)
-        await self.wait_for("Nodes discovered each other", nodes_discovered)
-        await self.wait_for("Master election resolved", master_elected)
-        await self.wait_for("API responding", api_responding)
-
-    async def _api(
-        self, method: str, path: str, body: dict | None = None, timeout: int = 30
-    ) -> dict:
-        """Make an API request to the cluster. Returns parsed JSON."""
-        url = f"http://localhost:52415{path}"
-        data = json.dumps(body).encode() if body else None
-        req = Request(
-            url, data=data, headers={"Content-Type": "application/json"}, method=method
-        )
-        loop = asyncio.get_event_loop()
-        resp_bytes = await loop.run_in_executor(
-            None, lambda: urlopen(req, timeout=timeout).read()
-        )
-        return json.loads(resp_bytes)
-
-    async def place_model(self, model: str, timeout: int = 600):
-        """Place a model instance on the cluster (triggers download) and wait until it's ready."""
-        await self._api("POST", "/place_instance", {"model_id": model})
-
-        async def model_ready():
-            try:
-                resp = await self._api("GET", "/v1/models")
-                return any(m.get("id") == model for m in resp.get("data", []))
-            except Exception:
-                return False
-
-        await self.wait_for(f"Model {model} ready", model_ready, timeout=timeout)
-
-    async def chat(
-        self, model: str, messages: list[dict], timeout: int = 600, **kwargs
-    ) -> dict:
-        """Send a chat completion request. Retries until model is downloaded and inference completes."""
-        body = json.dumps({"model": model, "messages": messages, **kwargs}).encode()
-        deadline = asyncio.get_event_loop().time() + timeout
-        last_error = None
-
-        while asyncio.get_event_loop().time() < deadline:
-            try:
-                req = Request(
-                    "http://localhost:52415/v1/chat/completions",
-                    data=body,
-                    headers={"Content-Type": "application/json"},
-                )
-                loop = asyncio.get_event_loop()
-                resp_bytes = await loop.run_in_executor(
-                    None, lambda r=req: urlopen(r, timeout=300).read()
-                )
-                return json.loads(resp_bytes)
-            except Exception as e:
-                last_error = e
-                await asyncio.sleep(5)
-
-        raise TimeoutError(f"Chat request failed after {timeout}s: {last_error}")

e2e/docker-compose.yml

@@ -1,20 +0,0 @@
-services:
-  exo-node-1:
-    image: exo-e2e:latest
-    build:
-      context: ..
-      dockerfile: e2e/Dockerfile
-    environment:
-      - EXO_LIBP2P_NAMESPACE=docker-e2e
-    command: [".venv/bin/exo", "-v"]
-    ports:
-      - "52415:52415"
-
-  exo-node-2:
-    image: exo-e2e:latest
-    build:
-      context: ..
-      dockerfile: e2e/Dockerfile
-    environment:
-      - EXO_LIBP2P_NAMESPACE=docker-e2e
-    command: [".venv/bin/exo", "-v"]

e2e/run_all.py

@@ -1,83 +0,0 @@
-#!/usr/bin/env python3
-"""Discovers and runs all E2E tests in e2e/test_*.py.
-
-Tests with '# slow' on the first line of their docstring are skipped
-unless --slow is passed or E2E_SLOW=1 is set.
-"""
-
-import os
-import subprocess
-import sys
-from pathlib import Path
-
-E2E_DIR = Path(__file__).parent.resolve()
-
-
-def is_slow(test_file: Path) -> bool:
-    """Check if the test file is marked as slow (has '# slow' in first 3 lines)."""
-    with open(test_file) as f:
-        for line in f:
-            if line.strip().startswith("#"):
-                continue
-            if line.strip().startswith('"""') or line.strip().startswith("'''"):
-                # Read into the docstring
-                for doc_line in f:
-                    if "slow" in doc_line.lower() and doc_line.strip().startswith(
-                        "slow"
-                    ):
-                        return True
-                    if '"""' in doc_line or "'''" in doc_line:
-                        break
-            break
-    return False
-
-
-def main():
-    run_slow = "--slow" in sys.argv or os.environ.get("E2E_SLOW") == "1"
-    if "--update-snapshots" in sys.argv:
-        os.environ["UPDATE_SNAPSHOTS"] = "1"
-    test_files = sorted(E2E_DIR.glob("test_*.py"))
-    if not test_files:
-        print("No test files found")
-        sys.exit(1)
-
-    passed = 0
-    failed = 0
-    skipped = 0
-    failures = []
-
-    for test_file in test_files:
-        name = test_file.stem
-        if is_slow(test_file) and not run_slow:
-            print(f"=== {name} === SKIPPED (slow, use --slow to run)")
-            skipped += 1
-            continue
-
-        print(f"=== {name} ===")
-        result = subprocess.run([sys.executable, str(test_file)])
-        if result.returncode == 0:
-            passed += 1
-        else:
-            # Retry once — Docker networking (mDNS) can be slow on first boot
-            print(f"\n=== {name} === RETRYING (attempt 2/2)")
-            result = subprocess.run([sys.executable, str(test_file)])
-            if result.returncode == 0:
-                passed += 1
-            else:
-                failed += 1
-                failures.append(name)
-        print()
-
-    total = passed + failed + skipped
-    print("================================")
-    print(
-        f"{passed}/{total} tests passed" + (f", {skipped} skipped" if skipped else "")
-    )
-
-    if failed:
-        print(f"Failed: {' '.join(failures)}")
-        sys.exit(1)
-
-
-if __name__ == "__main__":
-    main()

e2e/snapshot.py

@@ -1,78 +0,0 @@
-"""Snapshot testing infrastructure for E2E tests.
-
-Provides deterministic regression testing by comparing inference output
-against committed baseline snapshots. Tests FAIL if no baseline exists —
-baselines must be explicitly generated and committed.
-
-Generate baselines: UPDATE_SNAPSHOTS=1 python3 e2e/run_all.py --slow
-Update after intentional changes: UPDATE_SNAPSHOTS=1 python3 e2e/run_all.py --slow
-
-Snapshots are stored per-architecture (e.g. snapshots/x86_64/, snapshots/arm64/)
-since floating-point results differ between CPU architectures.
-"""
-
-import difflib
-import json
-import os
-import platform
-from pathlib import Path
-
-ARCH = platform.machine()
-SNAPSHOTS_DIR = Path(__file__).parent / "snapshots" / ARCH
-
-
-def assert_snapshot(
-    name: str,
-    content: str,
-    metadata: dict,
-) -> None:
-    """Compare content against a saved snapshot, or create one if missing.
-
-    Args:
-        name: Snapshot identifier (used as filename: snapshots/{arch}/{name}.json).
-        content: The actual inference output to compare.
-        metadata: Additional context stored alongside content (model, seed, etc.).
-                  Not used for comparison -- purely documentary.
-
-    Raises:
-        AssertionError: If content doesn't match the saved snapshot.
-
-    Environment:
-        UPDATE_SNAPSHOTS=1: Overwrite existing snapshot with actual content.
-    """
-    snapshot_file = SNAPSHOTS_DIR / f"{name}.json"
-    update = os.environ.get("UPDATE_SNAPSHOTS") == "1"
-
-    if update:
-        # Explicitly regenerate snapshot
-        SNAPSHOTS_DIR.mkdir(parents=True, exist_ok=True)
-        snapshot_data = {**metadata, "arch": ARCH, "content": content}
-        snapshot_file.write_text(json.dumps(snapshot_data, indent=2) + "\n")
-        print(f"  Updated snapshot: {ARCH}/{snapshot_file.name}")
-    elif not snapshot_file.exists():
-        raise AssertionError(
-            f"No baseline snapshot for '{name}' on {ARCH}.\n"
-            f"Expected file: {snapshot_file}\n\n"
-            f"Generate baselines with: UPDATE_SNAPSHOTS=1 python3 e2e/run_all.py --slow"
-        )
-    else:
-        snapshot = json.loads(snapshot_file.read_text())
-        expected = snapshot["content"]
-        if content != expected:
-            diff = "\n".join(
-                difflib.unified_diff(
-                    expected.splitlines(),
-                    content.splitlines(),
-                    fromfile=f"expected ({snapshot_file.relative_to(SNAPSHOTS_DIR.parent.parent)})",
-                    tofile="actual",
-                    lineterm="",
-                )
-            )
-            raise AssertionError(
-                f"Snapshot mismatch for '{name}' on {ARCH}!\n\n"
-                f"{diff}\n\n"
-                f"Expected: {expected!r}\n"
-                f"Actual:   {content!r}\n\n"
-                f"To update: UPDATE_SNAPSHOTS=1 python3 e2e/run_all.py --slow"
-            )
-        print(f"  Output matches snapshot ({ARCH}/{snapshot_file.name})")

e2e/test_cluster_formation.py

@@ -1,22 +0,0 @@
-"""Test: Basic cluster formation.
-
-Verifies two nodes discover each other, elect a master, and the API responds.
-"""
-
-import asyncio
-import sys
-
-sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent))
-from conftest import Cluster
-
-
-async def main():
-    async with Cluster("cluster_formation") as cluster:
-        await cluster.build()
-        await cluster.start()
-        await cluster.assert_healthy()
-        print("PASSED: cluster_formation")
-
-
-if __name__ == "__main__":
-    asyncio.run(main())

e2e/test_inference_snapshot.py

@@ -1,61 +0,0 @@
-"""Test: Deterministic inference output (snapshot test).
-
-Sends a chat completion request with a fixed seed,
-then verifies the output matches a known-good snapshot. This ensures
-inference produces consistent results across runs.
-
-Uses MLX CPU backend in Docker on x86 Linux.
-"""
-
-import asyncio
-import sys
-from pathlib import Path
-
-sys.path.insert(0, str(Path(__file__).parent))
-from snapshot import assert_snapshot
-
-from conftest import Cluster
-
-MODEL = "mlx-community/Qwen3-0.6B-4bit"
-SEED = 42
-PROMPT = "What is 2+2? Reply with just the number."
-MAX_TOKENS = 32
-
-
-async def main():
-    async with Cluster("inference_snapshot") as cluster:
-        await cluster.build()
-        await cluster.start()
-        await cluster.assert_healthy()
-
-        print(f"  Launching model {MODEL}...")
-        await cluster.place_model(MODEL)
-
-        print(f"  Sending chat completion (seed={SEED})...")
-        resp = await cluster.chat(
-            model=MODEL,
-            messages=[{"role": "user", "content": PROMPT}],
-            seed=SEED,
-            temperature=0,
-            max_tokens=MAX_TOKENS,
-        )
-
-        content = resp["choices"][0]["message"]["content"]
-        print(f"  Response: {content!r}")
-
-        assert_snapshot(
-            name="inference_snapshot",
-            content=content,
-            metadata={
-                "model": MODEL,
-                "seed": SEED,
-                "prompt": PROMPT,
-                "max_tokens": MAX_TOKENS,
-            },
-        )
-
-        print("PASSED: inference_snapshot")
-
-
-if __name__ == "__main__":
-    asyncio.run(main())

e2e/test_no_internet.py

@@ -1,47 +0,0 @@
-"""Test: Cluster works without internet access.
-
-Verifies exo functions correctly when containers can talk to each other
-but cannot reach the internet. Uses iptables to block all outbound traffic
-except private subnets and multicast (for mDNS discovery).
-"""
-
-import asyncio
-import sys
-
-sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent))
-from conftest import Cluster
-
-
-async def main():
-    async with Cluster(
-        "no_internet",
-        overrides=["tests/no_internet/docker-compose.override.yml"],
-    ) as cluster:
-        await cluster.build()
-        await cluster.start()
-        await cluster.assert_healthy()
-
-        # Verify internet is actually blocked from inside the containers
-        for node in ["exo-node-1", "exo-node-2"]:
-            rc, _ = await cluster.exec(
-                node,
-                "curl",
-                "-sf",
-                "--max-time",
-                "3",
-                "https://huggingface.co",
-                check=False,
-            )
-            assert rc != 0, f"{node} should not be able to reach the internet"
-            print(f"  {node}: internet correctly blocked")
-
-        # Verify exo detected no internet connectivity
-        log = await cluster.logs()
-        assert "Internet connectivity: False" in log, "exo should detect no internet"
-        print("  exo correctly detected no internet connectivity")
-
-        print("PASSED: no_internet")
-
-
-if __name__ == "__main__":
-    asyncio.run(main())

e2e/test_runner_chaos.py

@@ -1,65 +0,0 @@
-"""Test: Runner chaos — abrupt runner death detection.
-slow
-
-Sends a chat completion with the EXO_RUNNER_MUST_DIE trigger, which causes
-the runner process to call os._exit(1) (simulating an OOM kill). Verifies that
-the RunnerSupervisor health check detects the death and the system doesn't hang.
-
-Requires a machine that can run MLX inference at reasonable speed (Apple Silicon).
-Run with: python3 e2e/run_all.py --slow  or  E2E_SLOW=1 python3 e2e/run_all.py
-"""
-
-import asyncio
-import contextlib
-import sys
-from pathlib import Path
-
-sys.path.insert(0, str(Path(__file__).parent))
-from conftest import Cluster
-
-MODEL = "mlx-community/Qwen3-0.6B-4bit"
-
-
-async def main():
-    async with Cluster("runner_chaos") as cluster:
-        await cluster.build()
-        await cluster.start()
-        await cluster.assert_healthy()
-
-        # Place the model so a runner is loaded and ready
-        print(f"  Launching model {MODEL}...")
-        await cluster.place_model(MODEL)
-
-        # Send a chat request with the die trigger.
-        # The runner will call os._exit(1) mid-inference, simulating OOM kill.
-        # The chat request itself will fail — that's expected.
-        print("  Sending EXO_RUNNER_MUST_DIE trigger...")
-        with contextlib.suppress(Exception):
-            await cluster.chat(
-                model=MODEL,
-                messages=[{"role": "user", "content": "EXO RUNNER MUST DIE"}],
-                timeout=60,
-            )
-
-        # Wait for the health check to detect the death and emit RunnerFailed
-        async def health_check_detected():
-            log = await cluster.logs()
-            return "runner process died unexpectedly" in log
-
-        await cluster.wait_for(
-            "Health check detected runner death",
-            health_check_detected,
-            timeout=30,
-        )
-
-        # Verify RunnerFailed was emitted (visible in logs)
-        log = await cluster.logs()
-        assert "runner process died unexpectedly" in log, (
-            f"Expected health check to detect runner death but it didn't.\nLogs:\n{log}"
-        )
-
-        print("PASSED: runner_chaos")
-
-
-if __name__ == "__main__":
-    asyncio.run(main())

e2e/test_snapshot_code_gen.py

@@ -1,60 +0,0 @@
-"""Test: Code generation snapshot.
-slow
-
-Verifies deterministic output for a code generation prompt.
-"""
-
-import asyncio
-import sys
-from pathlib import Path
-
-sys.path.insert(0, str(Path(__file__).parent))
-from snapshot import assert_snapshot
-
-from conftest import Cluster
-
-MODEL = "mlx-community/Qwen3-0.6B-4bit"
-SEED = 42
-PROMPT = (
-    "Write a Python function to reverse a string. Only output the code, no explanation."
-)
-MAX_TOKENS = 64
-
-
-async def main():
-    async with Cluster("snapshot_code_gen") as cluster:
-        await cluster.build()
-        await cluster.start()
-        await cluster.assert_healthy()
-
-        print(f"  Launching model {MODEL}...")
-        await cluster.place_model(MODEL)
-
-        print(f"  Sending chat completion (seed={SEED})...")
-        resp = await cluster.chat(
-            model=MODEL,
-            messages=[{"role": "user", "content": PROMPT}],
-            seed=SEED,
-            temperature=0,
-            max_tokens=MAX_TOKENS,
-        )
-
-        content = resp["choices"][0]["message"]["content"]
-        print(f"  Response: {content!r}")
-
-        assert_snapshot(
-            name="snapshot_code_gen",
-            content=content,
-            metadata={
-                "model": MODEL,
-                "seed": SEED,
-                "prompt": PROMPT,
-                "max_tokens": MAX_TOKENS,
-            },
-        )
-
-        print("PASSED: snapshot_code_gen")
-
-
-if __name__ == "__main__":
-    asyncio.run(main())

e2e/test_snapshot_edge.py

@@ -1,65 +0,0 @@
-"""Test: Edge case snapshots.
-slow
-
-Verifies deterministic output for edge-case prompts: single word input,
-special characters, and unicode.
-"""
-
-import asyncio
-import sys
-from pathlib import Path
-
-sys.path.insert(0, str(Path(__file__).parent))
-from snapshot import assert_snapshot
-
-from conftest import Cluster
-
-MODEL = "mlx-community/Qwen3-0.6B-4bit"
-SEED = 42
-MAX_TOKENS = 32
-
-CASES = [
-    ("edge_single_word", "Hi"),
-    ("edge_special_chars", "What does 2 * (3 + 4) / 7 - 1 equal? Use <math> tags."),
-    ("edge_unicode", "Translate 'hello' to Japanese, Chinese, and Korean."),
-]
-
-
-async def main():
-    async with Cluster("snapshot_edge") as cluster:
-        await cluster.build()
-        await cluster.start()
-        await cluster.assert_healthy()
-
-        print(f"  Launching model {MODEL}...")
-        await cluster.place_model(MODEL)
-
-        for snapshot_name, prompt in CASES:
-            print(f"  [{snapshot_name}] Sending: {prompt!r}")
-            resp = await cluster.chat(
-                model=MODEL,
-                messages=[{"role": "user", "content": prompt}],
-                seed=SEED,
-                temperature=0,
-                max_tokens=MAX_TOKENS,
-            )
-
-            content = resp["choices"][0]["message"]["content"]
-            print(f"  [{snapshot_name}] Response: {content!r}")
-
-            assert_snapshot(
-                name=snapshot_name,
-                content=content,
-                metadata={
-                    "model": MODEL,
-                    "seed": SEED,
-                    "prompt": prompt,
-                    "max_tokens": MAX_TOKENS,
-                },
-            )
-
-        print("PASSED: snapshot_edge")
-
-
-if __name__ == "__main__":
-    asyncio.run(main())

e2e/test_snapshot_long_output.py

@@ -1,58 +0,0 @@
-"""Test: Longer output snapshot.
-slow
-
-Verifies deterministic output with a higher max_tokens (128).
-"""
-
-import asyncio
-import sys
-from pathlib import Path
-
-sys.path.insert(0, str(Path(__file__).parent))
-from snapshot import assert_snapshot
-
-from conftest import Cluster
-
-MODEL = "mlx-community/Qwen3-0.6B-4bit"
-SEED = 42
-PROMPT = "Explain how a binary search algorithm works."
-MAX_TOKENS = 128
-
-
-async def main():
-    async with Cluster("snapshot_long_output") as cluster:
-        await cluster.build()
-        await cluster.start()
-        await cluster.assert_healthy()
-
-        print(f"  Launching model {MODEL}...")
-        await cluster.place_model(MODEL)
-
-        print(f"  Sending chat completion (seed={SEED}, max_tokens={MAX_TOKENS})...")
-        resp = await cluster.chat(
-            model=MODEL,
-            messages=[{"role": "user", "content": PROMPT}],
-            seed=SEED,
-            temperature=0,
-            max_tokens=MAX_TOKENS,
-        )
-
-        content = resp["choices"][0]["message"]["content"]
-        print(f"  Response: {content!r}")
-
-        assert_snapshot(
-            name="snapshot_long_output",
-            content=content,
-            metadata={
-                "model": MODEL,
-                "seed": SEED,
-                "prompt": PROMPT,
-                "max_tokens": MAX_TOKENS,
-            },
-        )
-
-        print("PASSED: snapshot_long_output")
-
-
-if __name__ == "__main__":
-    asyncio.run(main())

e2e/test_snapshot_multi_model.py

@@ -1,73 +0,0 @@
-"""Test: Multi-model snapshot tests.
-slow
-
-Verifies deterministic output across different model architectures to catch
-model-specific regressions. Each model uses its own snapshot file.
-Run with: python3 e2e/run_all.py --slow  or  E2E_SLOW=1 python3 e2e/run_all.py
-"""
-
-import asyncio
-import sys
-from pathlib import Path
-
-sys.path.insert(0, str(Path(__file__).parent))
-from snapshot import assert_snapshot
-
-from conftest import Cluster
-
-SEED = 42
-PROMPT = "What is the capital of France?"
-MAX_TOKENS = 32
-
-MODELS = [
-    "mlx-community/SmolLM2-135M-Instruct",
-    "mlx-community/Llama-3.2-1B-Instruct-4bit",
-    "mlx-community/gemma-2-2b-it-4bit",
-]
-
-
-async def main():
-    async with Cluster("snapshot_multi_model") as cluster:
-        await cluster.build()
-        await cluster.start()
-        await cluster.assert_healthy()
-
-        for model in MODELS:
-            short_name = (
-                model.split("/")[-1].lower().replace("-", "_").replace(".", "_")
-            )
-            snapshot_name = f"snapshot_multi_{short_name}"
-
-            print(f"  Launching model {model}...")
-            await cluster.place_model(model)
-
-            print(f"  Sending chat completion (seed={SEED})...")
-            resp = await cluster.chat(
-                model=model,
-                messages=[{"role": "user", "content": PROMPT}],
-                seed=SEED,
-                temperature=0,
-                max_tokens=MAX_TOKENS,
-            )
-
-            content = resp["choices"][0]["message"]["content"]
-            print(f"  [{short_name}] Response: {content!r}")
-
-            assert_snapshot(
-                name=snapshot_name,
-                content=content,
-                metadata={
-                    "model": model,
-                    "seed": SEED,
-                    "prompt": PROMPT,
-                    "max_tokens": MAX_TOKENS,
-                },
-            )
-
-            print(f"  [{short_name}] PASSED")
-
-        print("PASSED: snapshot_multi_model")
-
-
-if __name__ == "__main__":
-    asyncio.run(main())

e2e/test_snapshot_reasoning.py

@@ -1,58 +0,0 @@
-"""Test: Reasoning/math snapshot.
-slow
-
-Verifies deterministic output for a simple reasoning prompt.
-"""
-
-import asyncio
-import sys
-from pathlib import Path
-
-sys.path.insert(0, str(Path(__file__).parent))
-from snapshot import assert_snapshot
-
-from conftest import Cluster
-
-MODEL = "mlx-community/Qwen3-0.6B-4bit"
-SEED = 42
-PROMPT = "If I have 3 apples and give away 1, how many do I have? Think step by step."
-MAX_TOKENS = 64
-
-
-async def main():
-    async with Cluster("snapshot_reasoning") as cluster:
-        await cluster.build()
-        await cluster.start()
-        await cluster.assert_healthy()
-
-        print(f"  Launching model {MODEL}...")
-        await cluster.place_model(MODEL)
-
-        print(f"  Sending chat completion (seed={SEED})...")
-        resp = await cluster.chat(
-            model=MODEL,
-            messages=[{"role": "user", "content": PROMPT}],
-            seed=SEED,
-            temperature=0,
-            max_tokens=MAX_TOKENS,
-        )
-
-        content = resp["choices"][0]["message"]["content"]
-        print(f"  Response: {content!r}")
-
-        assert_snapshot(
-            name="snapshot_reasoning",
-            content=content,
-            metadata={
-                "model": MODEL,
-                "seed": SEED,
-                "prompt": PROMPT,
-                "max_tokens": MAX_TOKENS,
-            },
-        )
-
-        print("PASSED: snapshot_reasoning")
-
-
-if __name__ == "__main__":
-    asyncio.run(main())

e2e/tests/no_internet/docker-compose.override.yml

@@ -1,32 +0,0 @@
-# Block all outbound internet traffic using iptables while preserving:
-#   - Multicast (224.0.0.0/4) for mDNS peer discovery
-#   - Private subnets (10/8, 172.16/12, 192.168/16) for inter-container communication
-#   - Loopback (127/8)
-# Requires NET_ADMIN capability for iptables.
-services:
-  exo-node-1:
-    cap_add:
-      - NET_ADMIN
-    entrypoint: ["/bin/sh", "-c"]
-    command:
-      - |
-        iptables -A OUTPUT -d 127.0.0.0/8 -j ACCEPT
-        iptables -A OUTPUT -d 10.0.0.0/8 -j ACCEPT
-        iptables -A OUTPUT -d 172.16.0.0/12 -j ACCEPT
-        iptables -A OUTPUT -d 192.168.0.0/16 -j ACCEPT
-        iptables -A OUTPUT -d 224.0.0.0/4 -j ACCEPT
-        iptables -A OUTPUT -j REJECT
-        exec .venv/bin/exo -v
-  exo-node-2:
-    cap_add:
-      - NET_ADMIN
-    entrypoint: ["/bin/sh", "-c"]
-    command:
-      - |
-        iptables -A OUTPUT -d 127.0.0.0/8 -j ACCEPT
-        iptables -A OUTPUT -d 10.0.0.0/8 -j ACCEPT
-        iptables -A OUTPUT -d 172.16.0.0/12 -j ACCEPT
-        iptables -A OUTPUT -d 192.168.0.0/16 -j ACCEPT
-        iptables -A OUTPUT -d 224.0.0.0/4 -j ACCEPT
-        iptables -A OUTPUT -j REJECT
-        exec .venv/bin/exo -v

python/parts.nix

@@ -14,9 +14,7 @@
 
       # Override overlay to inject Nix-built components
       exoOverlay = final: prev: {
-        # Replace workspace exo_pyo3_bindings with Nix-built wheel.
-        # Preserve passthru so mkVirtualEnv can resolve dependency groups.
-        # Copy .pyi stub + py.typed marker so basedpyright can find the types.
+        # Replace workspace exo_pyo3_bindings with Nix-built wheel
         exo-pyo3-bindings = pkgs.stdenv.mkDerivation {
           pname = "exo-pyo3-bindings";
           version = "0.1.0";
@@ -24,12 +22,6 @@
           # Install from pre-built wheel
           nativeBuildInputs = [ final.pyprojectWheelHook ];
           dontStrip = true;
-          passthru = prev.exo-pyo3-bindings.passthru or { };
-          postInstall = ''
-            local siteDir=$out/${final.python.sitePackages}/exo_pyo3_bindings
-            cp ${inputs.self}/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi $siteDir/
-            touch $siteDir/py.typed
-          '';
         };
       };
 
@@ -37,32 +29,17 @@
 
       # Overlay to provide build systems and custom packages
       buildSystemsOverlay = final: prev: {
+        # Use our pure Nix-built MLX with Metal support
+        mlx = self'.packages.mlx;
+
         # mlx-lm is a git dependency that needs setuptools
         mlx-lm = prev.mlx-lm.overrideAttrs (old: {
           nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
             final.setuptools
           ];
         });
-      } // lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin {
-        # Use our pure Nix-built MLX with Metal support (macOS only)
-        mlx = self'.packages.mlx;
       };
 
-      # Additional overlay for Linux-specific fixes (type checking env).
-      # Native wheels have shared lib dependencies we don't need at type-check time.
-      linuxOverlay = final: prev:
-        let
-          ignoreMissing = drv: drv.overrideAttrs { autoPatchelfIgnoreMissingDeps = [ "*" ]; };
-          nvidiaPackages = lib.filterAttrs (name: _: lib.hasPrefix "nvidia-" name) prev;
-        in
-        lib.optionalAttrs pkgs.stdenv.hostPlatform.isLinux (
-          (lib.mapAttrs (_: ignoreMissing) nvidiaPackages) // {
-            mlx = ignoreMissing prev.mlx;
-            torch = ignoreMissing prev.torch;
-            triton = ignoreMissing prev.triton;
-          }
-        );
-
       pythonSet = (pkgs.callPackage inputs.pyproject-nix.build.packages {
         inherit python;
       }).overrideScope (
@@ -71,7 +48,6 @@
           overlay
           exoOverlay
           buildSystemsOverlay
-          linuxOverlay
         ]
       );
       exoVenv = pythonSet.mkVirtualEnv "exo-env" workspace.deps.default;
@@ -142,21 +118,6 @@
           ${pkgs.ruff}/bin/ruff check ${inputs.self}
           touch $out
         '';
-
-        # Hermetic basedpyright type checking
-        typecheck = pkgs.runCommand "typecheck"
-          {
-            nativeBuildInputs = [
-              testVenv
-              pkgs.basedpyright
-            ];
-          }
-          ''
-            cd ${inputs.self}
-            export HOME=$TMPDIR
-            basedpyright --pythonpath ${testVenv}/bin/python
-            touch $out
-          '';
       };
     };
 }

resources/inference_model_cards/mlx-community--SmolLM2-135M-Instruct.toml

@@ -1,12 +0,0 @@
-model_id = "mlx-community/SmolLM2-135M-Instruct"
-n_layers = 30
-hidden_size = 576
-supports_tensor = true
-tasks = ["TextGeneration"]
-family = "llama"
-quantization = "bf16"
-base_model = "SmolLM2 135M"
-capabilities = ["text"]
-
-[storage_size]
-in_bytes = 269060381

resources/inference_model_cards/mlx-community--gemma-2-2b-it-4bit.toml

@@ -1,12 +0,0 @@
-model_id = "mlx-community/gemma-2-2b-it-4bit"
-n_layers = 26
-hidden_size = 2304
-supports_tensor = false
-tasks = ["TextGeneration"]
-family = "gemma2"
-quantization = "4bit"
-base_model = "Gemma 2 2B"
-capabilities = ["text"]
-
-[storage_size]
-in_bytes = 1492755242

src/exo/master/adapters/chat_completions.py

@@ -17,7 +17,6 @@ from exo.shared.types.api import (
     LogprobsContentItem,
     StreamingChoiceResponse,
     ToolCall,
-    Usage,
 )
 from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
 from exo.shared.types.common import CommandId
@@ -126,8 +125,6 @@ async def generate_chat_stream(
     chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
 ) -> AsyncGenerator[str, None]:
     """Generate Chat Completions API streaming events from chunks."""
-    last_usage: Usage | None = None
-
     async for chunk in chunk_stream:
         if isinstance(chunk, ErrorChunk):
             error_response = ErrorResponse(
@@ -141,8 +138,6 @@ async def generate_chat_stream(
             yield "data: [DONE]\n\n"
             return
 
-        last_usage = chunk.usage or last_usage
-
         if isinstance(chunk, ToolCallChunk):
             tool_call_deltas = [
                 ToolCall(
@@ -166,15 +161,12 @@ async def generate_chat_stream(
                         finish_reason="tool_calls",
                     )
                 ],
-                usage=last_usage,
             )
             yield f"data: {tool_response.model_dump_json()}\n\n"
             yield "data: [DONE]\n\n"
             return
 
         chunk_response = chunk_to_response(chunk, command_id)
-        if chunk.finish_reason is not None:
-            chunk_response = chunk_response.model_copy(update={"usage": last_usage})
         yield f"data: {chunk_response.model_dump_json()}\n\n"
 
         if chunk.finish_reason is not None:
@@ -184,9 +176,7 @@ async def generate_chat_stream(
 async def collect_chat_response(
     command_id: CommandId,
     chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
-) -> AsyncGenerator[str]:
-    # This is an AsyncGenerator[str] rather than returning a ChatCompletionReponse because
-    # FastAPI handles the cancellation better but wouldn't auto-serialize for some reason
+) -> ChatCompletionResponse:
     """Collect all token chunks and return a single ChatCompletionResponse."""
     text_parts: list[str] = []
     tool_calls: list[ToolCall] = []
@@ -194,7 +184,6 @@ async def collect_chat_response(
     model: str | None = None
     finish_reason: FinishReason | None = None
     error_message: str | None = None
-    last_usage: Usage | None = None
 
     async for chunk in chunk_stream:
         if isinstance(chunk, ErrorChunk):
@@ -204,8 +193,6 @@ async def collect_chat_response(
         if model is None:
             model = chunk.model
 
-        last_usage = chunk.usage or last_usage
-
         if isinstance(chunk, TokenChunk):
             text_parts.append(chunk.text)
             if chunk.logprob is not None:
@@ -236,7 +223,7 @@ async def collect_chat_response(
     combined_text = "".join(text_parts)
     assert model is not None
 
-    yield ChatCompletionResponse(
+    return ChatCompletionResponse(
         id=command_id,
         created=int(time.time()),
         model=model,
@@ -254,6 +241,4 @@ async def collect_chat_response(
                 finish_reason=finish_reason,
             )
         ],
-        usage=last_usage,
-    ).model_dump_json()
-    return
+    )

src/exo/master/adapters/claude.py

@@ -4,7 +4,7 @@ import json
 from collections.abc import AsyncGenerator
 from typing import Any
 
-from exo.shared.types.api import FinishReason, Usage
+from exo.shared.types.api import FinishReason
 from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
 from exo.shared.types.claude_api import (
     ClaudeContentBlock,
@@ -161,14 +161,12 @@ async def collect_claude_response(
     command_id: CommandId,
     model: str,
     chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
-) -> AsyncGenerator[str]:
-    # This is an AsyncGenerator[str] rather than returning a ChatCompletionReponse because
-    # FastAPI handles the cancellation better but wouldn't auto-serialize for some reason
+) -> ClaudeMessagesResponse:
     """Collect all token chunks and return a single ClaudeMessagesResponse."""
     text_parts: list[str] = []
     tool_use_blocks: list[ClaudeToolUseBlock] = []
     stop_reason: ClaudeStopReason | None = None
-    last_usage: Usage | None = None
+    last_stats = None
     error_message: str | None = None
 
     async for chunk in chunk_stream:
@@ -176,8 +174,6 @@ async def collect_claude_response(
             error_message = chunk.error_message or "Internal server error"
             break
 
-        last_usage = chunk.usage or last_usage
-
         if isinstance(chunk, ToolCallChunk):
             for tool in chunk.tool_calls:
                 tool_use_blocks.append(
@@ -187,10 +183,12 @@ async def collect_claude_response(
                         input=json.loads(tool.arguments),  # pyright: ignore[reportAny]
                     )
                 )
+            last_stats = chunk.stats or last_stats
             stop_reason = "tool_use"
             continue
 
         text_parts.append(chunk.text)
+        last_stats = chunk.stats or last_stats
 
         if chunk.finish_reason is not None:
             stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason)
@@ -210,11 +208,11 @@ async def collect_claude_response(
     if not content:
         content.append(ClaudeTextBlock(text=""))
 
-    # Use actual usage data if available
-    input_tokens = last_usage.prompt_tokens if last_usage else 0
-    output_tokens = last_usage.completion_tokens if last_usage else 0
+    # Use actual usage data from stats if available
+    input_tokens = last_stats.prompt_tokens if last_stats else 0
+    output_tokens = last_stats.generation_tokens if last_stats else 0
 
-    yield ClaudeMessagesResponse(
+    return ClaudeMessagesResponse(
         id=f"msg_{command_id}",
         model=model,
         content=content,
@@ -223,8 +221,7 @@ async def collect_claude_response(
             input_tokens=input_tokens,
             output_tokens=output_tokens,
         ),
-    ).model_dump_json()
-    return
+    )
 
 
 async def generate_claude_stream(
@@ -252,7 +249,7 @@ async def generate_claude_stream(
 
     output_tokens = 0
     stop_reason: ClaudeStopReason | None = None
-    last_usage: Usage | None = None
+    last_stats = None
     next_block_index = 1  # text block is 0, tool blocks start at 1
 
     async for chunk in chunk_stream:
@@ -260,9 +257,8 @@ async def generate_claude_stream(
             # Close text block and bail
             break
 
-        last_usage = chunk.usage or last_usage
-
         if isinstance(chunk, ToolCallChunk):
+            last_stats = chunk.stats or last_stats
             stop_reason = "tool_use"
 
             # Emit tool_use content blocks
@@ -294,6 +290,7 @@ async def generate_claude_stream(
             continue
 
         output_tokens += 1  # Count each chunk as one token
+        last_stats = chunk.stats or last_stats
 
         # content_block_delta
         delta_event = ClaudeContentBlockDeltaEvent(
@@ -305,9 +302,9 @@ async def generate_claude_stream(
         if chunk.finish_reason is not None:
             stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason)
 
-    # Use actual token count from usage if available
-    if last_usage is not None:
-        output_tokens = last_usage.completion_tokens
+    # Use actual token count from stats if available
+    if last_stats is not None:
+        output_tokens = last_stats.generation_tokens
 
     # content_block_stop for text block
     block_stop = ClaudeContentBlockStopEvent(index=0)

src/exo/master/adapters/responses.py

@@ -4,7 +4,6 @@ from collections.abc import AsyncGenerator
 from itertools import count
 from typing import Any
 
-from exo.shared.types.api import Usage
 from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
 from exo.shared.types.common import CommandId
 from exo.shared.types.openai_responses import (
@@ -122,15 +121,13 @@ async def collect_responses_response(
     command_id: CommandId,
     model: str,
     chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
-) -> AsyncGenerator[str]:
-    # This is an AsyncGenerator[str] rather than returning a ChatCompletionReponse because
-    # FastAPI handles the cancellation better but wouldn't auto-serialize for some reason
+) -> ResponsesResponse:
     """Collect all token chunks and return a single ResponsesResponse."""
     response_id = f"resp_{command_id}"
     item_id = f"item_{command_id}"
     accumulated_text = ""
     function_call_items: list[ResponseFunctionCallItem] = []
-    last_usage: Usage | None = None
+    last_stats = None
     error_message: str | None = None
 
     async for chunk in chunk_stream:
@@ -138,8 +135,6 @@ async def collect_responses_response(
             error_message = chunk.error_message or "Internal server error"
             break
 
-        last_usage = chunk.usage or last_usage
-
         if isinstance(chunk, ToolCallChunk):
             for tool in chunk.tool_calls:
                 function_call_items.append(
@@ -150,20 +145,22 @@ async def collect_responses_response(
                         arguments=tool.arguments,
                     )
                 )
+            last_stats = chunk.stats or last_stats
             continue
 
         accumulated_text += chunk.text
+        last_stats = chunk.stats or last_stats
 
     if error_message is not None:
         raise ValueError(error_message)
 
-    # Create usage from usage data if available
+    # Create usage from stats if available
     usage = None
-    if last_usage is not None:
+    if last_stats is not None:
         usage = ResponseUsage(
-            input_tokens=last_usage.prompt_tokens,
-            output_tokens=last_usage.completion_tokens,
-            total_tokens=last_usage.total_tokens,
+            input_tokens=last_stats.prompt_tokens,
+            output_tokens=last_stats.generation_tokens,
+            total_tokens=last_stats.prompt_tokens + last_stats.generation_tokens,
         )
 
     output: list[ResponseItem] = [
@@ -175,15 +172,14 @@ async def collect_responses_response(
     ]
     output.extend(function_call_items)
 
-    yield ResponsesResponse(
+    return ResponsesResponse(
         id=response_id,
         model=model,
         status="completed",
         output=output,
         output_text=accumulated_text,
         usage=usage,
-    ).model_dump_json()
-    return
+    )
 
 
 async def generate_responses_stream(
@@ -239,16 +235,15 @@ async def generate_responses_stream(
 
     accumulated_text = ""
     function_call_items: list[ResponseFunctionCallItem] = []
-    last_usage: Usage | None = None
+    last_stats = None
     next_output_index = 1  # message item is at 0
 
     async for chunk in chunk_stream:
         if isinstance(chunk, ErrorChunk):
             break
 
-        last_usage = chunk.usage or last_usage
-
         if isinstance(chunk, ToolCallChunk):
+            last_stats = chunk.stats or last_stats
             for tool in chunk.tool_calls:
                 fc_id = f"fc_{tool.id}"
                 call_id = f"call_{tool.id}"
@@ -307,6 +302,7 @@ async def generate_responses_stream(
             continue
 
         accumulated_text += chunk.text
+        last_stats = chunk.stats or last_stats
 
         # response.output_text.delta
         delta_event = ResponseTextDeltaEvent(
@@ -350,13 +346,13 @@ async def generate_responses_stream(
     )
     yield f"event: response.output_item.done\ndata: {item_done.model_dump_json()}\n\n"
 
-    # Create usage from usage data if available
+    # Create usage from stats if available
     usage = None
-    if last_usage is not None:
+    if last_stats is not None:
         usage = ResponseUsage(
-            input_tokens=last_usage.prompt_tokens,
-            output_tokens=last_usage.completion_tokens,
-            total_tokens=last_usage.total_tokens,
+            input_tokens=last_stats.prompt_tokens,
+            output_tokens=last_stats.generation_tokens,
+            total_tokens=last_stats.prompt_tokens + last_stats.generation_tokens,
         )
 
     # response.completed

src/exo/master/api.py

@@ -125,7 +125,6 @@ from exo.shared.types.commands import (
     PlaceInstance,
     SendInputChunk,
     StartDownload,
-    TaskCancelled,
     TaskFinished,
     TextGeneration,
 )
@@ -541,14 +540,16 @@ class API:
                         break
 
         except anyio.get_cancelled_exc_class():
-            command = TaskCancelled(cancelled_command_id=command_id)
-            with anyio.CancelScope(shield=True):
-                await self.command_sender.send(
-                    ForwarderCommand(origin=self.node_id, command=command)
-                )
+            # TODO: TaskCancelled
+            """
+            self.command_sender.send_nowait(
+                ForwarderCommand(origin=self.node_id, command=command)
+            )
+            """
             raise
         finally:
-            await self._send(TaskFinished(finished_command_id=command_id))
+            command = TaskFinished(finished_command_id=command_id)
+            await self._send(command)
             if command_id in self._text_generation_queues:
                 del self._text_generation_queues[command_id]
 
@@ -643,14 +644,11 @@ class API:
                     "X-Accel-Buffering": "no",
                 },
             )
-        else:
-            return StreamingResponse(
-                collect_chat_response(
-                    command.command_id,
-                    self._token_chunk_stream(command.command_id),
-                ),
-                media_type="application/json",
-            )
+
+        return await collect_chat_response(
+            command.command_id,
+            self._token_chunk_stream(command.command_id),
+        )
 
     async def bench_chat_completions(
         self, payload: BenchChatCompletionRequest
@@ -666,7 +664,8 @@ class API:
         command = TextGeneration(task_params=task_params)
         await self._send(command)
 
-        return await self._collect_text_generation_with_stats(command.command_id)
+        response = await self._collect_text_generation_with_stats(command.command_id)
+        return response
 
     async def _resolve_and_validate_text_model(self, model_id: ModelId) -> ModelId:
         """Validate a text model exists and return the resolved model ID.
@@ -884,11 +883,6 @@ class API:
                         del image_metadata[key]
 
         except anyio.get_cancelled_exc_class():
-            command = TaskCancelled(cancelled_command_id=command_id)
-            with anyio.CancelScope(shield=True):
-                await self.command_sender.send(
-                    ForwarderCommand(origin=self.node_id, command=command)
-                )
             raise
         finally:
             await self._send(TaskFinished(finished_command_id=command_id))
@@ -970,11 +964,6 @@ class API:
 
             return (images, stats if capture_stats else None)
         except anyio.get_cancelled_exc_class():
-            command = TaskCancelled(cancelled_command_id=command_id)
-            with anyio.CancelScope(shield=True):
-                await self.command_sender.send(
-                    ForwarderCommand(origin=self.node_id, command=command)
-                )
             raise
         finally:
             await self._send(TaskFinished(finished_command_id=command_id))
@@ -1232,15 +1221,12 @@ class API:
                     "X-Accel-Buffering": "no",
                 },
             )
-        else:
-            return StreamingResponse(
-                collect_claude_response(
-                    command.command_id,
-                    payload.model,
-                    self._token_chunk_stream(command.command_id),
-                ),
-                media_type="application/json",
-            )
+
+        return await collect_claude_response(
+            command.command_id,
+            payload.model,
+            self._token_chunk_stream(command.command_id),
+        )
 
     async def openai_responses(
         self, payload: ResponsesRequest
@@ -1268,15 +1254,11 @@ class API:
                 },
             )
 
-        else:
-            return StreamingResponse(
-                collect_responses_response(
-                    command.command_id,
-                    payload.model,
-                    self._token_chunk_stream(command.command_id),
-                ),
-                media_type="application/json",
-            )
+        return await collect_responses_response(
+            command.command_id,
+            payload.model,
+            self._token_chunk_stream(command.command_id),
+        )
 
     def _calculate_total_available_memory(self) -> Memory:
         """Calculate total available memory across all nodes in bytes."""

src/exo/master/main.py

@@ -24,7 +24,7 @@ from exo.shared.types.commands import (
     PlaceInstance,
     RequestEventLog,
     SendInputChunk,
-    TaskCancelled,
+    SetInstanceDraftModel,
     TaskFinished,
     TestCommand,
     TextGeneration,
@@ -36,11 +36,11 @@ from exo.shared.types.events import (
     IndexedEvent,
     InputChunkReceived,
     InstanceDeleted,
+    InstanceDraftModelUpdated,
     NodeGatheredInfo,
     NodeTimedOut,
     TaskCreated,
     TaskDeleted,
-    TaskStatusUpdated,
     TraceEventData,
     TracesCollected,
     TracesMerged,
@@ -281,7 +281,7 @@ class Master:
                         case DeleteInstance():
                             placement = delete_instance(command, self.state.instances)
                             transition_events = get_transition_events(
-                                self.state.instances, placement, self.state.tasks
+                                self.state.instances, placement
                             )
                             for cmd in cancel_unnecessary_downloads(
                                 placement, self.state.downloads
@@ -301,7 +301,7 @@ class Master:
                                 self.state.node_network,
                             )
                             transition_events = get_transition_events(
-                                self.state.instances, placement, self.state.tasks
+                                self.state.instances, placement
                             )
                             generated_events.extend(transition_events)
                         case CreateInstance():
@@ -311,7 +311,7 @@ class Master:
                                 self.state.instances,
                             )
                             transition_events = get_transition_events(
-                                self.state.instances, placement, self.state.tasks
+                                self.state.instances, placement
                             )
                             generated_events.extend(transition_events)
                         case SendInputChunk(chunk=chunk):
@@ -321,18 +321,14 @@ class Master:
                                     chunk=chunk,
                                 )
                             )
-                        case TaskCancelled():
-                            if (
-                                task_id := self.command_task_mapping.get(
-                                    command.cancelled_command_id
-                                )
-                            ) is not None:
-                                generated_events.append(
-                                    TaskStatusUpdated(
-                                        task_status=TaskStatus.Cancelled,
-                                        task_id=task_id,
-                                    )
+                        case SetInstanceDraftModel():
+                            generated_events.append(
+                                InstanceDraftModelUpdated(
+                                    instance_id=command.instance_id,
+                                    draft_model=command.draft_model,
+                                    num_draft_tokens=command.num_draft_tokens,
                                 )
+                            )
                         case TaskFinished():
                             generated_events.append(
                                 TaskDeleted(
@@ -341,9 +337,10 @@ class Master:
                                     ]
                                 )
                             )
-                            self.command_task_mapping.pop(
-                                command.finished_command_id, None
-                            )
+                            if command.finished_command_id in self.command_task_mapping:
+                                del self.command_task_mapping[
+                                    command.finished_command_id
+                                ]
                         case RequestEventLog():
                             # We should just be able to send everything, since other buffers will ignore old messages
                             # rate limit to 1000 at a time

src/exo/master/placement.py

@@ -22,15 +22,9 @@ from exo.shared.types.commands import (
     PlaceInstance,
 )
 from exo.shared.types.common import NodeId
-from exo.shared.types.events import (
-    Event,
-    InstanceCreated,
-    InstanceDeleted,
-    TaskStatusUpdated,
-)
+from exo.shared.types.events import Event, InstanceCreated, InstanceDeleted
 from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
-from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import (
     DownloadOngoing,
     DownloadProgress,
@@ -159,6 +153,8 @@ def place_instance(
                 shard_assignments=shard_assignments,
                 jaccl_devices=mlx_jaccl_devices,
                 jaccl_coordinators=mlx_jaccl_coordinators,
+                draft_model=command.draft_model,
+                num_draft_tokens=command.num_draft_tokens,
             )
         case InstanceMeta.MlxRing:
             ephemeral_port = random_ephemeral_port()
@@ -173,6 +169,8 @@ def place_instance(
                 shard_assignments=shard_assignments,
                 hosts_by_node=hosts_by_node,
                 ephemeral_port=ephemeral_port,
+                draft_model=command.draft_model,
+                num_draft_tokens=command.num_draft_tokens,
             )
 
     return target_instances
@@ -192,7 +190,6 @@ def delete_instance(
 def get_transition_events(
     current_instances: Mapping[InstanceId, Instance],
     target_instances: Mapping[InstanceId, Instance],
-    tasks: Mapping[TaskId, Task],
 ) -> Sequence[Event]:
     events: list[Event] = []
 
@@ -208,18 +205,6 @@ def get_transition_events(
     # find instances to delete
     for instance_id in current_instances:
         if instance_id not in target_instances:
-            for task in tasks.values():
-                if task.instance_id == instance_id and task.task_status in [
-                    TaskStatus.Pending,
-                    TaskStatus.Running,
-                ]:
-                    events.append(
-                        TaskStatusUpdated(
-                            task_status=TaskStatus.Cancelled,
-                            task_id=task.task_id,
-                        )
-                    )
-
             events.append(
                 InstanceDeleted(
                     instance_id=instance_id,

src/exo/master/tests/test_claude_tool_use.py

@@ -4,11 +4,7 @@ import json
 from collections.abc import AsyncGenerator
 from typing import Any, cast
 
-from exo.master.adapters.claude import (
-    ClaudeMessagesResponse,
-    collect_claude_response,
-    generate_claude_stream,
-)
+from exo.master.adapters.claude import collect_claude_response, generate_claude_stream
 from exo.shared.types.api import ToolCallItem
 from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
 from exo.shared.types.common import CommandId, ModelId
@@ -21,18 +17,6 @@ async def _chunks_to_stream(
         yield chunk
 
 
-async def _collect_response(
-    command_id: CommandId,
-    model: str,
-    chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
-) -> ClaudeMessagesResponse:
-    """Helper to consume the async generator and parse the JSON response."""
-    parts: list[str] = []
-    async for part in collect_claude_response(command_id, model, chunk_stream):
-        parts.append(part)
-    return ClaudeMessagesResponse.model_validate_json("".join(parts))
-
-
 MODEL = ModelId("test-model")
 COMMAND_ID = CommandId("cmd_test123")
 
@@ -63,7 +47,7 @@ class TestCollectClaudeResponseToolUse:
                 ],
             ),
         ]
-        response = await _collect_response(
+        response = await collect_claude_response(
             COMMAND_ID, "test-model", _chunks_to_stream(chunks)
         )
 
@@ -93,7 +77,7 @@ class TestCollectClaudeResponseToolUse:
                 ],
             ),
         ]
-        response = await _collect_response(
+        response = await collect_claude_response(
             COMMAND_ID, "test-model", _chunks_to_stream(chunks)
         )
 
@@ -118,7 +102,7 @@ class TestCollectClaudeResponseToolUse:
                 ],
             ),
         ]
-        response = await _collect_response(
+        response = await collect_claude_response(
             COMMAND_ID, "test-model", _chunks_to_stream(chunks)
         )
 
@@ -132,7 +116,7 @@ class TestCollectClaudeResponseToolUse:
 
     async def test_no_content_produces_empty_text_block(self):
         chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = []
-        response = await _collect_response(
+        response = await collect_claude_response(
             COMMAND_ID, "test-model", _chunks_to_stream(chunks)
         )
         assert len(response.content) == 1

src/exo/master/tests/test_placement.py

@@ -239,7 +239,7 @@ def test_get_transition_events_no_change(instance: Instance):
     target_instances = {instance_id: instance}
 
     # act
-    events = get_transition_events(current_instances, target_instances, {})
+    events = get_transition_events(current_instances, target_instances)
 
     # assert
     assert len(events) == 0
@@ -252,7 +252,7 @@ def test_get_transition_events_create_instance(instance: Instance):
     target_instances: dict[InstanceId, Instance] = {instance_id: instance}
 
     # act
-    events = get_transition_events(current_instances, target_instances, {})
+    events = get_transition_events(current_instances, target_instances)
 
     # assert
     assert len(events) == 1
@@ -266,7 +266,7 @@ def test_get_transition_events_delete_instance(instance: Instance):
     target_instances: dict[InstanceId, Instance] = {}
 
     # act
-    events = get_transition_events(current_instances, target_instances, {})
+    events = get_transition_events(current_instances, target_instances)
 
     # assert
     assert len(events) == 1

src/exo/shared/apply.py

@@ -12,6 +12,7 @@ from exo.shared.types.events import (
     InputChunkReceived,
     InstanceCreated,
     InstanceDeleted,
+    InstanceDraftModelUpdated,
     NodeDownloadProgress,
     NodeGatheredInfo,
     NodeTimedOut,
@@ -72,6 +73,8 @@ def event_apply(event: Event, state: State) -> State:
             return apply_instance_created(event, state)
         case InstanceDeleted():
             return apply_instance_deleted(event, state)
+        case InstanceDraftModelUpdated():
+            return apply_instance_draft_model_updated(event, state)
         case NodeTimedOut():
             return apply_node_timed_out(event, state)
         case NodeDownloadProgress():
@@ -190,6 +193,25 @@ def apply_instance_deleted(event: InstanceDeleted, state: State) -> State:
     return state.model_copy(update={"instances": new_instances})
 
 
+def apply_instance_draft_model_updated(
+    event: InstanceDraftModelUpdated, state: State
+) -> State:
+    if event.instance_id not in state.instances:
+        return state
+    instance = state.instances[event.instance_id]
+    updated_instance = instance.model_copy(
+        update={
+            "draft_model": event.draft_model,
+            "num_draft_tokens": event.num_draft_tokens,
+        }
+    )
+    new_instances: Mapping[InstanceId, Instance] = {
+        **state.instances,
+        event.instance_id: updated_instance,
+    }
+    return state.model_copy(update={"instances": new_instances})
+
+
 def apply_runner_status_updated(event: RunnerStatusUpdated, state: State) -> State:
     new_runners: Mapping[RunnerId, RunnerStatus] = {
         **state.runners,

src/exo/shared/types/api.py

@@ -3,7 +3,8 @@ from collections.abc import Generator
 from typing import Annotated, Any, Literal
 from uuid import uuid4
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, field_validator
+from pydantic_core import PydanticUseDefault
 
 from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.common import CommandId, NodeId
@@ -227,6 +228,13 @@ class PlaceInstanceParams(BaseModel):
     instance_meta: InstanceMeta = InstanceMeta.MlxRing
     min_nodes: int = 1
 
+    @field_validator("sharding", "instance_meta", mode="plain")
+    @classmethod
+    def use_default(cls, v: object):
+        if not v or not isinstance(v, (Sharding, InstanceMeta)):
+            raise PydanticUseDefault()
+        return v
+
 
 class CreateInstanceParams(BaseModel):
     instance: Instance

src/exo/shared/types/commands.py

@@ -38,6 +38,8 @@ class PlaceInstance(BaseCommand):
     sharding: Sharding
     instance_meta: InstanceMeta
     min_nodes: int
+    draft_model: ModelId | None = None
+    num_draft_tokens: int = 4
 
 
 class CreateInstance(BaseCommand):
@@ -48,10 +50,6 @@ class DeleteInstance(BaseCommand):
     instance_id: InstanceId
 
 
-class TaskCancelled(BaseCommand):
-    cancelled_command_id: CommandId
-
-
 class TaskFinished(BaseCommand):
     finished_command_id: CommandId
 
@@ -76,6 +74,14 @@ class DeleteDownload(BaseCommand):
     model_id: ModelId
 
 
+class SetInstanceDraftModel(BaseCommand):
+    """Set or update the draft model for an existing instance."""
+
+    instance_id: InstanceId
+    draft_model: ModelId | None  # None to disable speculative decoding
+    num_draft_tokens: int = 4
+
+
 class CancelDownload(BaseCommand):
     target_node_id: NodeId
     model_id: ModelId
@@ -93,7 +99,7 @@ Command = (
     | PlaceInstance
     | CreateInstance
     | DeleteInstance
-    | TaskCancelled
+    | SetInstanceDraftModel
     | TaskFinished
     | SendInputChunk
 )

src/exo/shared/types/events.py

@@ -5,7 +5,7 @@ from pydantic import Field
 
 from exo.shared.topology import Connection
 from exo.shared.types.chunks import GenerationChunk, InputImageChunk
-from exo.shared.types.common import CommandId, Id, NodeId, SessionId
+from exo.shared.types.common import CommandId, Id, ModelId, NodeId, SessionId
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
@@ -68,6 +68,14 @@ class InstanceDeleted(BaseEvent):
     instance_id: InstanceId
 
 
+class InstanceDraftModelUpdated(BaseEvent):
+    """Draft model updated on an existing instance."""
+
+    instance_id: InstanceId
+    draft_model: ModelId | None
+    num_draft_tokens: int
+
+
 class RunnerStatusUpdated(BaseEvent):
     runner_id: RunnerId
     runner_status: RunnerStatus
@@ -141,6 +149,7 @@ Event = (
     | TaskAcknowledged
     | InstanceCreated
     | InstanceDeleted
+    | InstanceDraftModelUpdated
     | RunnerStatusUpdated
     | RunnerDeleted
     | NodeTimedOut

src/exo/shared/types/tasks.py

@@ -24,7 +24,6 @@ class TaskStatus(str, Enum):
     Complete = "Complete"
     TimedOut = "TimedOut"
     Failed = "Failed"
-    Cancelled = "Cancelled"
 
 
 class BaseTask(TaggedModel):
@@ -41,6 +40,12 @@ class DownloadModel(BaseTask):  # emitted by Worker
     shard_metadata: ShardMetadata
 
 
+class DownloadDraftModel(BaseTask):  # emitted by Worker
+    """Download a draft model for speculative decoding (rank 0 only)."""
+
+    model_id: str  # HuggingFace model ID
+
+
 class LoadModel(BaseTask):  # emitted by Worker
     pass
 
@@ -61,11 +66,6 @@ class TextGeneration(BaseTask):  # emitted by Master
     error_message: str | None = Field(default=None)
 
 
-class CancelTask(BaseTask):
-    cancelled_task_id: TaskId
-    runner_id: RunnerId
-
-
 class ImageGeneration(BaseTask):  # emitted by Master
     command_id: CommandId
     task_params: ImageGenerationTaskParams
@@ -86,15 +86,23 @@ class Shutdown(BaseTask):  # emitted by Worker
     runner_id: RunnerId
 
 
+class SetDraftModel(BaseTask):  # emitted by Worker
+    """Load or clear a draft model on an already-running instance."""
+
+    model_id: str | None  # HuggingFace model ID, or None to clear
+    num_draft_tokens: int = 4
+
+
 Task = (
     CreateRunner
     | DownloadModel
+    | DownloadDraftModel
     | ConnectToGroup
     | LoadModel
     | StartWarmup
     | TextGeneration
-    | CancelTask
     | ImageGeneration
     | ImageEdits
     | Shutdown
+    | SetDraftModel
 )

src/exo/shared/types/worker/instances.py

@@ -2,7 +2,7 @@ from enum import Enum
 
 from pydantic import model_validator
 
-from exo.shared.types.common import Host, Id, NodeId
+from exo.shared.types.common import Host, Id, ModelId, NodeId
 from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardMetadata
 from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
 
@@ -19,6 +19,8 @@ class InstanceMeta(str, Enum):
 class BaseInstance(TaggedModel):
     instance_id: InstanceId
     shard_assignments: ShardAssignments
+    draft_model: ModelId | None = None  # For speculative decoding (rank 0 only)
+    num_draft_tokens: int = 4  # Tokens to draft per iteration (when draft_model is set)
 
     def shard(self, runner_id: RunnerId) -> ShardMetadata | None:
         return self.shard_assignments.runner_to_shard.get(runner_id, None)

src/exo/shared/types/worker/runner_response.py

@@ -62,7 +62,6 @@ class PartialImageResponse(BaseRunnerResponse):
 class ToolCallResponse(BaseRunnerResponse):
     tool_calls: list[ToolCallItem]
     usage: Usage | None
-    stats: GenerationStats | None = None
 
 
 class FinishedResponse(BaseRunnerResponse):

src/exo/utils/channels.py

@@ -125,9 +125,7 @@ class MpSender[T]:
             self._state.buffer.put(item, block=True)
 
     async def send_async(self, item: T) -> None:
-        await to_thread.run_sync(
-            self.send, item, limiter=CapacityLimiter(1), abandon_on_cancel=True
-        )
+        await to_thread.run_sync(self.send, item, limiter=CapacityLimiter(1))
 
     def close(self) -> None:
         if not self._state.closed.is_set():

src/exo/worker/engines/mlx/generator/generate.py

@@ -393,11 +393,10 @@ def mlx_generate(
                     f"Model generated unexpected finish_reason: {out.finish_reason}"
                 )
 
-            total_prompt_tokens = len(all_prompt_tokens)
             usage = Usage(
-                prompt_tokens=total_prompt_tokens,
+                prompt_tokens=int(out.prompt_tokens),
                 completion_tokens=completion_tokens,
-                total_tokens=total_prompt_tokens + completion_tokens,
+                total_tokens=int(out.prompt_tokens) + completion_tokens,
                 prompt_tokens_details=PromptTokensDetails(
                     cached_tokens=prefix_hit_length
                 ),

src/exo/worker/engines/mlx/utils_mlx.py

@@ -64,6 +64,8 @@ from exo.worker.runner.bootstrap import logger
 Group = mx.distributed.Group
 
 
+# TODO: Test this
+#  ALSO https://github.com/exo-explore/exo/pull/233#discussion_r2549683673
 def get_weights_size(model_shard_meta: ShardMetadata) -> Memory:
     return Memory.from_float_kb(
         (model_shard_meta.end_layer - model_shard_meta.start_layer)
@@ -81,6 +83,30 @@ class ModelLoadingTimeoutError(Exception):
     pass
 
 
+def mx_barrier(group: Group | None = None):
+    mx.eval(
+        mx.distributed.all_sum(
+            mx.array(1.0),
+            stream=mx.default_stream(mx.Device(mx.cpu)),
+            group=group,
+        )
+    )
+
+
+def broadcast_from_zero(value: int, group: Group | None = None):
+    if group is None:
+        return value
+
+    if group.rank() == 0:
+        a = mx.array([value], dtype=mx.int32)
+    else:
+        a = mx.array([0], dtype=mx.int32)
+
+    m = mx.distributed.all_sum(a, stream=mx.Device(mx.DeviceType.cpu), group=group)
+    mx.eval(m)
+    return int(m.item())
+
+
 class HostList(RootModel[list[str]]):
     @classmethod
     def from_hosts(cls, hosts: list[Host]) -> "HostList":
@@ -197,6 +223,27 @@ def load_mlx_items(
     return cast(Model, model), tokenizer
 
 
+def load_draft_model(model_id: ModelId) -> nn.Module:
+    """Load a draft model for speculative decoding (rank 0 only).
+
+    Draft models are small models (typically 0.5B-2B parameters) used to
+    generate candidate tokens quickly, which are then verified by the main
+    model in a single forward pass.
+
+    Assumes the model has already been downloaded by the worker.
+
+    Args:
+        model_id: HuggingFace model ID for the draft model
+
+    Returns:
+        The loaded draft model
+    """
+    model_path = build_model_path(model_id)
+    draft_model, _ = load_model(model_path, strict=True)
+    logger.info(f"Loaded draft model from {model_path}")
+    return draft_model
+
+
 def shard_and_load(
     shard_metadata: ShardMetadata,
     group: Group,
@@ -565,23 +612,3 @@ def mlx_cleanup(
     import gc
 
     gc.collect()
-
-
-def mx_any(bool_: bool, group: Group | None) -> bool:
-    if group is None:
-        return bool_
-    num_true = mx.distributed.all_sum(
-        mx.array(bool_), group=group, stream=mx.default_stream(mx.Device(mx.cpu))
-    )
-    mx.eval(num_true)
-    return num_true.item() > 0
-
-
-def mx_barrier(group: Group | None):
-    if group is None:
-        return
-    mx.eval(
-        mx.distributed.all_sum(
-            mx.array(1.0), group=group, stream=mx.default_stream(mx.Device(mx.cpu))
-        )
-    )

src/exo/worker/main.py

@@ -33,7 +33,6 @@ from exo.shared.types.events import (
 from exo.shared.types.multiaddr import Multiaddr
 from exo.shared.types.state import State
 from exo.shared.types.tasks import (
-    CancelTask,
     CreateRunner,
     DownloadModel,
     ImageEdits,
@@ -225,22 +224,15 @@ class Worker:
                         )
                     )
                 case Shutdown(runner_id=runner_id):
-                    runner = self.runners.pop(runner_id)
                     try:
                         with fail_after(3):
-                            await runner.start_task(task)
+                            await self.runners.pop(runner_id).start_task(task)
                     except TimeoutError:
                         await self.event_sender.send(
                             TaskStatusUpdated(
                                 task_id=task.task_id, task_status=TaskStatus.TimedOut
                             )
                         )
-                    finally:
-                        runner.shutdown()
-                case CancelTask(
-                    cancelled_task_id=cancelled_task_id, runner_id=runner_id
-                ):
-                    await self.runners[runner_id].cancel_task(cancelled_task_id)
                 case ImageEdits() if task.task_params.total_input_chunks > 0:
                     # Assemble image from chunks and inject into task
                     cmd_id = task.command_id
@@ -278,18 +270,18 @@ class Worker:
                         del self.input_chunk_buffer[cmd_id]
                     if cmd_id in self.input_chunk_counts:
                         del self.input_chunk_counts[cmd_id]
-                    await self._start_runner_task(modified_task)
+                    await self.runners[self._task_to_runner_id(task)].start_task(
+                        modified_task
+                    )
                 case task:
-                    await self._start_runner_task(task)
+                    await self.runners[self._task_to_runner_id(task)].start_task(task)
 
     def shutdown(self):
         self._tg.cancel_scope.cancel()
 
-    async def _start_runner_task(self, task: Task):
-        if (instance := self.state.instances.get(task.instance_id)) is not None:
-            await self.runners[
-                instance.shard_assignments.node_to_runner[self.node_id]
-            ].start_task(task)
+    def _task_to_runner_id(self, task: Task):
+        instance = self.state.instances[task.instance_id]
+        return instance.shard_assignments.node_to_runner[self.node_id]
 
     async def _nack_request(self, since_idx: int) -> None:
         # We request all events after (and including) the missing index.
@@ -328,6 +320,8 @@ class Worker:
             for event in self.out_for_delivery.copy().values():
                 await self.local_event_sender.send(event)
 
+    ## Op Executors
+
     def _create_supervisor(self, task: CreateRunner) -> RunnerSupervisor:
         """Creates and stores a new AssignedRunner with initial downloading status."""
         runner = RunnerSupervisor.create(

src/exo/worker/plan.py

@@ -4,7 +4,6 @@ from collections.abc import Mapping, Sequence
 
 from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.tasks import (
-    CancelTask,
     ConnectToGroup,
     CreateRunner,
     DownloadModel,
@@ -54,14 +53,13 @@ def plan(
 ) -> Task | None:
     # Python short circuiting OR logic should evaluate these sequentially.
     return (
-        _cancel_tasks(runners, tasks)
-        or _kill_runner(runners, all_runners, instances)
+        _kill_runner(runners, all_runners, instances)
         or _create_runner(node_id, runners, instances)
         or _model_needs_download(node_id, runners, global_download_status)
         or _init_distributed_backend(runners, all_runners)
         or _load_model(runners, all_runners, global_download_status)
         or _ready_to_warmup(runners, all_runners)
-        or _pending_tasks(runners, tasks, all_runners, input_chunk_buffer or {})
+        or _pending_tasks(runners, tasks, all_runners, input_chunk_buffer)
     )
 
 
@@ -272,7 +270,7 @@ def _pending_tasks(
     runners: Mapping[RunnerId, RunnerSupervisor],
     tasks: Mapping[TaskId, Task],
     all_runners: Mapping[RunnerId, RunnerStatus],
-    input_chunk_buffer: Mapping[CommandId, dict[int, str]],
+    input_chunk_buffer: Mapping[CommandId, dict[int, str]] | None = None,
 ) -> Task | None:
     for task in tasks.values():
         # for now, just forward chat completions
@@ -286,7 +284,7 @@ def _pending_tasks(
         if isinstance(task, ImageEdits) and task.task_params.total_input_chunks > 0:
             cmd_id = task.command_id
             expected = task.task_params.total_input_chunks
-            received = len(input_chunk_buffer.get(cmd_id, {}))
+            received = len((input_chunk_buffer or {}).get(cmd_id, {}))
             if received < expected:
                 continue  # Wait for all chunks to arrive
 
@@ -294,33 +292,16 @@ def _pending_tasks(
             if task.instance_id != runner.bound_instance.instance.instance_id:
                 continue
 
-            # the task status _should_ be set to completed by the LAST runner
-            # it is currently set by the first
-            # this is definitely a hack
+            # I have a design point here; this is a state race in disguise as the task status doesn't get updated to completed fast enough
+            # however, realistically the task status should be set to completed by the LAST runner, so this is a true race
+            # the actual solution is somewhat deeper than this bypass - TODO!
             if task.task_id in runner.completed:
                 continue
 
+            # TODO: Check ordering aligns with MLX distributeds expectations.
+
             if isinstance(runner.status, RunnerReady) and all(
                 isinstance(all_runners[global_runner_id], (RunnerReady, RunnerRunning))
                 for global_runner_id in runner.bound_instance.instance.shard_assignments.runner_to_shard
             ):
                 return task
-
-
-def _cancel_tasks(
-    runners: Mapping[RunnerId, RunnerSupervisor],
-    tasks: Mapping[TaskId, Task],
-) -> Task | None:
-    for task in tasks.values():
-        if task.task_status != TaskStatus.Cancelled:
-            continue
-        for runner_id, runner in runners.items():
-            if task.instance_id != runner.bound_instance.instance.instance_id:
-                continue
-            if task.task_id in runner.cancelled:
-                continue
-            return CancelTask(
-                instance_id=task.instance_id,
-                cancelled_task_id=task.task_id,
-                runner_id=runner_id,
-            )

src/exo/worker/runner/bootstrap.py

@@ -1,36 +1,21 @@
-from __future__ import annotations
-
 import os
-import threading
-from multiprocessing.sharedctypes import Synchronized
 
 import loguru
 
 from exo.shared.types.events import Event, RunnerStatusUpdated
-from exo.shared.types.tasks import Task, TaskId
+from exo.shared.types.tasks import Task
 from exo.shared.types.worker.instances import BoundInstance, MlxJacclInstance
 from exo.shared.types.worker.runners import RunnerFailed
 from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender
 
 logger: "loguru.Logger" = loguru.logger
 
-HEARTBEAT_INTERVAL_SECONDS = 0.5
-
-
-def _heartbeat_loop(heartbeat: Synchronized[int], stop: threading.Event) -> None:
-    """Daemon thread that periodically increments the heartbeat counter."""
-    while not stop.is_set():
-        heartbeat.value += 1
-        stop.wait(HEARTBEAT_INTERVAL_SECONDS)
-
 
 def entrypoint(
     bound_instance: BoundInstance,
     event_sender: MpSender[Event],
     task_receiver: MpReceiver[Task],
-    cancel_receiver: MpReceiver[TaskId],
     _logger: "loguru.Logger",
-    heartbeat: Synchronized[int] | None = None,
 ) -> None:
     fast_synch_override = os.environ.get("EXO_FAST_SYNCH")
     if fast_synch_override == "on" or (
@@ -49,22 +34,11 @@ def entrypoint(
 
     logger.info(f"Fast synch flag: {os.environ['MLX_METAL_FAST_SYNCH']}")
 
-    # Start heartbeat thread so the supervisor can detect if we freeze.
-    stop_heartbeat = threading.Event()
-    heartbeat_thread: threading.Thread | None = None
-    if heartbeat is not None:
-        heartbeat_thread = threading.Thread(
-            target=_heartbeat_loop,
-            args=(heartbeat, stop_heartbeat),
-            daemon=True,
-        )
-        heartbeat_thread.start()
-
     # Import main after setting global logger - this lets us just import logger from this module
     try:
         from exo.worker.runner.runner import main
 
-        main(bound_instance, event_sender, task_receiver, cancel_receiver)
+        main(bound_instance, event_sender, task_receiver)
     except ClosedResourceError:
         logger.warning("Runner communication closed unexpectedly")
     except Exception as e:
@@ -78,9 +52,6 @@ def entrypoint(
             )
         )
     finally:
-        stop_heartbeat.set()
-        if heartbeat_thread is not None:
-            heartbeat_thread.join(timeout=1)
         try:
             event_sender.close()
             task_receiver.close()

src/exo/worker/runner/runner.py

@@ -1,7 +1,5 @@
 import base64
 import json
-import math
-import os
 import resource
 import time
 from collections.abc import Generator
@@ -90,7 +88,6 @@ from exo.worker.engines.mlx.utils_mlx import (
     initialize_mlx,
     load_mlx_items,
     mlx_force_oom,
-    mx_any,
 )
 from exo.worker.runner.bootstrap import logger
 
@@ -115,7 +112,6 @@ def main(
     bound_instance: BoundInstance,
     event_sender: MpSender[Event],
     task_receiver: MpReceiver[Task],
-    cancel_receiver: MpReceiver[TaskId],
 ):
     soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
     resource.setrlimit(resource.RLIMIT_NOFILE, (min(max(soft, 2048), hard), hard))
@@ -133,15 +129,11 @@ def main(
         time.sleep(timeout)
 
     setup_start_time = time.time()
-    cancelled_tasks = set[TaskId]()
 
-    # type checker was unhappy with me - splitting these fixed it
-    inference_model: Model | None = None
-    image_model: DistributedImageModel | None = None
+    model: Model | DistributedImageModel | None = None
     tokenizer = None
     group = None
     kv_prefix_cache: KVPrefixCache | None = None
-    check_for_cancel_every: int | None = None
 
     current_status: RunnerStatus = RunnerIdle()
     logger.info("runner created")
@@ -154,7 +146,6 @@ def main(
             if task.task_id in seen:
                 logger.warning("repeat task - potential error")
             seen.add(task.task_id)
-            cancelled_tasks.discard(TaskId("CANCEL_CURRENT_TASK"))
             event_sender.send(
                 TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Running)
             )
@@ -200,7 +191,7 @@ def main(
                         time.sleep(0.5)
 
                     if ModelTask.TextGeneration in shard_metadata.model_card.tasks:
-                        inference_model, tokenizer = load_mlx_items(
+                        model, tokenizer = load_mlx_items(
                             bound_instance, group, on_timeout=on_model_load_timeout
                         )
                         logger.info(
@@ -212,7 +203,7 @@ def main(
                         ModelTask.TextToImage in shard_metadata.model_card.tasks
                         or ModelTask.ImageToImage in shard_metadata.model_card.tasks
                     ):
-                        image_model = initialize_image_model(bound_instance)
+                        model = initialize_image_model(bound_instance)
                     else:
                         raise ValueError(
                             f"Unknown model task(s): {shard_metadata.model_card.tasks}"
@@ -220,6 +211,8 @@ def main(
                     current_status = RunnerLoaded()
                     logger.info("runner loaded")
                 case StartWarmup() if isinstance(current_status, RunnerLoaded):
+                    assert model
+
                     current_status = RunnerWarmingUp()
                     logger.info("runner warming up")
                     event_sender.send(
@@ -231,31 +224,16 @@ def main(
 
                     logger.info(f"warming up inference for instance: {instance}")
                     if ModelTask.TextGeneration in shard_metadata.model_card.tasks:
-                        assert inference_model
+                        assert not isinstance(model, DistributedImageModel)
                         assert tokenizer
 
-                        t = time.monotonic()
                         toks = warmup_inference(
-                            model=inference_model,
+                            model=model,
                             tokenizer=tokenizer,
                             group=group,
+                            # kv_prefix_cache=kv_prefix_cache,  # supply for warmup-time prefix caching
                         )
                         logger.info(f"warmed up by generating {toks} tokens")
-                        check_for_cancel_every = min(
-                            math.ceil(toks / min(time.monotonic() - t, 0.001)), 100
-                        )
-                        if group is not None:
-                            check_for_cancel_every = int(
-                                mx.max(
-                                    mx.distributed.all_gather(
-                                        mx.array([check_for_cancel_every]), group=group
-                                    )
-                                ).item()
-                            )
-
-                        logger.info(
-                            f"runner checking for cancellation every {check_for_cancel_every} tokens"
-                        )
                         logger.info(
                             f"runner initialized in {time.time() - setup_start_time} seconds"
                         )
@@ -263,8 +241,8 @@ def main(
                         ModelTask.TextToImage in shard_metadata.model_card.tasks
                         or ModelTask.ImageToImage in shard_metadata.model_card.tasks
                     ):
-                        assert image_model
-                        image = warmup_image_generator(model=image_model)
+                        assert isinstance(model, DistributedImageModel)
+                        image = warmup_image_generator(model=model)
                         if image is not None:
                             logger.info(f"warmed up by generating {image.size} image")
                         else:
@@ -284,9 +262,9 @@ def main(
                         )
                     )
                     event_sender.send(TaskAcknowledged(task_id=task.task_id))
-                    assert inference_model
+
+                    assert model and not isinstance(model, DistributedImageModel)
                     assert tokenizer
-                    assert check_for_cancel_every
 
                     try:
                         _check_for_debug_prompts(task_params)
@@ -296,7 +274,7 @@ def main(
 
                         # Generate responses using the actual MLX generation
                         mlx_generator = mlx_generate(
-                            model=inference_model,
+                            model=model,
                             tokenizer=tokenizer,
                             task=task_params,
                             prompt=prompt,
@@ -321,11 +299,11 @@ def main(
                             patch_glm_tokenizer(tokenizer)
 
                         # GPT-OSS specific parsing to match other model formats.
-                        elif isinstance(inference_model, GptOssModel):
+                        elif isinstance(model, GptOssModel):
                             mlx_generator = parse_gpt_oss(mlx_generator)
 
                         if tokenizer.has_tool_calling and not isinstance(
-                            inference_model, GptOssModel
+                            model, GptOssModel
                         ):
                             assert tokenizer.tool_call_start
                             assert tokenizer.tool_call_end
@@ -338,18 +316,7 @@ def main(
                             )
 
                         completion_tokens = 0
-                        tokens_since_last_cancel_check = 0
                         for response in mlx_generator:
-                            tokens_since_last_cancel_check += 1
-                            if tokens_since_last_cancel_check >= check_for_cancel_every:
-                                tokens_since_last_cancel_check = 0
-                                cancelled_tasks.update(cancel_receiver.collect())
-                                want_to_cancel = (task.task_id in cancelled_tasks) or (
-                                    TaskId("CANCEL_CURRENT_TASK") in cancelled_tasks
-                                )
-                                if mx_any(want_to_cancel, group):
-                                    break
-
                             match response:
                                 case GenerationResponse():
                                     completion_tokens += 1
@@ -397,7 +364,6 @@ def main(
                                                     tool_calls=response.tool_calls,
                                                     model=shard_metadata.model_card.model_id,
                                                     usage=response.usage,
-                                                    stats=response.stats,
                                                 ),
                                             )
                                         )
@@ -422,7 +388,7 @@ def main(
                 case ImageGeneration(
                     task_params=task_params, command_id=command_id
                 ) if isinstance(current_status, RunnerReady):
-                    assert image_model
+                    assert isinstance(model, DistributedImageModel)
                     logger.info(f"received image generation request: {str(task)[:500]}")
                     current_status = RunnerRunning()
                     logger.info("runner running")
@@ -435,9 +401,7 @@ def main(
 
                     try:
                         image_index = 0
-                        for response in generate_image(
-                            model=image_model, task=task_params
-                        ):
+                        for response in generate_image(model=model, task=task_params):
                             is_primary_output = _is_primary_output_node(shard_metadata)
 
                             if is_primary_output:
@@ -487,7 +451,7 @@ def main(
                 case ImageEdits(task_params=task_params, command_id=command_id) if (
                     isinstance(current_status, RunnerReady)
                 ):
-                    assert image_model
+                    assert isinstance(model, DistributedImageModel)
                     logger.info(f"received image edits request: {str(task)[:500]}")
                     current_status = RunnerRunning()
                     logger.info("runner running")
@@ -500,9 +464,7 @@ def main(
 
                     try:
                         image_index = 0
-                        for response in generate_image(
-                            model=image_model, task=task_params
-                        ):
+                        for response in generate_image(model=model, task=task_params):
                             if _is_primary_output_node(shard_metadata):
                                 match response:
                                     case PartialImageResponse():
@@ -561,20 +523,14 @@ def main(
                     raise ValueError(
                         f"Received {task.__class__.__name__} outside of state machine in {current_status=}"
                     )
-            was_cancelled = (task.task_id in cancelled_tasks) or (
-                TaskId("CANCEL_CURRENT_TASK") in cancelled_tasks
+            event_sender.send(
+                TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Complete)
             )
-            if not was_cancelled:
-                event_sender.send(
-                    TaskStatusUpdated(
-                        task_id=task.task_id, task_status=TaskStatus.Complete
-                    )
-                )
             event_sender.send(
                 RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
             )
             if isinstance(current_status, RunnerShutdown):
-                del inference_model, image_model, tokenizer, group
+                del model, tokenizer, group
                 mx.clear_cache()
                 import gc
 
@@ -808,9 +764,7 @@ def parse_tool_calls(
                     tools = [_validate_single_tool(tool) for tool in parsed]
                 else:
                     tools = [_validate_single_tool(parsed)]
-                yield ToolCallResponse(
-                    tool_calls=tools, usage=response.usage, stats=response.stats
-                )
+                yield ToolCallResponse(tool_calls=tools, usage=response.usage)
 
             except (
                 json.JSONDecodeError,
@@ -841,8 +795,7 @@ def parse_tool_calls(
                     text=tool_call_start + "".join(tool_call_text_parts),
                     token=0,
                     finish_reason=response.finish_reason,
-                    usage=response.usage,
-                    stats=response.stats,
+                    usage=None,
                 )
             continue
         # fallthrough
@@ -1000,7 +953,6 @@ def _validate_single_tool(obj: dict[str, Any]) -> ToolCallItem:
 EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL"
 EXO_RUNNER_MUST_OOM = "EXO RUNNER MUST OOM"
 EXO_RUNNER_MUST_TIMEOUT = "EXO RUNNER MUST TIMEOUT"
-EXO_RUNNER_MUST_DIE = "EXO RUNNER MUST DIE"
 
 
 def _check_for_debug_prompts(task_params: TextGenerationTaskParams) -> None:
@@ -1016,9 +968,6 @@ def _check_for_debug_prompts(task_params: TextGenerationTaskParams) -> None:
     if not prompt:
         return
 
-    if EXO_RUNNER_MUST_DIE in prompt:
-        logger.info("Abrupt process death triggered (simulates OOM kill)")
-        os._exit(1)
     if EXO_RUNNER_MUST_FAIL in prompt:
         logger.info("raising exception")
         raise Exception("Artificial runner exception - for testing purposes only.")

src/exo/worker/runner/runner_supervisor.py

@@ -1,17 +1,12 @@
-from __future__ import annotations
-
 import contextlib
-import multiprocessing
 import signal
 from dataclasses import dataclass, field
 from multiprocessing import Process
-from multiprocessing.sharedctypes import Synchronized
 from typing import Self
 
 import anyio
 from anyio import (
     BrokenResourceError,
-    CancelScope,
     ClosedResourceError,
     to_thread,
 )
@@ -31,7 +26,6 @@ from exo.shared.types.worker.runners import (
     RunnerIdle,
     RunnerLoading,
     RunnerRunning,
-    RunnerShutdown,
     RunnerShuttingDown,
     RunnerStatus,
     RunnerWarmingUp,
@@ -42,8 +36,6 @@ from exo.worker.runner.bootstrap import entrypoint
 
 PREFILL_TIMEOUT_SECONDS = 60
 DECODE_TIMEOUT_SECONDS = 5
-HEALTH_CHECK_INTERVAL_SECONDS = 1
-HEARTBEAT_STALE_CHECKS = 10
 
 
 @dataclass(eq=False)
@@ -55,15 +47,9 @@ class RunnerSupervisor:
     _ev_recv: MpReceiver[Event]
     _task_sender: MpSender[Task]
     _event_sender: Sender[Event]
-    _cancel_sender: MpSender[TaskId]
-    _heartbeat: Synchronized[int]
     status: RunnerStatus = field(default_factory=RunnerIdle, init=False)
     pending: dict[TaskId, anyio.Event] = field(default_factory=dict, init=False)
     completed: set[TaskId] = field(default_factory=set, init=False)
-    cancelled: set[TaskId] = field(default_factory=set, init=False)
-    _death_handled: bool = field(default=False, init=False)
-    _last_heartbeat_value: int = field(default=0, init=False)
-    _heartbeat_stale_count: int = field(default=0, init=False)
 
     @classmethod
     def create(
@@ -74,10 +60,8 @@ class RunnerSupervisor:
         initialize_timeout: float = 400,
     ) -> Self:
         ev_send, ev_recv = mp_channel[Event]()
+        # A task is kind of a runner command
         task_sender, task_recv = mp_channel[Task]()
-        cancel_sender, cancel_recv = mp_channel[TaskId]()
-
-        heartbeat: Synchronized[int] = multiprocessing.Value("Q", 0)
 
         runner_process = Process(
             target=entrypoint,
@@ -85,9 +69,7 @@ class RunnerSupervisor:
                 bound_instance,
                 ev_send,
                 task_recv,
-                cancel_recv,
                 logger,
-                heartbeat,
             ),
             daemon=True,
         )
@@ -101,26 +83,20 @@ class RunnerSupervisor:
             initialize_timeout=initialize_timeout,
             _ev_recv=ev_recv,
             _task_sender=task_sender,
-            _cancel_sender=cancel_sender,
             _event_sender=event_sender,
-            _heartbeat=heartbeat,
         )
 
         return self
 
     async def run(self):
         self.runner_process.start()
-        async with anyio.create_task_group() as tg:
-            tg.start_soon(self._forward_events)
-            tg.start_soon(self._health_check, tg.cancel_scope)
+        await self._forward_events()
 
     def shutdown(self):
         logger.info("Runner supervisor shutting down")
         self._ev_recv.close()
         self._task_sender.close()
         self._event_sender.close()
-        self._cancel_sender.send(TaskId("CANCEL_CURRENT_TASK"))
-        self._cancel_sender.close()
         self.runner_process.join(1)
         if not self.runner_process.is_alive():
             logger.info("Runner process succesfully terminated")
@@ -136,6 +112,14 @@ class RunnerSupervisor:
         logger.critical("Runner process didn't respond to SIGTERM, killing")
         self.runner_process.kill()
 
+        self.runner_process.join(1)
+        if not self.runner_process.is_alive():
+            return
+
+        logger.critical(
+            "Runner process didn't respond to SIGKILL. System resources may have leaked"
+        )
+
     async def start_task(self, task: Task):
         if task.task_id in self.pending:
             logger.warning(
@@ -157,17 +141,6 @@ class RunnerSupervisor:
             return
         await event.wait()
 
-    async def cancel_task(self, task_id: TaskId):
-        if task_id in self.completed:
-            logger.info(f"Unable to cancel {task_id} as it has been completed")
-            return
-        self.cancelled.add(task_id)
-        with anyio.move_on_after(0.5) as scope:
-            await self._cancel_sender.send_async(task_id)
-        if scope.cancel_called:
-            logger.error("RunnerSupervisor cancel pipe blocked")
-            await self._check_runner(TimeoutError("cancel pipe blocked"))
-
     async def _forward_events(self):
         with self._ev_recv as events:
             try:
@@ -195,99 +168,9 @@ class RunnerSupervisor:
                         self.completed.add(event.task_id)
                     await self._event_sender.send(event)
             except (ClosedResourceError, BrokenResourceError) as e:
-                if not self._death_handled:
-                    self._death_handled = True
-                    await self._check_runner(e)
-                    for tid in self.pending:
-                        self.pending[tid].set()
-
-    async def _health_check(self, cancel_scope: CancelScope) -> None:
-        """Periodically check if the runner process is alive and responsive.
-
-        Detects two failure modes:
-        1. Process death (e.g. OOM kill) without cleanly closing the event
-           channel, which would leave _forward_events blocked on queue.get().
-        2. Unresponsive process (e.g. frozen by OS memory pressure, deadlock)
-           detected via a stale heartbeat counter.
-        """
-        while True:
-            await anyio.sleep(HEALTH_CHECK_INTERVAL_SECONDS)
-
-            if not self.runner_process.is_alive():
-                self._handle_process_exit(cancel_scope)
-                return
-
-            # Check heartbeat counter — if it hasn't changed between
-            # consecutive checks, the subprocess may be frozen.
-            current = self._heartbeat.value
-            if current > 0:
-                if current == self._last_heartbeat_value:
-                    self._heartbeat_stale_count += 1
-                    if self._heartbeat_stale_count >= HEARTBEAT_STALE_CHECKS:
-                        logger.error(
-                            f"Health check: runner process unresponsive "
-                            f"(heartbeat stale for {self._heartbeat_stale_count} checks), killing"
-                        )
-                        self._handle_unresponsive(cancel_scope)
-                        return
-                else:
-                    self._heartbeat_stale_count = 0
-                self._last_heartbeat_value = current
-
-    def _handle_process_exit(self, cancel_scope: CancelScope) -> None:
-        """Handle runner process that has exited."""
-        if not self._death_handled:
-            self._death_handled = True
-            if isinstance(
-                self.status, (RunnerShutdown, RunnerShuttingDown, RunnerFailed)
-            ):
-                logger.info("Health check: runner process exited (expected)")
-            else:
-                rc = self.runner_process.exitcode
-                if isinstance(rc, int) and rc < 0:
-                    sig = -rc
-                    try:
-                        cause = f"signal={sig} ({signal.strsignal(sig)})"
-                    except Exception:
-                        cause = f"signal={sig}"
-                else:
-                    cause = f"exitcode={rc}"
-
-                logger.error(
-                    f"Health check: runner process died unexpectedly ({cause})"
-                )
-                self._event_sender.send_nowait(
-                    RunnerStatusUpdated(
-                        runner_id=self.bound_instance.bound_runner_id,
-                        runner_status=RunnerFailed(
-                            error_message=f"Terminated ({cause})"
-                        ),
-                    )
-                )
-                self.shutdown()
-
-            for tid in self.pending:
-                self.pending[tid].set()
-
-        cancel_scope.cancel()
-
-    def _handle_unresponsive(self, cancel_scope: CancelScope) -> None:
-        """Handle runner process that is alive but unresponsive."""
-        if not self._death_handled:
-            self._death_handled = True
-            self._event_sender.send_nowait(
-                RunnerStatusUpdated(
-                    runner_id=self.bound_instance.bound_runner_id,
-                    runner_status=RunnerFailed(
-                        error_message="Runner process unresponsive (heartbeat timeout)"
-                    ),
-                )
-            )
-            for tid in self.pending:
-                self.pending[tid].set()
-            self.shutdown()
-
-        cancel_scope.cancel()
+                await self._check_runner(e)
+                for tid in self.pending:
+                    self.pending[tid].set()
 
     def __del__(self) -> None:
         if self.runner_process.is_alive():

src/exo/worker/tests/unittests/test_runner/test_event_ordering.py

@@ -1,9 +1,7 @@
 # Check tasks are complete before runner is ever ready.
-import unittest.mock
 from collections.abc import Iterable
 from typing import Callable
 
-import mlx.core as mx
 import pytest
 
 import exo.worker.runner.runner as mlx_runner
@@ -21,7 +19,6 @@ from exo.shared.types.tasks import (
     Shutdown,
     StartWarmup,
     Task,
-    TaskId,
     TaskStatus,
     TextGeneration,
 )
@@ -116,7 +113,6 @@ def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(mlx_runner, "load_mlx_items", make_nothin((1, MockTokenizer)))
     monkeypatch.setattr(mlx_runner, "warmup_inference", make_nothin(1))
     monkeypatch.setattr(mlx_runner, "_check_for_debug_prompts", nothin)
-    monkeypatch.setattr(mlx_runner, "mx_any", make_nothin(False))
     # Mock apply_chat_template since we're using a fake tokenizer (integer 1).
     # Returns a prompt without thinking tag so detect_thinking_prompt_suffix returns None.
     monkeypatch.setattr(mlx_runner, "apply_chat_template", make_nothin("test prompt"))
@@ -167,7 +163,6 @@ def _run(tasks: Iterable[Task]):
     )
 
     task_sender, task_receiver = mp_channel[Task]()
-    _cancel_sender, cancel_receiver = mp_channel[TaskId]()
     event_sender = EventCollector()
 
     with task_sender:
@@ -178,16 +173,8 @@ def _run(tasks: Iterable[Task]):
         # this is some c++ nonsense
         task_receiver.close = nothin
         task_receiver.join = nothin
-        with unittest.mock.patch(
-            "exo.worker.runner.runner.mx.distributed.all_gather",
-            make_nothin(mx.array([1])),
-        ):
-            mlx_runner.main(
-                bound_instance,
-                event_sender,  # pyright: ignore[reportArgumentType]
-                task_receiver,
-                cancel_receiver,
-            )
+
+        mlx_runner.main(bound_instance, event_sender, task_receiver)  # type: ignore[arg-type]
 
         return event_sender.events
 

src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py

@@ -1,204 +1 @@
-from __future__ import annotations
-
-import multiprocessing
-import os
-import signal as signal_module
-from collections.abc import Callable
-from multiprocessing.sharedctypes import Synchronized
-from typing import Any
-
-import anyio
-
-from exo.shared.types.events import Event, RunnerStatusUpdated
-from exo.shared.types.tasks import Task, TaskId
-from exo.shared.types.worker.runners import RunnerFailed, RunnerIdle, RunnerShutdown
-from exo.utils.channels import Receiver, Sender, channel, mp_channel
-from exo.worker.runner.runner_supervisor import (
-    HEALTH_CHECK_INTERVAL_SECONDS,
-    HEARTBEAT_STALE_CHECKS,
-    RunnerSupervisor,
-)
-
-from ...constants import (
-    INSTANCE_1_ID,
-    MODEL_A_ID,
-    NODE_A,
-    RUNNER_1_ID,
-)
-from ..conftest import get_bound_mlx_ring_instance
-
-
-def _die_immediately() -> None:
-    """Subprocess target that exits with a non-zero code."""
-    os._exit(1)
-
-
-def _die_with_signal() -> None:
-    """Subprocess target that kills itself with SIGKILL (simulates OOM)."""
-    os.kill(os.getpid(), signal_module.SIGKILL)
-
-
-def _exit_cleanly() -> None:
-    """Subprocess target that exits with code 0."""
-    os._exit(0)
-
-
-def _hang_forever() -> None:
-    """Subprocess target that hangs without updating heartbeat (simulates freeze)."""
-    import time
-
-    # Write one heartbeat so the supervisor starts tracking, then stop.
-    time.sleep(100000)
-
-
-def _build_supervisor(
-    event_sender: Sender[Event],
-    target: Callable[..., Any],
-) -> RunnerSupervisor:
-    """Build a RunnerSupervisor with a custom subprocess target.
-
-    Uses a clone of event_sender (matching real Worker behavior) so that
-    closing the supervisor's copy doesn't close the test's receiver.
-    """
-    bound_instance = get_bound_mlx_ring_instance(
-        instance_id=INSTANCE_1_ID,
-        model_id=MODEL_A_ID,
-        runner_id=RUNNER_1_ID,
-        node_id=NODE_A,
-    )
-
-    _ev_send, ev_recv = mp_channel[Event]()
-    task_sender, _task_recv = mp_channel[Task]()
-    cancel_sender, _cancel_recv = mp_channel[TaskId]()
-    runner_process = multiprocessing.Process(target=target, daemon=True)
-    heartbeat: Synchronized[int] = multiprocessing.Value("Q", 0)
-
-    return RunnerSupervisor(
-        bound_instance=bound_instance,
-        shard_metadata=bound_instance.bound_shard,
-        runner_process=runner_process,
-        initialize_timeout=10,
-        _ev_recv=ev_recv,
-        _task_sender=task_sender,
-        _cancel_sender=cancel_sender,
-        _event_sender=event_sender.clone(),
-        _heartbeat=heartbeat,
-    )
-
-
-def _collect_failed_events(
-    event_receiver: Receiver[Event],
-) -> list[RunnerFailed]:
-    """Drain the receiver and return all RunnerFailed statuses."""
-    out: list[RunnerFailed] = []
-    while True:
-        try:
-            event = event_receiver.receive_nowait()
-        except Exception:
-            break
-        if isinstance(event, RunnerStatusUpdated) and isinstance(
-            event.runner_status, RunnerFailed
-        ):
-            out.append(event.runner_status)
-    return out
-
-
-async def test_health_check_detects_dead_process():
-    """When the runner process dies with a non-zero exit code, the health check
-    should emit a RunnerFailed event and run() should return."""
-    event_sender, event_receiver = channel[Event]()
-    supervisor = _build_supervisor(event_sender, _die_immediately)
-
-    with anyio.fail_after(HEALTH_CHECK_INTERVAL_SECONDS + 5):
-        await supervisor.run()
-
-    failures = _collect_failed_events(event_receiver)
-    assert len(failures) == 1
-    assert failures[0].error_message is not None
-    assert "exitcode=1" in failures[0].error_message
-
-
-async def test_health_check_detects_signal_death():
-    """When the runner process is killed by a signal (e.g. OOM -> SIGKILL),
-    the health check should report the signal in the failure message."""
-    event_sender, event_receiver = channel[Event]()
-    supervisor = _build_supervisor(event_sender, _die_with_signal)
-
-    with anyio.fail_after(HEALTH_CHECK_INTERVAL_SECONDS + 5):
-        await supervisor.run()
-
-    failures = _collect_failed_events(event_receiver)
-    assert len(failures) == 1
-    assert failures[0].error_message is not None
-    assert "signal=9" in failures[0].error_message
-
-
-async def test_health_check_releases_pending_tasks():
-    """When the runner dies, any pending start_task() waiters should be unblocked."""
-    event_sender, _event_receiver = channel[Event]()
-    supervisor = _build_supervisor(event_sender, _die_immediately)
-
-    # Register a pending waiter as if start_task() was waiting for acknowledgement
-    task_event = anyio.Event()
-    tid = TaskId("pending-task")
-    supervisor.pending[tid] = task_event
-
-    with anyio.fail_after(HEALTH_CHECK_INTERVAL_SECONDS + 5):
-        await supervisor.run()
-
-    assert task_event.is_set()
-
-
-async def test_clean_exit_no_failure_when_shutdown_status():
-    """When the runner was in RunnerShutdown status and exits with code 0,
-    no RunnerFailed event should be emitted."""
-    event_sender, event_receiver = channel[Event]()
-    supervisor = _build_supervisor(event_sender, _exit_cleanly)
-
-    # Simulate that the runner had already reported shutdown via events
-    supervisor.status = RunnerShutdown()
-
-    with anyio.fail_after(HEALTH_CHECK_INTERVAL_SECONDS + 5):
-        await supervisor.run()
-
-    failures = _collect_failed_events(event_receiver)
-    assert len(failures) == 0
-
-
-async def test_unexpected_exit_code_zero_emits_failure():
-    """When the runner exits with code 0 but was NOT in a shutdown state,
-    this is unexpected and should still emit RunnerFailed."""
-    event_sender, event_receiver = channel[Event]()
-    supervisor = _build_supervisor(event_sender, _exit_cleanly)
-
-    assert isinstance(supervisor.status, RunnerIdle)
-
-    with anyio.fail_after(HEALTH_CHECK_INTERVAL_SECONDS + 5):
-        await supervisor.run()
-
-    failures = _collect_failed_events(event_receiver)
-    assert len(failures) == 1
-    assert failures[0].error_message is not None
-    assert "exitcode=0" in failures[0].error_message
-
-
-async def test_heartbeat_timeout_detects_unresponsive_process():
-    """When the runner process is alive but its heartbeat goes stale,
-    the health check should kill it and emit RunnerFailed."""
-    event_sender, event_receiver = channel[Event]()
-    supervisor = _build_supervisor(event_sender, _hang_forever)
-
-    # Pre-seed the heartbeat counter with a non-zero value and set the
-    # supervisor's last-seen value to match so it appears stale immediately.
-    # Set stale count to HEARTBEAT_STALE_CHECKS - 1 so a single check triggers.
-    supervisor._heartbeat.value = 42  # pyright: ignore[reportPrivateUsage]
-    supervisor._last_heartbeat_value = 42  # pyright: ignore[reportPrivateUsage]
-    supervisor._heartbeat_stale_count = HEARTBEAT_STALE_CHECKS - 1  # pyright: ignore[reportPrivateUsage]
-
-    with anyio.fail_after(HEALTH_CHECK_INTERVAL_SECONDS + 5):
-        await supervisor.run()
-
-    failures = _collect_failed_events(event_receiver)
-    assert len(failures) == 1
-    assert failures[0].error_message is not None
-    assert "unresponsive" in failures[0].error_message.lower()
+# TODO: