fix(dashboard): only suggest RDMA for TB5 hardware, not all Thunderbolt

Parse linkSpeed from system_profiler to distinguish TB5 (>40 Gb/s) from
TB3/TB4 (40 Gb/s). Fixes #1636.

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

.mlx_typings/mlx_lm/generate.pyi

@@ -73,9 +73,11 @@ class GenerationResponse:
     finish_reason: Optional[str] = ...
 
 def maybe_quantize_kv_cache(
-    prompt_cache, quantized_kv_start, kv_group_size, kv_bits
-):  # -> None:
-    ...
+    prompt_cache: Any,
+    quantized_kv_start: int | None,
+    kv_group_size: int | None,
+    kv_bits: int | None,
+) -> None: ...
 def generate_step(
     prompt: mx.array,
     model: nn.Module,

.mlx_typings/mlx_lm/models/cache.pyi

@@ -16,7 +16,7 @@ class Cache(Protocol):
         self, keys: mx.array, values: mx.array
     ) -> tuple[mx.array, mx.array]: ...
     @property
-    def state(self) -> tuple[mx.array, mx.array]: ...
+    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
     @state.setter
     def state(self, v) -> None: ...
 
@@ -92,13 +92,14 @@ class _BaseCache(Cache):
     values: mx.array
     offset: int
     @property
-    def state(self) -> tuple[mx.array, mx.array]: ...
+    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
     @state.setter
     def state(self, v) -> None: ...
     @property
     def meta_state(self) -> Literal[""]: ...
     @meta_state.setter
     def meta_state(self, v) -> None: ...
+    def trim(self, n: int) -> int: ...
     def is_trimmable(self) -> Literal[False]: ...
     @classmethod
     def from_state(cls, state, meta_state) -> Self: ...
@@ -114,15 +115,13 @@ class ConcatenateKVCache(_BaseCache):
     def update_and_fetch(self, keys, values):  # -> tuple[Any | array, Any | array]:
         ...
     @property
-    def state(self):  # -> tuple[Any | array | None, Any | array | None]:
-        ...
+    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
     @state.setter
     def state(self, v):  # -> None:
         ...
     def is_trimmable(self):  # -> Literal[True]:
         ...
-    def trim(self, n):  # -> int:
-        ...
+    def trim(self, n: int) -> int: ...
     def make_mask(self, *args, **kwargs):  # -> array | Literal['causal'] | None:
         ...
 
@@ -132,10 +131,7 @@ class QuantizedKVCache(_BaseCache):
     def update_and_fetch(self, keys, values):  # -> Any:
         ...
     @property
-    def state(
-        self,
-    ):  # -> tuple[Any | tuple[array, array, array] | None, Any | tuple[array, array, array] | None] | Any:
-        ...
+    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
     @state.setter
     def state(self, v):  # -> None:
         ...
@@ -147,8 +143,7 @@ class QuantizedKVCache(_BaseCache):
         ...
     def is_trimmable(self):  # -> Literal[True]:
         ...
-    def trim(self, n):  # -> int:
-        ...
+    def trim(self, n: int) -> int: ...
     def make_mask(self, *args, **kwargs):  # -> array | Literal['causal'] | None:
         ...
 
@@ -160,18 +155,18 @@ class KVCache(_BaseCache):
     @property
     def state(
         self,
-    ) -> tuple[array, array]: ...
+    ) -> tuple[mx.array | None, mx.array | None]: ...
     @state.setter
     def state(self, v) -> None: ...
     def is_trimmable(self):  # -> Literal[True]:
         ...
-    def trim(self, n):  # -> int:
-        ...
+    def trim(self, n: int) -> int: ...
     def to_quantized(
         self, group_size: int = ..., bits: int = ...
     ) -> QuantizedKVCache: ...
-    def make_mask(self, *args, **kwargs):  # -> array | Literal['causal'] | None:
-        ...
+    def make_mask(
+        self, *args: Any, **kwargs: Any
+    ) -> mx.array | Literal["causal"] | None: ...
 
 class RotatingKVCache(_BaseCache):
     step = ...
@@ -183,8 +178,7 @@ class RotatingKVCache(_BaseCache):
     @property
     def state(
         self,
-    ):  # -> tuple[Any | array, Any | array] | tuple[Any | array | None, Any | array | None]:
-        ...
+    ) -> tuple[mx.array | None, mx.array | None]: ...
     @state.setter
     def state(self, v):  # -> None:
         ...
@@ -196,8 +190,7 @@ class RotatingKVCache(_BaseCache):
         ...
     def is_trimmable(self):  # -> bool:
         ...
-    def trim(self, n):  # -> int:
-        ...
+    def trim(self, n: int) -> int: ...
     def to_quantized(
         self, group_size: int = ..., bits: int = ...
     ) -> QuantizedKVCache: ...
@@ -212,8 +205,7 @@ class ArraysCache(_BaseCache):
         ...
     def __getitem__(self, idx): ...
     @property
-    def state(self):  # -> list[Any | array] | list[array]:
-        ...
+    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
     @state.setter
     def state(self, v):  # -> None:
         ...
@@ -227,8 +219,7 @@ class ArraysCache(_BaseCache):
         In-place extend this cache with the other cache.
         """
 
-    def make_mask(self, N: int):  # -> array | None:
-        ...
+    def make_mask(self, N: int) -> mx.array | None: ...
 
 class MambaCache(ArraysCache):
     def __init__(self, left_padding: Optional[List[int]] = ...) -> None: ...
@@ -239,8 +230,7 @@ class ChunkedKVCache(KVCache):
         ...
     def update_and_fetch(self, keys, values):  # -> tuple[array, array]:
         ...
-    def trim(self, n):  # -> int:
-        ...
+    def trim(self, n: int) -> int: ...
     @property
     def meta_state(self):  # -> tuple[str, ...]:
         ...
@@ -253,10 +243,9 @@ class CacheList(_BaseCache):
     def __getitem__(self, idx): ...
     def is_trimmable(self):  # -> bool:
         ...
-    def trim(self, n): ...
+    def trim(self, n: int) -> int: ...
     @property
-    def state(self):  # -> list[Any]:
-        ...
+    def state(self) -> list[tuple[mx.array | None, mx.array | None]]: ...
     @state.setter
     def state(self, v):  # -> None:
         ...

.mlx_typings/mlx_lm/models/qwen3_5.pyi

@@ -0,0 +1,153 @@
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .cache import ArraysCache, KVCache
+from .qwen3_next import (
+    Qwen3NextAttention as Attention,
+    Qwen3NextMLP as MLP,
+    Qwen3NextRMSNormGated as RMSNormGated,
+    Qwen3NextSparseMoeBlock,
+)
+
+SparseMoeBlock = Qwen3NextSparseMoeBlock
+from .switch_layers import SwitchGLU
+
+@dataclass
+class TextModelArgs:
+    model_type: str
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    rms_norm_eps: float
+    vocab_size: int
+    num_key_value_heads: int
+    max_position_embeddings: int
+    linear_num_value_heads: int
+    linear_num_key_heads: int
+    linear_key_head_dim: int
+    linear_value_head_dim: int
+    linear_conv_kernel_dim: int
+    tie_word_embeddings: bool
+    attention_bias: bool
+    head_dim: Optional[int]
+    full_attention_interval: int
+    num_experts: int
+    num_experts_per_tok: int
+    decoder_sparse_step: int
+    shared_expert_intermediate_size: int
+    moe_intermediate_size: int
+    norm_topk_prob: bool
+    rope_parameters: Optional[dict[str, Any]]
+    partial_rotary_factor: float
+    rope_theta: float
+    rope_scaling: Optional[dict[str, Any]]
+
+    @classmethod
+    def from_dict(cls, params: dict[str, Any]) -> TextModelArgs: ...
+    def __post_init__(self) -> None: ...
+
+class GatedDeltaNet(nn.Module):
+    hidden_size: int
+    num_v_heads: int
+    num_k_heads: int
+    head_k_dim: int
+    head_v_dim: int
+    key_dim: int
+    value_dim: int
+    conv_kernel_size: int
+    conv_dim: int
+    conv1d: nn.Conv1d
+    in_proj_qkv: nn.Linear
+    in_proj_z: nn.Linear
+    in_proj_b: nn.Linear
+    in_proj_a: nn.Linear
+    dt_bias: mx.array
+    A_log: mx.array
+    norm: RMSNormGated
+    out_proj: nn.Linear
+
+    def __init__(self, config: TextModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class DecoderLayer(nn.Module):
+    is_linear: bool
+    linear_attn: GatedDeltaNet
+    self_attn: Attention
+    input_layernorm: nn.RMSNorm
+    post_attention_layernorm: nn.RMSNorm
+    mlp: MLP | SparseMoeBlock
+
+    def __init__(self, args: TextModelArgs, layer_idx: int) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class Qwen3_5TextModel(nn.Module):
+    embed_tokens: nn.Embedding
+    layers: list[DecoderLayer]
+    norm: nn.RMSNorm
+    ssm_idx: int
+    fa_idx: int
+
+    def __init__(self, args: TextModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+        input_embeddings: Optional[mx.array] = None,
+    ) -> mx.array: ...
+
+class TextModel(nn.Module):
+    args: TextModelArgs
+    model_type: str
+    model: Qwen3_5TextModel
+    lm_head: nn.Linear
+
+    def __init__(self, args: TextModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+        input_embeddings: Optional[mx.array] = None,
+    ) -> mx.array: ...
+    @property
+    def layers(self) -> list[DecoderLayer]: ...
+    def make_cache(self) -> list[ArraysCache | KVCache]: ...
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
+
+@dataclass
+class ModelArgs:
+    model_type: str
+    text_config: dict[str, Any]
+
+    @classmethod
+    def from_dict(cls, params: dict[str, Any]) -> ModelArgs: ...
+
+class Model(nn.Module):
+    args: ModelArgs
+    model_type: str
+    language_model: TextModel
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+        input_embeddings: Optional[mx.array] = None,
+    ) -> mx.array: ...
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
+    @property
+    def layers(self) -> list[DecoderLayer]: ...
+    def make_cache(self) -> list[ArraysCache | KVCache]: ...

.mlx_typings/mlx_lm/models/qwen3_5_moe.pyi

@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .cache import ArraysCache, KVCache
+from .qwen3_5 import DecoderLayer, Model as Qwen3_5Model, TextModel
+
+@dataclass
+class ModelArgs:
+    model_type: str
+    text_config: dict[str, Any]
+
+    @classmethod
+    def from_dict(cls, params: dict[str, Any]) -> ModelArgs: ...
+
+class Model(Qwen3_5Model):
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...

.mlx_typings/mlx_lm/models/qwen3_next.pyi

@@ -7,6 +7,15 @@ import mlx.nn as nn
 
 from .switch_layers import SwitchGLU
 
+class Qwen3NextRMSNormGated(nn.Module):
+    eps: float
+    weight: mx.array
+
+    def __init__(self, hidden_size: int, eps: float = ...) -> None: ...
+    def __call__(
+        self, hidden_states: mx.array, gate: mx.array | None = None
+    ) -> mx.array: ...
+
 class Qwen3NextMLP(nn.Module):
     gate_proj: nn.Linear
     down_proj: nn.Linear

Cargo.lock

@@ -216,6 +216,28 @@ dependencies = [
  "windows-sys 0.61.2",
 ]
 
+[[package]]
+name = "async-stream"
+version = "0.3.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476"
+dependencies = [
+ "async-stream-impl",
+ "futures-core",
+ "pin-project-lite",
+]
+
+[[package]]
+name = "async-stream-impl"
+version = "0.3.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.111",
+]
+
 [[package]]
 name = "async-trait"
 version = "0.1.89"
@@ -2759,6 +2781,7 @@ dependencies = [
 name = "networking"
 version = "0.0.1"
 dependencies = [
+ "async-stream",
  "delegate",
  "either",
  "extend",
@@ -2767,6 +2790,7 @@ dependencies = [
  "keccak-const",
  "libp2p",
  "log",
+ "pin-project",
  "tokio",
  "tracing-subscriber",
  "util",

Cargo.toml

@@ -1,10 +1,6 @@
 [workspace]
 resolver = "3"
-members = [
-    "rust/networking",
-    "rust/exo_pyo3_bindings",
-    "rust/util",
-]
+members = ["rust/networking", "rust/exo_pyo3_bindings", "rust/util"]
 
 [workspace.package]
 version = "0.0.1"
@@ -34,6 +30,7 @@ delegate = "0.13"
 keccak-const = "0.2"
 
 # Async dependencies
+async-stream = "0.3"
 tokio = "1.46"
 futures-lite = "2.6.1"
 futures-timer = "3.0"

bench/bench.toml

@@ -2,6 +2,4 @@
 #
 # Lists the suite files to include. Each file defines benchmarks
 # with shared constraints, topology, and default args.
-include = [
-    "single-m3-ultra.toml",
-]
+include = ["single-m3-ultra.toml"]

bench/pyproject.toml

@@ -4,13 +4,13 @@ version = "0.1.0"
 description = "Benchmarking tool for exo distributed inference"
 requires-python = ">=3.13"
 dependencies = [
-    "httpx>=0.27.0",
-    "loguru>=0.7.3",
-    "transformers>=5.0.0",
-    "huggingface-hub>=0.33.4",
-    "tiktoken>=0.12.0",
-    "jinja2>=3.1.0",
-    "protobuf>=5.29.0",
+  "httpx>=0.27.0",
+  "loguru>=0.7.3",
+  "transformers>=5.0.0",
+  "huggingface-hub>=0.33.4",
+  "tiktoken>=0.12.0",
+  "jinja2>=3.1.0",
+  "protobuf>=5.29.0",
 ]
 
 [build-system]

bench/single-m3-ultra.toml

@@ -2,10 +2,10 @@
 #
 # Shared constraints applied to ALL benchmarks in this file.
 constraints = [
-    "All(MacOsBuild(=25D125))",
-    "Hosts(=1)",
-    "All(Chip(m3_ultra))",
-    "All(GpuCores(=80))",
+  "All(MacOsBuild(=25D125))",
+  "Hosts(=1)",
+  "All(Chip(m3_ultra))",
+  "All(GpuCores(=80))",
 ]
 
 [topology]

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

@@ -3158,6 +3158,23 @@ class AppStore {
     return (await response.json()) as TraceStatsResponse;
   }
 
+  /**
+   * Delete traces by task IDs
+   */
+  async deleteTraces(
+    taskIds: string[],
+  ): Promise<{ deleted: string[]; notFound: string[] }> {
+    const response = await fetch("/v1/traces/delete", {
+      method: "POST",
+      headers: { "Content-Type": "application/json" },
+      body: JSON.stringify({ taskIds }),
+    });
+    if (!response.ok) {
+      throw new Error(`Failed to delete traces: ${response.status}`);
+    }
+    return await response.json();
+  }
+
   /**
    * Get the URL for the raw trace file (for Perfetto)
    */
@@ -3301,3 +3318,5 @@ export const fetchTraceStats = (taskId: string) =>
   appStore.fetchTraceStats(taskId);
 export const getTraceRawUrl = (taskId: string) =>
   appStore.getTraceRawUrl(taskId);
+export const deleteTraces = (taskIds: string[]) =>
+  appStore.deleteTraces(taskIds);

dashboard/src/routes/+page.svelte

@@ -135,9 +135,14 @@
     if (!rdmaCtl) return false;
     const ids = tbIdentifiers;
     if (!ids) return false;
-    // Find nodes with TB5 hardware (any TB interface)
+    // Find nodes with TB5 hardware (link speed > 40 Gb/s)
     const tb5NodeIds = Object.entries(ids)
-      .filter(([_, node]) => node.interfaces.length > 0)
+      .filter(([_, node]) =>
+        node.interfaces.some((iface) => {
+          const match = iface.linkSpeed.match(/(\d+)\s*Gb/i);
+          return match != null && parseInt(match[1], 10) > 40;
+        }),
+      )
       .map(([id]) => id);
     if (tb5NodeIds.length < 2) return false;
     // At least one TB5 node has RDMA disabled
@@ -6050,12 +6055,11 @@
               </div>
 
               <div
-                class="relative aspect-square bg-exo-dark-gray rounded-lg overflow-hidden"
+                class="relative aspect-square bg-exo-dark-gray rounded-lg overflow-hidden pointer-events-none"
               >
                 <TopologyGraph
                   highlightedNodes={highlightedNodes()}
                   filteredNodes={nodeFilter}
-                  onNodeClick={togglePreviewNodeFilter}
                 />
 
                 {@render clusterWarningsCompact()}

dashboard/src/routes/traces/+page.svelte

@@ -3,6 +3,7 @@
   import {
     listTraces,
     getTraceRawUrl,
+    deleteTraces,
     type TraceListItem,
   } from "$lib/stores/app.svelte";
   import HeaderNav from "$lib/components/HeaderNav.svelte";
@@ -10,6 +11,51 @@
   let traces = $state<TraceListItem[]>([]);
   let loading = $state(true);
   let error = $state<string | null>(null);
+  let selectedIds = $state<Set<string>>(new Set());
+  let deleting = $state(false);
+
+  let allSelected = $derived(
+    traces.length > 0 && selectedIds.size === traces.length,
+  );
+
+  function toggleSelect(taskId: string) {
+    const next = new Set(selectedIds);
+    if (next.has(taskId)) {
+      next.delete(taskId);
+    } else {
+      next.add(taskId);
+    }
+    selectedIds = next;
+  }
+
+  function toggleSelectAll() {
+    if (allSelected) {
+      selectedIds = new Set();
+    } else {
+      selectedIds = new Set(traces.map((t) => t.taskId));
+    }
+  }
+
+  async function handleDelete() {
+    if (selectedIds.size === 0) return;
+    const count = selectedIds.size;
+    if (
+      !confirm(
+        `Delete ${count} trace${count === 1 ? "" : "s"}? This cannot be undone.`,
+      )
+    )
+      return;
+    deleting = true;
+    try {
+      await deleteTraces([...selectedIds]);
+      selectedIds = new Set();
+      await refresh();
+    } catch (e) {
+      error = e instanceof Error ? e.message : "Failed to delete traces";
+    } finally {
+      deleting = false;
+    }
+  }
 
   function formatBytes(bytes: number): string {
     if (!bytes || bytes <= 0) return "0B";
@@ -109,6 +155,16 @@
         </h1>
       </div>
       <div class="flex items-center gap-3">
+        {#if selectedIds.size > 0}
+          <button
+            type="button"
+            class="text-xs font-mono text-red-400 hover:text-red-300 transition-colors uppercase border border-red-500/40 px-2 py-1 rounded"
+            onclick={handleDelete}
+            disabled={deleting}
+          >
+            {deleting ? "Deleting..." : `Delete (${selectedIds.size})`}
+          </button>
+        {/if}
         <button
           type="button"
           class="text-xs font-mono text-exo-light-gray hover:text-exo-yellow transition-colors uppercase border border-exo-medium-gray/40 px-2 py-1 rounded"
@@ -143,14 +199,41 @@
       </div>
     {:else}
       <div class="space-y-3">
+        <div class="flex items-center gap-2 px-1">
+          <button
+            type="button"
+            class="text-xs font-mono uppercase transition-colors {allSelected
+              ? 'text-exo-yellow'
+              : 'text-exo-light-gray hover:text-exo-yellow'}"
+            onclick={toggleSelectAll}
+          >
+            {allSelected ? "Deselect all" : "Select all"}
+          </button>
+        </div>
         {#each traces as trace}
+          {@const isSelected = selectedIds.has(trace.taskId)}
+          <!-- svelte-ignore a11y_no_static_element_interactions -->
           <div
-            class="rounded border border-exo-medium-gray/30 bg-exo-black/30 p-4 flex items-center justify-between gap-4"
+            role="button"
+            tabindex="0"
+            class="w-full text-left rounded border-l-2 border-r border-t border-b transition-all p-4 flex items-center justify-between gap-4 cursor-pointer {isSelected
+              ? 'bg-exo-yellow/10 border-l-exo-yellow border-r-exo-medium-gray/30 border-t-exo-medium-gray/30 border-b-exo-medium-gray/30'
+              : 'bg-exo-black/30 border-l-transparent border-r-exo-medium-gray/30 border-t-exo-medium-gray/30 border-b-exo-medium-gray/30 hover:bg-white/[0.03]'}"
+            onclick={() => toggleSelect(trace.taskId)}
+            onkeydown={(e) => {
+              if (e.key === "Enter" || e.key === " ") {
+                e.preventDefault();
+                toggleSelect(trace.taskId);
+              }
+            }}
           >
             <div class="min-w-0 flex-1">
               <a
                 href="#/traces/{trace.taskId}"
-                class="text-sm font-mono text-white hover:text-exo-yellow transition-colors truncate block"
+                class="text-sm font-mono transition-colors truncate block {isSelected
+                  ? 'text-exo-yellow'
+                  : 'text-white hover:text-exo-yellow'}"
+                onclick={(e) => e.stopPropagation()}
               >
                 {trace.taskId}
               </a>
@@ -160,7 +243,11 @@
                 )}
               </div>
             </div>
-            <div class="flex items-center gap-2 shrink-0">
+            <!-- svelte-ignore a11y_click_events_have_key_events -->
+            <div
+              class="flex items-center gap-2 shrink-0"
+              onclick={(e) => e.stopPropagation()}
+            >
               <a
                 href="#/traces/{trace.taskId}"
                 class="text-xs font-mono text-exo-light-gray hover:text-exo-yellow transition-colors uppercase border border-exo-medium-gray/40 px-2 py-1 rounded"

docs/api.md

@@ -183,6 +183,27 @@ Same schema as `/v1/chat/completions`.
 **Response:**
 Chat completion plus benchmarking metrics.
 
+### Cancel Command
+
+**POST** `/v1/cancel/{command_id}`
+
+Cancels an active generation command (text or image). Notifies workers and closes the stream.
+
+**Path parameters:**
+
+* `command_id`: string, ID of the command to cancel
+
+**Response (example):**
+
+```json
+{
+  "message": "Command cancelled.",
+  "command_id": "cmd-abc-123"
+}
+```
+
+Returns 404 if the command is not found or already completed.
+
 ## 5. Image Generation & Editing
 
 ### Image Generation
@@ -266,6 +287,7 @@ GET     /v1/models
 
 POST    /v1/chat/completions
 POST    /bench/chat/completions
+POST    /v1/cancel/{command_id}
 
 POST    /v1/images/generations
 POST    /bench/images/generations

flake.nix

@@ -108,6 +108,7 @@
                 package = pkgsSwift.swiftPackages.swift-format;
               };
               shfmt.enable = true;
+              taplo.enable = true;
             };
           };
 

pyproject.toml

@@ -5,31 +5,31 @@ description = "Exo"
 readme = "README.md"
 requires-python = ">=3.13"
 dependencies = [
-    "aiofiles>=24.1.0",
-    "aiohttp>=3.12.14",
-    "types-aiofiles>=24.1.0.20250708",
-    "pydantic>=2.11.7",
-    "fastapi>=0.116.1",
-    "filelock>=3.18.0",
-    "rustworkx>=0.17.1",
-    "huggingface-hub>=0.33.4",
-    "psutil>=7.0.0",
-    "loguru>=0.7.3",
-    "exo_pyo3_bindings", # rust bindings
-    "anyio==4.11.0",
-    "mlx; sys_platform == 'darwin'",
-    "mlx[cpu]==0.30.6; sys_platform == 'linux'",
-    "mlx-lm==0.30.7",
-    "tiktoken>=0.12.0", # required for kimi k2 tokenizer
-    "hypercorn>=0.18.0",
-    "openai-harmony>=0.0.8",
-    "httpx>=0.28.1",
-    "tomlkit>=0.14.0",
-    "pillow>=11.0,<12.0", # compatibility with mflux
-    "mflux==0.15.5",
-    "python-multipart>=0.0.21",
-    "msgspec>=0.19.0",
-    "zstandard>=0.23.0",
+  "aiofiles>=24.1.0",
+  "aiohttp>=3.12.14",
+  "types-aiofiles>=24.1.0.20250708",
+  "pydantic>=2.11.7",
+  "fastapi>=0.116.1",
+  "filelock>=3.18.0",
+  "rustworkx>=0.17.1",
+  "huggingface-hub>=0.33.4",
+  "psutil>=7.0.0",
+  "loguru>=0.7.3",
+  "exo_pyo3_bindings",                         # rust bindings
+  "anyio==4.11.0",
+  "mlx; sys_platform == 'darwin'",
+  "mlx[cpu]==0.30.6; sys_platform == 'linux'",
+  "mlx-lm",
+  "tiktoken>=0.12.0",                          # required for kimi k2 tokenizer
+  "hypercorn>=0.18.0",
+  "openai-harmony>=0.0.8",
+  "httpx>=0.28.1",
+  "tomlkit>=0.14.0",
+  "pillow>=11.0,<12.0",                        # compatibility with mflux
+  "mflux==0.15.5",
+  "python-multipart>=0.0.21",
+  "msgspec>=0.19.0",
+  "zstandard>=0.23.0",
 ]
 
 [project.scripts]
@@ -38,12 +38,12 @@ exo = "exo.main:main"
 # dependencies only required for development
 [dependency-groups]
 dev = [
-    "basedpyright>=1.29.0",
-    "pyinstaller>=6.17.0",
-    "pytest>=8.4.0",
-    "pytest-asyncio>=1.0.0",
-    "pytest-env",
-    "ruff>=0.11.13",
+  "basedpyright>=1.29.0",
+  "pyinstaller>=6.17.0",
+  "pytest>=8.4.0",
+  "pytest-asyncio>=1.0.0",
+  "pytest-env",
+  "ruff>=0.11.13",
 ]
 
 # mlx[cuda] requires a newer version of mlx. the ideal on linux is: default to mlx[cpu] unless[cuda] specified.
@@ -57,15 +57,12 @@ dev = [
 ###
 
 [tool.uv.workspace]
-members = [
-    "rust/exo_pyo3_bindings",
-    "bench",
-]
+members = ["rust/exo_pyo3_bindings", "bench"]
 
 [tool.uv.sources]
 exo_pyo3_bindings = { workspace = true }
 mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
-#mlx-lm = { git = "https://github.com/davidmcc73/mlx-lm", branch = "stable" }
+mlx-lm = { git = "https://github.com/ml-explore/mlx-lm", rev = "834fac934c4e04de9b3d723e2b9287a2c60cfd4a" }
 # Uncomment to use local mlx/mlx-lm development versions:
 # mlx = { path = "/Users/Shared/mlx", editable=true }
 # mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }
@@ -95,7 +92,15 @@ reportUnnecessaryTypeIgnoreComment = "error"
 pythonVersion = "3.13"
 pythonPlatform = "Darwin"
 
-exclude = ["**/.venv", "**/venv", "**/__pycache__", "**/exo_scripts", "**/.direnv", "**/rust", "**/.github"]
+exclude = [
+  "**/.venv",
+  "**/venv",
+  "**/__pycache__",
+  "**/exo_scripts",
+  "**/.direnv",
+  "**/rust",
+  "**/.github",
+]
 stubPath = ".mlx_typings"
 
 [[tool.basedpyright.executionEnvironments]]
@@ -109,17 +114,18 @@ root = "src"
 [tool.uv]
 required-version = ">=0.8.6"
 prerelease = "allow"
-environments = [
-    "sys_platform == 'darwin'",
-    "sys_platform == 'linux'",
-]
+environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
 
 ###
 # ruff configuration
 ###
 
 [tool.ruff]
-extend-exclude = ["shared/protobufs/**", "*mlx_typings/**", "rust/exo_pyo3_bindings/**"]
+extend-exclude = [
+  "shared/protobufs/**",
+  "*mlx_typings/**",
+  "rust/exo_pyo3_bindings/**",
+]
 
 [tool.ruff.lint]
 extend-select = ["I", "N", "B", "A", "PIE", "SIM"]
@@ -127,13 +133,7 @@ extend-select = ["I", "N", "B", "A", "PIE", "SIM"]
 [tool.pytest.ini_options]
 pythonpath = "."
 asyncio_mode = "auto"
-markers = [
-    "slow: marks tests as slow (deselected by default)"
-]
-env = [
-  "EXO_TESTS=1"
-]
+markers = ["slow: marks tests as slow (deselected by default)"]
+env = ["EXO_TESTS=1"]
 addopts = "-m 'not slow' --ignore=tests/start_distributed_test.py"
-filterwarnings = [
-    "ignore:builtin type Swig:DeprecationWarning",
-]
+filterwarnings = ["ignore:builtin type Swig:DeprecationWarning"]

resources/inference_model_cards/mlx-community--Qwen3.5-122B-A10B-4bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-122B-A10B-4bit"
+n_layers = 48
+hidden_size = 3072
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "4bit"
+base_model = "Qwen3.5 122B A10B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 69593314272

resources/inference_model_cards/mlx-community--Qwen3.5-122B-A10B-6bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-122B-A10B-6bit"
+n_layers = 48
+hidden_size = 3072
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "6bit"
+base_model = "Qwen3.5 122B A10B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 100120675296

resources/inference_model_cards/mlx-community--Qwen3.5-122B-A10B-8bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-122B-A10B-8bit"
+n_layers = 48
+hidden_size = 3072
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "8bit"
+base_model = "Qwen3.5 122B A10B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 130648036320

resources/inference_model_cards/mlx-community--Qwen3.5-122B-A10B-bf16.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-122B-A10B-bf16"
+n_layers = 48
+hidden_size = 3072
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "bf16"
+base_model = "Qwen3.5 122B A10B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 245125640160

resources/inference_model_cards/mlx-community--Qwen3.5-27B-4bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-27B-4bit"
+n_layers = 64
+hidden_size = 5120
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "4bit"
+base_model = "Qwen3.5 27B"
+capabilities = ["text", "thinking"]
+
+[storage_size]
+in_bytes = 16054266848

resources/inference_model_cards/mlx-community--Qwen3.5-27B-8bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-27B-8bit"
+n_layers = 64
+hidden_size = 5120
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "8bit"
+base_model = "Qwen3.5 27B"
+capabilities = ["text", "thinking"]
+
+[storage_size]
+in_bytes = 29500943328

resources/inference_model_cards/mlx-community--Qwen3.5-2B-MLX-8bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-2B-MLX-8bit"
+n_layers = 24
+hidden_size = 2048
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "8bit"
+base_model = "Qwen3.5 2B"
+capabilities = ["text", "thinking"]
+
+[storage_size]
+in_bytes = 2662787264

resources/inference_model_cards/mlx-community--Qwen3.5-35B-A3B-4bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-35B-A3B-4bit"
+n_layers = 40
+hidden_size = 2048
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "4bit"
+base_model = "Qwen3.5 35B A3B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 20391405152

resources/inference_model_cards/mlx-community--Qwen3.5-35B-A3B-8bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-35B-A3B-8bit"
+n_layers = 40
+hidden_size = 2048
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "8bit"
+base_model = "Qwen3.5 35B A3B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 37721130592

resources/inference_model_cards/mlx-community--Qwen3.5-397B-A17B-4bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-397B-A17B-4bit"
+n_layers = 60
+hidden_size = 4096
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "4bit"
+base_model = "Qwen3.5 397B A17B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 223860768352

resources/inference_model_cards/mlx-community--Qwen3.5-397B-A17B-6bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-397B-A17B-6bit"
+n_layers = 60
+hidden_size = 4096
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "6bit"
+base_model = "Qwen3.5 397B A17B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 322946674272

resources/inference_model_cards/mlx-community--Qwen3.5-397B-A17B-8bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-397B-A17B-8bit"
+n_layers = 60
+hidden_size = 4096
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "8bit"
+base_model = "Qwen3.5 397B A17B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 422032580192

resources/inference_model_cards/mlx-community--Qwen3.5-9B-4bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-9B-4bit"
+n_layers = 32
+hidden_size = 4096
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "4bit"
+base_model = "Qwen3.5 9B"
+capabilities = ["text", "thinking"]
+
+[storage_size]
+in_bytes = 5950062560

resources/inference_model_cards/mlx-community--Qwen3.5-9B-8bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-9B-8bit"
+n_layers = 32
+hidden_size = 4096
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "8bit"
+base_model = "Qwen3.5 9B"
+capabilities = ["text", "thinking"]
+
+[storage_size]
+in_bytes = 10426433504

rust/exo_pyo3_bindings/Cargo.toml

@@ -26,20 +26,24 @@ networking = { workspace = true }
 
 # interop
 pyo3 = { version = "0.27.2", features = [
-    # "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13
-    # "nightly", # enables better-supported GIL integration
-    "experimental-async", # async support in #[pyfunction] & #[pymethods]
-    #"experimental-inspect", # inspection of generated binary => easier to automate type-hint generation
-    #"py-clone", # adding Clone-ing of `Py<T>` without GIL (may cause panics - remove if panics happen)
-    # "multiple-pymethods", # allows multiple #[pymethods] sections per class
+  # "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13
+  # "nightly", # enables better-supported GIL integration
+  "experimental-async", # async support in #[pyfunction] & #[pymethods]
+  #"experimental-inspect", # inspection of generated binary => easier to automate type-hint generation
+  #"py-clone", # adding Clone-ing of `Py<T>` without GIL (may cause panics - remove if panics happen)
+  # "multiple-pymethods", # allows multiple #[pymethods] sections per class
 
-    # integrations with other libraries
-    # "arc_lock", "bigdecimal", "either", "hashbrown", "indexmap", "num-bigint", "num-complex", "num-rational",
-    # "ordered-float", "rust_decimal", "smallvec",
-    # "anyhow", "chrono", "chrono-local", "chrono-tz", "eyre", "jiff-02", "lock_api", "parking-lot", "time",  "serde",
+  # integrations with other libraries
+  # "arc_lock", "bigdecimal", "either", "hashbrown", "indexmap", "num-bigint", "num-complex", "num-rational",
+  # "ordered-float", "rust_decimal", "smallvec",
+  # "anyhow", "chrono", "chrono-local", "chrono-tz", "eyre", "jiff-02", "lock_api", "parking-lot", "time",  "serde",
 ] }
 pyo3-stub-gen = { version = "0.17.2" }
-pyo3-async-runtimes = { version = "0.27.0", features = ["attributes", "tokio-runtime", "testing"] }
+pyo3-async-runtimes = { version = "0.27.0", features = [
+  "attributes",
+  "tokio-runtime",
+  "testing",
+] }
 pyo3-log = "0.13.2"
 
 # macro dependencies

rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi

@@ -2,7 +2,6 @@
 # ruff: noqa: E501, F401
 
 import builtins
-import enum
 import typing
 
 @typing.final
@@ -11,29 +10,6 @@ class AllQueuesFullError(builtins.Exception):
     def __repr__(self) -> builtins.str: ...
     def __str__(self) -> builtins.str: ...
 
-@typing.final
-class ConnectionUpdate:
-    @property
-    def update_type(self) -> ConnectionUpdateType:
-        r"""
-        Whether this is a connection or disconnection event
-        """
-    @property
-    def peer_id(self) -> builtins.str:
-        r"""
-        Identity of the peer that we have connected to or disconnected from.
-        """
-    @property
-    def remote_ipv4(self) -> builtins.str:
-        r"""
-        Remote connection's IPv4 address.
-        """
-    @property
-    def remote_tcp_port(self) -> builtins.int:
-        r"""
-        Remote connection's TCP port.
-        """
-
 @typing.final
 class Keypair:
     r"""
@@ -58,21 +34,15 @@ class Keypair:
         Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
         """
 
+@typing.final
+class MessageTooLargeError(builtins.Exception):
+    def __new__(cls, *args: typing.Any) -> MessageTooLargeError: ...
+    def __repr__(self) -> builtins.str: ...
+    def __str__(self) -> builtins.str: ...
+
 @typing.final
 class NetworkingHandle:
     def __new__(cls, identity: Keypair) -> NetworkingHandle: ...
-    async def connection_update_recv(self) -> ConnectionUpdate:
-        r"""
-        Receives the next `ConnectionUpdate` from networking.
-        """
-    async def connection_update_recv_many(self, limit: builtins.int) -> builtins.list[ConnectionUpdate]:
-        r"""
-        Receives at most `limit` `ConnectionUpdate`s from networking and returns them.
-        
-        For `limit = 0`, an empty collection of `ConnectionUpdate`s will be returned immediately.
-        For `limit > 0`, if there are no `ConnectionUpdate`s in the channel's queue this method
-        will sleep until a `ConnectionUpdate`s is sent.
-        """
     async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
         r"""
         Subscribe to a `GossipSub` topic.
@@ -91,24 +61,7 @@ class NetworkingHandle:
         
         If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
         """
-    async def gossipsub_recv(self) -> tuple[builtins.str, bytes]:
-        r"""
-        Receives the next message from the `GossipSub` network.
-        """
-    async def gossipsub_recv_many(self, limit: builtins.int) -> builtins.list[tuple[builtins.str, bytes]]:
-        r"""
-        Receives at most `limit` messages from the `GossipSub` network and returns them.
-        
-        For `limit = 0`, an empty collection of messages will be returned immediately.
-        For `limit > 0`, if there are no messages in the channel's queue this method
-        will sleep until a message is sent.
-        """
-
-@typing.final
-class MessageTooLargeError(builtins.Exception):
-    def __new__(cls, *args: typing.Any) -> MessageTooLargeError: ...
-    def __repr__(self) -> builtins.str: ...
-    def __str__(self) -> builtins.str: ...
+    async def recv(self) -> PyFromSwarm: ...
 
 @typing.final
 class NoPeersSubscribedToTopicError(builtins.Exception):
@@ -116,11 +69,26 @@ class NoPeersSubscribedToTopicError(builtins.Exception):
     def __repr__(self) -> builtins.str: ...
     def __str__(self) -> builtins.str: ...
 
-@typing.final
-class ConnectionUpdateType(enum.Enum):
-    r"""
-    Connection or disconnection event discriminant type.
-    """
-    Connected = ...
-    Disconnected = ...
+class PyFromSwarm:
+    @typing.final
+    class Connection(PyFromSwarm):
+        __match_args__ = ("peer_id", "connected",)
+        @property
+        def peer_id(self) -> builtins.str: ...
+        @property
+        def connected(self) -> builtins.bool: ...
+        def __new__(cls, peer_id: builtins.str, connected: builtins.bool) -> PyFromSwarm.Connection: ...
+    
+    @typing.final
+    class Message(PyFromSwarm):
+        __match_args__ = ("origin", "topic", "data",)
+        @property
+        def origin(self) -> builtins.str: ...
+        @property
+        def topic(self) -> builtins.str: ...
+        @property
+        def data(self) -> bytes: ...
+        def __new__(cls, origin: builtins.str, topic: builtins.str, data: bytes) -> PyFromSwarm.Message: ...
+    
+    ...
 

rust/exo_pyo3_bindings/pyproject.toml

@@ -4,21 +4,18 @@ build-backend = "maturin"
 
 [project]
 name = "exo_pyo3_bindings"
-version = "0.1.0"
+version = "0.2.0"
 description = "Add your description here"
 readme = "README.md"
 authors = [
-    { name = "Andrei Cravtov", email = "the.andrei.cravtov@gmail.com" }
+  { name = "Andrei Cravtov", email = "the.andrei.cravtov@gmail.com" },
+  { name = "Evan Quiney", email = "evanev7@gmail.com" },
 ]
 requires-python = ">=3.13"
 dependencies = []
 
 [dependency-groups]
-dev = [
-    "exo_pyo3_bindings",
-    "pytest>=8.4.0",
-    "pytest-asyncio>=1.0.0",
-]
+dev = ["exo_pyo3_bindings", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"]
 
 [tool.maturin]
 #purelib = true

rust/exo_pyo3_bindings/src/lib.rs

@@ -155,6 +155,9 @@ pub(crate) mod ext {
 fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     // install logger
     pyo3_log::init();
+    let mut builder = tokio::runtime::Builder::new_multi_thread();
+    builder.enable_all();
+    pyo3_async_runtimes::tokio::init(builder);
 
     // TODO: for now this is all NOT a submodule, but figure out how to make the submodule system
     //       work with maturin, where the types generate correctly, in the right folder, without

rust/exo_pyo3_bindings/src/networking.rs

@@ -1,26 +1,24 @@
-#![allow(
-    clippy::multiple_inherent_impl,
-    clippy::unnecessary_wraps,
-    clippy::unused_self,
-    clippy::needless_pass_by_value
-)]
+use std::pin::Pin;
+use std::sync::Arc;
 
 use crate::r#const::MPSC_CHANNEL_SIZE;
 use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
-use crate::ext::{ResultExt as _, TokioMpscReceiverExt as _, TokioMpscSenderExt as _};
+use crate::ext::{ResultExt as _, TokioMpscSenderExt as _};
 use crate::ident::PyKeypair;
+use crate::networking::exception::{
+    PyAllQueuesFullError, PyMessageTooLargeError, PyNoPeersSubscribedToTopicError,
+};
 use crate::pyclass;
-use libp2p::futures::StreamExt as _;
-use libp2p::gossipsub;
-use libp2p::gossipsub::{IdentTopic, Message, MessageId, PublishError};
-use libp2p::swarm::SwarmEvent;
-use networking::discovery;
-use networking::swarm::create_swarm;
+use futures_lite::{Stream, StreamExt as _};
+use libp2p::gossipsub::PublishError;
+use networking::swarm::{FromSwarm, ToSwarm, create_swarm};
+use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::{PyModule, PyModuleMethods as _};
 use pyo3::types::PyBytes;
-use pyo3::{Bound, Py, PyErr, PyResult, PyTraverseError, PyVisit, Python, pymethods};
-use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_enum, gen_stub_pymethods};
-use std::net::IpAddr;
+use pyo3::{Bound, Py, PyAny, PyErr, PyResult, Python, pymethods};
+use pyo3_stub_gen::derive::{
+    gen_methods_from_python, gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods,
+};
 use tokio::sync::{Mutex, mpsc, oneshot};
 
 mod exception {
@@ -131,237 +129,45 @@ mod exception {
     }
 }
 
-/// Connection or disconnection event discriminant type.
-#[gen_stub_pyclass_enum]
-#[pyclass(eq, eq_int, name = "ConnectionUpdateType")]
-#[derive(Debug, Clone, PartialEq)]
-enum PyConnectionUpdateType {
-    Connected = 0,
-    Disconnected,
-}
-
-#[gen_stub_pyclass]
-#[pyclass(frozen, name = "ConnectionUpdate")]
-#[derive(Debug, Clone)]
-struct PyConnectionUpdate {
-    /// Whether this is a connection or disconnection event
-    #[pyo3(get)]
-    update_type: PyConnectionUpdateType,
-
-    /// Identity of the peer that we have connected to or disconnected from.
-    #[pyo3(get)]
-    peer_id: String,
-
-    /// Remote connection's IPv4 address.
-    #[pyo3(get)]
-    remote_ipv4: String,
-
-    /// Remote connection's TCP port.
-    #[pyo3(get)]
-    remote_tcp_port: u16,
-}
-
-enum ToTask {
-    GossipsubSubscribe {
-        topic: String,
-        result_tx: oneshot::Sender<PyResult<bool>>,
-    },
-    GossipsubUnsubscribe {
-        topic: String,
-        result_tx: oneshot::Sender<bool>,
-    },
-    GossipsubPublish {
-        topic: String,
-        data: Vec<u8>,
-        result_tx: oneshot::Sender<PyResult<MessageId>>,
-    },
-}
-
-#[allow(clippy::enum_glob_use)]
-async fn networking_task(
-    mut swarm: networking::swarm::Swarm,
-    mut to_task_rx: mpsc::Receiver<ToTask>,
-    connection_update_tx: mpsc::Sender<PyConnectionUpdate>,
-    gossipsub_message_tx: mpsc::Sender<(String, Vec<u8>)>,
-) {
-    use SwarmEvent::*;
-    use ToTask::*;
-    use networking::swarm::BehaviourEvent::*;
-
-    log::info!("RUST: networking task started");
-
-    loop {
-        tokio::select! {
-            message = to_task_rx.recv() => {
-                // handle closed channel
-                let Some(message) = message else {
-                    log::info!("RUST: channel closed");
-                    break;
-                };
-
-                // dispatch incoming messages
-                match message {
-                    GossipsubSubscribe { topic, result_tx } => {
-                        // try to subscribe
-                        let result = swarm.behaviour_mut()
-                            .gossipsub.subscribe(&IdentTopic::new(topic));
-
-                        // send response oneshot
-                        if let Err(e) = result_tx.send(result.pyerr()) {
-                            log::error!("RUST: could not subscribe to gossipsub topic since channel already closed: {e:?}");
-                            continue;
-                        }
-                    }
-                    GossipsubUnsubscribe { topic, result_tx } => {
-                        // try to unsubscribe from the topic
-                        let result = swarm.behaviour_mut()
-                            .gossipsub.unsubscribe(&IdentTopic::new(topic));
-
-                        // send response oneshot (or exit if connection closed)
-                        if let Err(e) = result_tx.send(result) {
-                            log::error!("RUST: could not unsubscribe from gossipsub topic since channel already closed: {e:?}");
-                            continue;
-                        }
-                    }
-                    GossipsubPublish { topic, data, result_tx } => {
-                        // try to publish the data -> catch NoPeersSubscribedToTopic error & convert to correct exception
-                        let result = swarm.behaviour_mut().gossipsub.publish(
-                            IdentTopic::new(topic), data);
-                        let pyresult: PyResult<MessageId> = if let Err(PublishError::NoPeersSubscribedToTopic) = result {
-                            Err(exception::PyNoPeersSubscribedToTopicError::new_err())
-                        } else if let Err(PublishError::AllQueuesFull(_)) = result {
-                            Err(exception::PyAllQueuesFullError::new_err())
-                        } else if let Err(PublishError::MessageTooLarge) = result {
-                            Err(exception::PyMessageTooLargeError::new_err())
-                        } else {
-                            result.pyerr()
-                        };
-
-                        // send response oneshot (or exit if connection closed)
-                        if let Err(e) = result_tx.send(pyresult) {
-                            log::error!("RUST: could not publish gossipsub message since channel already closed: {e:?}");
-                            continue;
-                        }
-                    }
-                }
-            }
-
-            // architectural solution to this problem:
-            // create keep_alive behavior who's job it is to dial peers discovered by mDNS (and drop when expired)
-            //   -> it will emmit TRUE connected/disconnected events consumable elsewhere
-            //
-            // gossipsub will feed off-of dial attempts created by networking, and that will bootstrap its' peers list
-            // then for actual communication it will dial those peers if need-be
-            swarm_event = swarm.select_next_some() => {
-                match swarm_event {
-                    Behaviour(Gossipsub(gossipsub::Event::Message {
-                        message: Message {
-                            topic,
-                            data,
-                            ..
-                        },
-                        ..
-                    })) => {
-                        // topic-ID is just the topic hash!!! (since we used identity hasher)
-                        let message = (topic.into_string(), data);
-
-                        // send incoming message to channel (or exit if connection closed)
-                        if let Err(e) = gossipsub_message_tx.send(message).await {
-                            log::error!("RUST: could not send incoming gossipsub message since channel already closed: {e}");
-                            continue;
-                        }
-                    },
-                    Behaviour(Discovery(discovery::Event::ConnectionEstablished { peer_id, remote_ip, remote_tcp_port, .. })) => {
-                        // grab IPv4 string
-                        let remote_ipv4 = match remote_ip {
-                            IpAddr::V4(ip) => ip.to_string(),
-                            IpAddr::V6(ip) => {
-                                log::warn!("RUST: ignoring connection to IPv6 address: {ip}");
-                                continue;
-                            }
-                        };
-
-                        // send connection event to channel (or exit if connection closed)
-                        if let Err(e) = connection_update_tx.send(PyConnectionUpdate {
-                            update_type: PyConnectionUpdateType::Connected,
-                            peer_id: peer_id.to_base58(),
-                            remote_ipv4,
-                            remote_tcp_port,
-                        }).await {
-                            log::error!("RUST: could not send connection update since channel already closed: {e}");
-                            continue;
-                        }
-                    },
-                    Behaviour(Discovery(discovery::Event::ConnectionClosed { peer_id, remote_ip, remote_tcp_port, .. })) => {
-                        // grab IPv4 string
-                        let remote_ipv4 = match remote_ip {
-                            IpAddr::V4(ip) => ip.to_string(),
-                            IpAddr::V6(ip) => {
-                                log::warn!("RUST: ignoring disconnection from IPv6 address: {ip}");
-                                continue;
-                            }
-                        };
-
-                        // send disconnection event to channel (or exit if connection closed)
-                        if let Err(e) = connection_update_tx.send(PyConnectionUpdate {
-                            update_type: PyConnectionUpdateType::Disconnected,
-                            peer_id: peer_id.to_base58(),
-                            remote_ipv4,
-                            remote_tcp_port,
-                        }).await {
-                            log::error!("RUST: could not send connection update since channel already closed: {e}");
-                            continue;
-                        }
-                    },
-                    e => {
-                        log::info!("RUST: other event {e:?}");
-                    }
-                }
-            }
-        }
-    }
-
-    log::info!("RUST: networking task stopped");
-}
-
 #[gen_stub_pyclass]
 #[pyclass(name = "NetworkingHandle")]
-#[derive(Debug)]
 struct PyNetworkingHandle {
     // channels
-    to_task_tx: Option<mpsc::Sender<ToTask>>,
-    connection_update_rx: Mutex<mpsc::Receiver<PyConnectionUpdate>>,
-    gossipsub_message_rx: Mutex<mpsc::Receiver<(String, Vec<u8>)>>,
+    pub to_swarm: mpsc::Sender<ToSwarm>,
+    pub swarm: Arc<Mutex<Pin<Box<dyn Stream<Item = FromSwarm> + Send>>>>,
 }
 
-impl Drop for PyNetworkingHandle {
-    fn drop(&mut self) {
-        // TODO: may or may not need to await a "kill-signal" oneshot channel message,
-        //       to ensure that the networking task is done BEFORE exiting the clear function...
-        //       but this may require GIL?? and it may not be safe to call GIL here??
-        self.to_task_tx = None; // Using Option<T> as a trick to force channel to be dropped
-    }
+#[gen_stub_pyclass_complex_enum]
+#[pyclass]
+enum PyFromSwarm {
+    Connection {
+        peer_id: String,
+        connected: bool,
+    },
+    Message {
+        origin: String,
+        topic: String,
+        data: Py<PyBytes>,
+    },
 }
-
-#[allow(clippy::expect_used)]
-impl PyNetworkingHandle {
-    fn new(
-        to_task_tx: mpsc::Sender<ToTask>,
-        connection_update_rx: mpsc::Receiver<PyConnectionUpdate>,
-        gossipsub_message_rx: mpsc::Receiver<(String, Vec<u8>)>,
-    ) -> Self {
-        Self {
-            to_task_tx: Some(to_task_tx),
-            connection_update_rx: Mutex::new(connection_update_rx),
-            gossipsub_message_rx: Mutex::new(gossipsub_message_rx),
+impl From<FromSwarm> for PyFromSwarm {
+    fn from(value: FromSwarm) -> Self {
+        match value {
+            FromSwarm::Discovered { peer_id } => Self::Connection {
+                peer_id: peer_id.to_base58(),
+                connected: true,
+            },
+            FromSwarm::Expired { peer_id } => Self::Connection {
+                peer_id: peer_id.to_base58(),
+                connected: false,
+            },
+            FromSwarm::Message { from, topic, data } => Self::Message {
+                origin: from.to_base58(),
+                topic: topic,
+                data: data.pybytes(),
+            },
         }
     }
-
-    const fn to_task_tx(&self) -> &mpsc::Sender<ToTask> {
-        self.to_task_tx
-            .as_ref()
-            .expect("The sender should only be None after de-initialization.")
-    }
 }
 
 #[gen_stub_pymethods]
@@ -375,97 +181,36 @@ impl PyNetworkingHandle {
 
     #[new]
     fn py_new(identity: Bound<'_, PyKeypair>) -> PyResult<Self> {
-        use pyo3_async_runtimes::tokio::get_runtime;
-
         // create communication channels
-        let (to_task_tx, to_task_rx) = mpsc::channel(MPSC_CHANNEL_SIZE);
-        let (connection_update_tx, connection_update_rx) = mpsc::channel(MPSC_CHANNEL_SIZE);
-        let (gossipsub_message_tx, gossipsub_message_rx) = mpsc::channel(MPSC_CHANNEL_SIZE);
+        let (to_swarm, from_client) = mpsc::channel(MPSC_CHANNEL_SIZE);
 
         // get identity
         let identity = identity.borrow().0.clone();
 
         // create networking swarm (within tokio context!! or it crashes)
-        let swarm = get_runtime()
-            .block_on(async { create_swarm(identity) })
-            .pyerr()?;
+        let _guard = pyo3_async_runtimes::tokio::get_runtime().enter();
+        let swarm = create_swarm(identity, from_client).pyerr()?.into_stream();
 
-        // spawn tokio task running the networking logic
-        get_runtime().spawn(async move {
-            networking_task(
-                swarm,
-                to_task_rx,
-                connection_update_tx,
-                gossipsub_message_tx,
-            )
-            .await;
-        });
-        Ok(Self::new(
-            to_task_tx,
-            connection_update_rx,
-            gossipsub_message_rx,
-        ))
+        Ok(Self {
+            swarm: Arc::new(Mutex::new(swarm)),
+            to_swarm,
+        })
     }
 
     #[gen_stub(skip)]
-    const fn __traverse__(&self, _visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
-        Ok(()) // This is needed purely so `__clear__` can work
+    fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        let swarm = Arc::clone(&self.swarm);
+        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+            swarm
+                .try_lock()
+                .map_err(|_| PyRuntimeError::new_err("called recv twice concurrently"))?
+                .next()
+                .await
+                .ok_or(PyErr::receiver_channel_closed())
+                .map(PyFromSwarm::from)
+        })
     }
 
-    #[gen_stub(skip)]
-    fn __clear__(&mut self) {
-        // TODO: may or may not need to await a "kill-signal" oneshot channel message,
-        //       to ensure that the networking task is done BEFORE exiting the clear function...
-        //       but this may require GIL?? and it may not be safe to call GIL here??
-        self.to_task_tx = None; // Using Option<T> as a trick to force channel to be dropped
-    }
-
-    // ---- Connection update receiver methods ----
-
-    /// Receives the next `ConnectionUpdate` from networking.
-    async fn connection_update_recv(&self) -> PyResult<PyConnectionUpdate> {
-        self.connection_update_rx
-            .lock()
-            .allow_threads_py() // allow-threads-aware async call
-            .await
-            .recv_py()
-            .allow_threads_py() // allow-threads-aware async call
-            .await
-    }
-
-    /// Receives at most `limit` `ConnectionUpdate`s from networking and returns them.
-    ///
-    /// For `limit = 0`, an empty collection of `ConnectionUpdate`s will be returned immediately.
-    /// For `limit > 0`, if there are no `ConnectionUpdate`s in the channel's queue this method
-    /// will sleep until a `ConnectionUpdate`s is sent.
-    async fn connection_update_recv_many(&self, limit: usize) -> PyResult<Vec<PyConnectionUpdate>> {
-        self.connection_update_rx
-            .lock()
-            .allow_threads_py() // allow-threads-aware async call
-            .await
-            .recv_many_py(limit)
-            .allow_threads_py() // allow-threads-aware async call
-            .await
-    }
-
-    // TODO: rn this blocks main thread if anything else is awaiting the channel (bc its a mutex)
-    //       so its too dangerous to expose just yet. figure out a better semantics for handling this,
-    //       so things don't randomly block
-    // /// Tries to receive the next `ConnectionUpdate` from networking.
-    // fn connection_update_try_recv(&self) -> PyResult<Option<PyConnectionUpdate>> {
-    //     self.connection_update_rx.blocking_lock().try_recv_py()
-    // }
-    //
-    // /// Checks if the `ConnectionUpdate` channel is empty.
-    // fn connection_update_is_empty(&self) -> bool {
-    //     self.connection_update_rx.blocking_lock().is_empty()
-    // }
-    //
-    // /// Returns the number of `ConnectionUpdate`s in the channel.
-    // fn connection_update_len(&self) -> usize {
-    //     self.connection_update_rx.blocking_lock().len()
-    // }
-
     // ---- Gossipsub management methods ----
 
     /// Subscribe to a `GossipSub` topic.
@@ -475,10 +220,10 @@ impl PyNetworkingHandle {
         let (tx, rx) = oneshot::channel();
 
         // send off request to subscribe
-        self.to_task_tx()
-            .send_py(ToTask::GossipsubSubscribe {
+        self.to_swarm
+            .send_py(ToSwarm::Subscribe {
                 topic,
-                result_tx: tx,
+                result_sender: tx,
             })
             .allow_threads_py() // allow-threads-aware async call
             .await?;
@@ -487,6 +232,7 @@ impl PyNetworkingHandle {
         rx.allow_threads_py() // allow-threads-aware async call
             .await
             .map_err(|_| PyErr::receiver_channel_closed())?
+            .pyerr()
     }
 
     /// Unsubscribes from a `GossipSub` topic.
@@ -496,10 +242,10 @@ impl PyNetworkingHandle {
         let (tx, rx) = oneshot::channel();
 
         // send off request to unsubscribe
-        self.to_task_tx()
-            .send_py(ToTask::GossipsubUnsubscribe {
+        self.to_swarm
+            .send_py(ToSwarm::Unsubscribe {
                 topic,
-                result_tx: tx,
+                result_sender: tx,
             })
             .allow_threads_py() // allow-threads-aware async call
             .await?;
@@ -518,11 +264,11 @@ impl PyNetworkingHandle {
 
         // send off request to subscribe
         let data = Python::attach(|py| Vec::from(data.as_bytes(py)));
-        self.to_task_tx()
-            .send_py(ToTask::GossipsubPublish {
+        self.to_swarm
+            .send_py(ToSwarm::Publish {
                 topic,
                 data,
-                result_tx: tx,
+                result_sender: tx,
             })
             .allow_threads_py() // allow-threads-aware async call
             .await?;
@@ -531,64 +277,26 @@ impl PyNetworkingHandle {
         let _ = rx
             .allow_threads_py() // allow-threads-aware async call
             .await
-            .map_err(|_| PyErr::receiver_channel_closed())??;
+            .map_err(|_| PyErr::receiver_channel_closed())?
+            .map_err(|e| match e {
+                PublishError::AllQueuesFull(_) => PyAllQueuesFullError::new_err(),
+                PublishError::MessageTooLarge => PyMessageTooLargeError::new_err(),
+                PublishError::NoPeersSubscribedToTopic => {
+                    PyNoPeersSubscribedToTopicError::new_err()
+                }
+                e => PyRuntimeError::new_err(e.to_string()),
+            })?;
         Ok(())
     }
+}
 
-    // ---- Gossipsub message receiver methods ----
-
-    /// Receives the next message from the `GossipSub` network.
-    async fn gossipsub_recv(&self) -> PyResult<(String, Py<PyBytes>)> {
-        self.gossipsub_message_rx
-            .lock()
-            .allow_threads_py() // allow-threads-aware async call
-            .await
-            .recv_py()
-            .allow_threads_py() // allow-threads-aware async call
-            .await
-            .map(|(t, d)| (t, d.pybytes()))
+pyo3_stub_gen::inventory::submit! {
+    gen_methods_from_python! {
+        r#"
+            class PyNetworkingHandle:
+                async def recv() -> PyFromSwarm: ...
+        "#
     }
-
-    /// Receives at most `limit` messages from the `GossipSub` network and returns them.
-    ///
-    /// For `limit = 0`, an empty collection of messages will be returned immediately.
-    /// For `limit > 0`, if there are no messages in the channel's queue this method
-    /// will sleep until a message is sent.
-    async fn gossipsub_recv_many(&self, limit: usize) -> PyResult<Vec<(String, Py<PyBytes>)>> {
-        Ok(self
-            .gossipsub_message_rx
-            .lock()
-            .allow_threads_py() // allow-threads-aware async call
-            .await
-            .recv_many_py(limit)
-            .allow_threads_py() // allow-threads-aware async call
-            .await?
-            .into_iter()
-            .map(|(t, d)| (t, d.pybytes()))
-            .collect())
-    }
-
-    // TODO: rn this blocks main thread if anything else is awaiting the channel (bc its a mutex)
-    //       so its too dangerous to expose just yet. figure out a better semantics for handling this,
-    //       so things don't randomly block
-    // /// Tries to receive the next message from the `GossipSub` network.
-    // fn gossipsub_try_recv(&self) -> PyResult<Option<(String, Py<PyBytes>)>> {
-    //     Ok(self
-    //         .gossipsub_message_rx
-    //         .blocking_lock()
-    //         .try_recv_py()?
-    //         .map(|(t, d)| (t, d.pybytes())))
-    // }
-    //
-    // /// Checks if the `GossipSub` message channel is empty.
-    // fn gossipsub_is_empty(&self) -> bool {
-    //     self.gossipsub_message_rx.blocking_lock().is_empty()
-    // }
-    //
-    // /// Returns the number of `GossipSub` messages in the channel.
-    // fn gossipsub_len(&self) -> usize {
-    //     self.gossipsub_message_rx.blocking_lock().len()
-    // }
 }
 
 pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
@@ -596,10 +304,8 @@ pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<exception::PyAllQueuesFullError>()?;
     m.add_class::<exception::PyMessageTooLargeError>()?;
 
-    m.add_class::<PyConnectionUpdateType>()?;
-    m.add_class::<PyConnectionUpdate>()?;
-    m.add_class::<PyConnectionUpdateType>()?;
     m.add_class::<PyNetworkingHandle>()?;
+    m.add_class::<PyFromSwarm>()?;
 
     Ok(())
 }

rust/exo_pyo3_bindings/tests/test_python.py

@@ -1,16 +1,20 @@
 import asyncio
 
 import pytest
-from exo_pyo3_bindings import Keypair, NetworkingHandle, NoPeersSubscribedToTopicError
+from exo_pyo3_bindings import (
+    Keypair,
+    NetworkingHandle,
+    NoPeersSubscribedToTopicError,
+    PyFromSwarm,
+)
 
 
 @pytest.mark.asyncio
 async def test_sleep_on_multiple_items() -> None:
     print("PYTHON: starting handle")
-    h = NetworkingHandle(Keypair.generate_ed25519())
+    h = NetworkingHandle(Keypair.generate())
 
-    ct = asyncio.create_task(_await_cons(h))
-    mt = asyncio.create_task(_await_msg(h))
+    rt = asyncio.create_task(_await_recv(h))
 
     # sleep for 4 ticks
     for i in range(4):
@@ -22,13 +26,11 @@ async def test_sleep_on_multiple_items() -> None:
             print("caught it", e)
 
 
-async def _await_cons(h: NetworkingHandle):
+async def _await_recv(h: NetworkingHandle):
     while True:
-        c = await h.connection_update_recv()
-        print(f"PYTHON: connection update: {c}")
-
-
-async def _await_msg(h: NetworkingHandle):
-    while True:
-        m = await h.gossipsub_recv()
-        print(f"PYTHON: message: {m}")
+        event = await h.recv()
+        match event:
+            case PyFromSwarm.Connection() as c:
+                print(f"PYTHON: connection update: {c}")
+            case PyFromSwarm.Message() as m:
+                print(f"PYTHON: message: {m}")

rust/networking/Cargo.toml

@@ -21,13 +21,17 @@ extend = { workspace = true }
 delegate = { workspace = true }
 
 # async
-tokio = { workspace = true, features = ["full"] }
+async-stream = { workspace = true }
 futures-lite = { workspace = true }
 futures-timer = { workspace = true }
+tokio = { workspace = true, features = ["full"] }
 
 # utility dependencies
 util = { workspace = true }
-tracing-subscriber = { version = "0.3.19", features = ["default", "env-filter"] }
+tracing-subscriber = { version = "0.3.19", features = [
+  "default",
+  "env-filter",
+] }
 keccak-const = { workspace = true }
 
 # tracing/logging
@@ -35,3 +39,4 @@ log = { workspace = true }
 
 # networking
 libp2p = { workspace = true, features = ["full"] }
+pin-project = "1.1.10"

rust/networking/examples/chatroom.rs

@@ -1,7 +1,9 @@
 use futures_lite::StreamExt;
-use libp2p::{gossipsub, identity, swarm::SwarmEvent};
-use networking::{discovery, swarm};
-use tokio::{io, io::AsyncBufReadExt as _, select};
+use libp2p::identity;
+use networking::swarm;
+use networking::swarm::{FromSwarm, ToSwarm};
+use tokio::sync::{mpsc, oneshot};
+use tokio::{io, io::AsyncBufReadExt as _};
 use tracing_subscriber::EnvFilter;
 use tracing_subscriber::filter::LevelFilter;
 
@@ -11,64 +13,69 @@ async fn main() {
         .with_env_filter(EnvFilter::from_default_env().add_directive(LevelFilter::INFO.into()))
         .try_init();
 
+    let (to_swarm, from_client) = mpsc::channel(20);
+
     // Configure swarm
-    let mut swarm =
-        swarm::create_swarm(identity::Keypair::generate_ed25519()).expect("Swarm creation failed");
+    let mut swarm = swarm::create_swarm(identity::Keypair::generate_ed25519(), from_client)
+        .expect("Swarm creation failed")
+        .into_stream();
 
     // Create a Gossipsub topic & subscribe
-    let topic = gossipsub::IdentTopic::new("test-net");
-    swarm
-        .behaviour_mut()
-        .gossipsub
-        .subscribe(&topic)
-        .expect("Subscribing to topic failed");
+    let (tx, rx) = oneshot::channel();
+    _ = to_swarm
+        .send(ToSwarm::Subscribe {
+            topic: "test-net".to_string(),
+            result_sender: tx,
+        })
+        .await
+        .expect("should send");
 
     // Read full lines from stdin
     let mut stdin = io::BufReader::new(io::stdin()).lines();
     println!("Enter messages via STDIN and they will be sent to connected peers using Gossipsub");
 
+    tokio::task::spawn(async move {
+        rx.await
+            .expect("tx not dropped")
+            .expect("subscribe shouldn't fail");
+        loop {
+            if let Ok(Some(line)) = stdin.next_line().await {
+                let (tx, rx) = oneshot::channel();
+                if let Err(e) = to_swarm
+                    .send(swarm::ToSwarm::Publish {
+                        topic: "test-net".to_string(),
+                        data: line.as_bytes().to_vec(),
+                        result_sender: tx,
+                    })
+                    .await
+                {
+                    println!("Send error: {e:?}");
+                    return;
+                };
+                match rx.await {
+                    Ok(Err(e)) => println!("Publish error: {e:?}"),
+                    Err(e) => println!("Publish error: {e:?}"),
+                    Ok(_) => {}
+                }
+            }
+        }
+    });
+
     // Kick it off
     loop {
-        select! {
-            // on gossipsub outgoing
-            Ok(Some(line)) = stdin.next_line() => {
-                if let Err(e) = swarm
-                    .behaviour_mut().gossipsub
-                    .publish(topic.clone(), line.as_bytes()) {
-                    println!("Publish error: {e:?}");
-                }
+        // on gossipsub outgoing
+        match swarm.next().await {
+            // on gossipsub incoming
+            Some(FromSwarm::Discovered { peer_id }) => {
+                println!("\n\nconnected to {peer_id}\n\n")
             }
-            event = swarm.next() => match event {
-                // on gossipsub incoming
-                Some(SwarmEvent::Behaviour(swarm::BehaviourEvent::Gossipsub(gossipsub::Event::Message {
-                    propagation_source: peer_id,
-                    message_id: id,
-                    message,
-                }))) => println!(
-                        "\n\nGot message: '{}' with id: {id} from peer: {peer_id}\n\n",
-                        String::from_utf8_lossy(&message.data),
-                    ),
-
-                // on discovery
-                Some(SwarmEvent::Behaviour(swarm::BehaviourEvent::Discovery(e)) )=> match e {
-                    discovery::Event::ConnectionEstablished {
-                        peer_id, connection_id, remote_ip, remote_tcp_port
-                    } => {
-                        println!("\n\nConnected to: {peer_id}; connection ID: {connection_id}; remote IP: {remote_ip}; remote TCP port: {remote_tcp_port}\n\n");
-                    }
-                    discovery::Event::ConnectionClosed {
-                        peer_id, connection_id, remote_ip, remote_tcp_port
-                    } => {
-                        eprintln!("\n\nDisconnected from: {peer_id}; connection ID: {connection_id}; remote IP: {remote_ip}; remote TCP port: {remote_tcp_port}\n\n");
-                    }
-                }
-
-                // ignore outgoing errors: those are normal
-                e@Some(SwarmEvent::OutgoingConnectionError { .. }) => { log::debug!("Outgoing connection error: {e:?}"); }
-
-                // otherwise log any other event
-                e => { log::info!("Other event {e:?}"); }
+            Some(FromSwarm::Expired { peer_id }) => {
+                println!("\n\ndisconnected from {peer_id}\n\n")
             }
+            Some(FromSwarm::Message { from, topic, data }) => {
+                println!("{topic}/{from}:\n{}", String::from_utf8_lossy(&data))
+            }
+            None => {}
         }
     }
 }

rust/networking/src/swarm.rs

@@ -1,9 +1,11 @@
-use crate::alias;
-use crate::swarm::transport::tcp_transport;
-pub use behaviour::{Behaviour, BehaviourEvent};
-use libp2p::{SwarmBuilder, identity};
+use std::pin::Pin;
 
-pub type Swarm = libp2p::Swarm<Behaviour>;
+use crate::swarm::transport::tcp_transport;
+use crate::{alias, discovery};
+pub use behaviour::{Behaviour, BehaviourEvent};
+use futures_lite::{Stream, StreamExt};
+use libp2p::{PeerId, SwarmBuilder, gossipsub, identity, swarm::SwarmEvent};
+use tokio::sync::{mpsc, oneshot};
 
 /// The current version of the network: this prevents devices running different versions of the
 /// software from interacting with each other.
@@ -15,8 +17,136 @@ pub type Swarm = libp2p::Swarm<Behaviour>;
 pub const NETWORK_VERSION: &[u8] = b"v0.0.1";
 pub const OVERRIDE_VERSION_ENV_VAR: &str = "EXO_LIBP2P_NAMESPACE";
 
+// Uses oneshot senders to emulate function calling apis while avoiding requiring unique ownership
+// of the Swarm.
+pub enum ToSwarm {
+    Unsubscribe {
+        topic: String,
+        result_sender: oneshot::Sender<bool>,
+    },
+    Subscribe {
+        topic: String,
+        result_sender: oneshot::Sender<Result<bool, gossipsub::SubscriptionError>>,
+    },
+    Publish {
+        topic: String,
+        data: Vec<u8>,
+        result_sender: oneshot::Sender<Result<gossipsub::MessageId, gossipsub::PublishError>>,
+    },
+}
+pub enum FromSwarm {
+    Message {
+        from: PeerId,
+        topic: String,
+        data: Vec<u8>,
+    },
+    Discovered {
+        peer_id: PeerId,
+    },
+    Expired {
+        peer_id: PeerId,
+    },
+}
+
+pub struct Swarm {
+    swarm: libp2p::Swarm<Behaviour>,
+    from_client: mpsc::Receiver<ToSwarm>,
+}
+
+impl Swarm {
+    pub fn into_stream(self) -> Pin<Box<dyn Stream<Item = FromSwarm> + Send>> {
+        let Swarm {
+            mut swarm,
+            mut from_client,
+        } = self;
+        let stream = async_stream::stream! {
+            loop {
+                tokio::select! {
+                    msg = from_client.recv() => {
+                        let Some(msg) = msg else { break };
+                        on_message(&mut swarm, msg);
+                    }
+                    event = swarm.next() => {
+                        let Some(event) = event else { break };
+                        if let Some(item) = filter_swarm_event(event) {
+                            yield item;
+                        }
+                    }
+                }
+            }
+        };
+        Box::pin(stream)
+    }
+}
+
+fn on_message(swarm: &mut libp2p::Swarm<Behaviour>, message: ToSwarm) {
+    match message {
+        ToSwarm::Subscribe {
+            topic,
+            result_sender,
+        } => {
+            let result = swarm
+                .behaviour_mut()
+                .gossipsub
+                .subscribe(&gossipsub::IdentTopic::new(topic));
+            _ = result_sender.send(result);
+        }
+        ToSwarm::Unsubscribe {
+            topic,
+            result_sender,
+        } => {
+            let result = swarm
+                .behaviour_mut()
+                .gossipsub
+                .unsubscribe(&gossipsub::IdentTopic::new(topic));
+            _ = result_sender.send(result);
+        }
+        ToSwarm::Publish {
+            topic,
+            data,
+            result_sender,
+        } => {
+            let result = swarm
+                .behaviour_mut()
+                .gossipsub
+                .publish(gossipsub::IdentTopic::new(topic), data);
+            _ = result_sender.send(result);
+        }
+    }
+}
+
+fn filter_swarm_event(event: SwarmEvent<BehaviourEvent>) -> Option<FromSwarm> {
+    match event {
+        SwarmEvent::Behaviour(BehaviourEvent::Gossipsub(gossipsub::Event::Message {
+            message:
+                gossipsub::Message {
+                    source: Some(peer_id),
+                    topic,
+                    data,
+                    ..
+                },
+            ..
+        })) => Some(FromSwarm::Message {
+            from: peer_id,
+            topic: topic.into_string(),
+            data,
+        }),
+        SwarmEvent::Behaviour(BehaviourEvent::Discovery(
+            discovery::Event::ConnectionEstablished { peer_id, .. },
+        )) => Some(FromSwarm::Discovered { peer_id }),
+        SwarmEvent::Behaviour(BehaviourEvent::Discovery(discovery::Event::ConnectionClosed {
+            peer_id,
+            ..
+        })) => Some(FromSwarm::Expired { peer_id }),
+        _ => None,
+    }
+}
+
 /// Create and configure a swarm which listens to all ports on OS
-pub fn create_swarm(keypair: identity::Keypair) -> alias::AnyResult<Swarm> {
+pub fn create_swarm(
+    keypair: identity::Keypair,
+    from_client: mpsc::Receiver<ToSwarm>,
+) -> alias::AnyResult<Swarm> {
     let mut swarm = SwarmBuilder::with_existing_identity(keypair)
         .with_tokio()
         .with_other_transport(tcp_transport)?
@@ -25,7 +155,7 @@ pub fn create_swarm(keypair: identity::Keypair) -> alias::AnyResult<Swarm> {
 
     // Listen on all interfaces and whatever port the OS assigns
     swarm.listen_on("/ip4/0.0.0.0/tcp/0".parse()?)?;
-    Ok(swarm)
+    Ok(Swarm { swarm, from_client })
 }
 
 mod transport {

src/exo/download/coordinator.py

@@ -1,6 +1,5 @@
 import asyncio
 from dataclasses import dataclass, field
-from random import random
 
 import anyio
 from anyio import current_time
@@ -21,13 +20,9 @@ from exo.shared.types.commands import (
     ForwarderDownloadCommand,
     StartDownload,
 )
-from exo.shared.types.common import NodeId, SessionId, SystemId
+from exo.shared.types.common import NodeId
 from exo.shared.types.events import (
     Event,
-    EventId,
-    # TODO(evan): just for acks, should delete this ASAP
-    GlobalForwarderEvent,
-    LocalForwarderEvent,
     NodeDownloadProgress,
 )
 from exo.shared.types.worker.downloads import (
@@ -38,40 +33,28 @@ from exo.shared.types.worker.downloads import (
     DownloadProgress,
 )
 from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
-from exo.utils.channels import Receiver, Sender, channel
+from exo.utils.channels import Receiver, Sender
 from exo.utils.task_group import TaskGroup
 
 
 @dataclass
 class DownloadCoordinator:
     node_id: NodeId
-    session_id: SessionId
     shard_downloader: ShardDownloader
     download_command_receiver: Receiver[ForwarderDownloadCommand]
-    local_event_sender: Sender[LocalForwarderEvent]
-
-    # ack stuff
-    _global_event_receiver: Receiver[GlobalForwarderEvent]
-    _out_for_delivery: dict[EventId, LocalForwarderEvent] = field(default_factory=dict)
-
+    event_sender: Sender[Event]
     offline: bool = False
 
-    _system_id: SystemId = field(default_factory=SystemId)
-
     # Local state
     download_status: dict[ModelId, DownloadProgress] = field(default_factory=dict)
     active_downloads: dict[ModelId, asyncio.Task[None]] = field(default_factory=dict)
 
-    # Internal event channel for forwarding (initialized in __post_init__)
-    event_sender: Sender[Event] = field(init=False)
-    event_receiver: Receiver[Event] = field(init=False)
     _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
 
     # Per-model throttle for download progress events
     _last_progress_time: dict[ModelId, float] = field(default_factory=dict)
 
     def __post_init__(self) -> None:
-        self.event_sender, self.event_receiver = channel[Event]()
         self.shard_downloader.on_progress(self._download_progress_callback)
 
     def _model_dir(self, model_id: ModelId) -> str:
@@ -123,10 +106,7 @@ class DownloadCoordinator:
         try:
             async with self._tg as tg:
                 tg.start_soon(self._command_processor)
-                tg.start_soon(self._forward_events)
                 tg.start_soon(self._emit_existing_download_progress)
-                tg.start_soon(self._resend_out_for_delivery)
-                tg.start_soon(self._clear_ofd)
         finally:
             for task in self.active_downloads.values():
                 task.cancel()
@@ -134,20 +114,6 @@ class DownloadCoordinator:
     def shutdown(self) -> None:
         self._tg.cancel_tasks()
 
-    # directly copied from worker
-    async def _resend_out_for_delivery(self) -> None:
-        # This can also be massively tightened, we should check events are at least a certain age before resending.
-        # Exponential backoff would also certainly help here.
-        while True:
-            await anyio.sleep(1 + random())
-            for event in self._out_for_delivery.copy().values():
-                await self.local_event_sender.send(event)
-
-    async def _clear_ofd(self) -> None:
-        with self._global_event_receiver as events:
-            async for event in events:
-                self._out_for_delivery.pop(event.event.event_id, None)
-
     async def _command_processor(self) -> None:
         with self.download_command_receiver as commands:
             async for cmd in commands:
@@ -320,23 +286,6 @@ class DownloadCoordinator:
             )
             del self.download_status[model_id]
 
-    async def _forward_events(self) -> None:
-        idx = 0
-        with self.event_receiver as events:
-            async for event in events:
-                fe = LocalForwarderEvent(
-                    origin_idx=idx,
-                    origin=self._system_id,
-                    session=self.session_id,
-                    event=event,
-                )
-                idx += 1
-                logger.debug(
-                    f"DownloadCoordinator published event {idx}: {str(event)[:100]}"
-                )
-                await self.local_event_sender.send(fe)
-                self._out_for_delivery[event.event_id] = fe
-
     async def _emit_existing_download_progress(self) -> None:
         try:
             while True:

src/exo/download/download_utils.py

@@ -314,9 +314,13 @@ async def fetch_file_list_with_cache(
         _fetched_file_lists_this_session.add(cache_key)
         return file_list
     except Exception as e:
+        logger.opt(exception=e).warning(
+            "Ran into exception when fetching file list from HF."
+        )
+
         if await aios.path.exists(cache_file):
             logger.warning(
-                f"No internet and no cached file list for {model_id} - using local file list"
+                f"No cached file list for {model_id} - using local file list"
             )
             async with aiofiles.open(cache_file, "r") as f:
                 return TypeAdapter(list[FileListEntry]).validate_json(await f.read())

src/exo/download/tests/test_coordinator_ack.py

@@ -1,98 +0,0 @@
-from typing import Any
-
-import anyio
-import pytest
-
-from exo.download.coordinator import DownloadCoordinator
-from exo.download.shard_downloader import NoopShardDownloader
-from exo.shared.models.model_cards import ModelCard, ModelTask
-from exo.shared.types.common import ModelId, NodeId, SessionId
-from exo.shared.types.events import (
-    GlobalForwarderEvent,
-    LocalForwarderEvent,
-    NodeDownloadProgress,
-)
-from exo.shared.types.memory import Memory
-from exo.shared.types.worker.downloads import (
-    DownloadPending,
-)
-from exo.shared.types.worker.shards import PipelineShardMetadata
-from exo.utils.channels import channel
-
-# Use the built‑in NoopShardDownloader directly – it already implements the required abstract interface.
-# No additional subclass is needed for this test.
-
-
-@pytest.mark.anyio
-async def test_ack_behaviour():
-    # Create channels (type Any for simplicity)
-    _, command_receiver = channel[Any]()
-    local_sender, _ = channel[Any]()
-    global_sender, global_receiver = channel[Any]()
-
-    # Minimal identifiers
-    node_id = NodeId()
-    session_id = SessionId(master_node_id=node_id, election_clock=0)
-
-    # Create a dummy model card and shard metadata
-    model_id = ModelId("test/model")
-    model_card = ModelCard(
-        model_id=model_id,
-        storage_size=Memory.from_bytes(0),
-        n_layers=1,
-        hidden_size=1,
-        supports_tensor=True,
-        tasks=[ModelTask.TextGeneration],
-    )
-    shard = PipelineShardMetadata(
-        model_card=model_card,
-        device_rank=0,
-        world_size=1,
-        start_layer=0,
-        end_layer=1,
-        n_layers=1,
-    )
-
-    # Instantiate the coordinator with the dummy downloader
-    coord = DownloadCoordinator(
-        node_id=node_id,
-        session_id=session_id,
-        shard_downloader=NoopShardDownloader(),
-        download_command_receiver=command_receiver,
-        local_event_sender=local_sender,
-        _global_event_receiver=global_receiver,
-    )
-
-    async with anyio.create_task_group() as tg:
-        # Start the forwarding and ack‑clearing loops
-        tg.start_soon(coord._forward_events)  # pyright: ignore[reportPrivateUsage]
-        tg.start_soon(coord._clear_ofd)  # pyright: ignore[reportPrivateUsage]
-
-        # Send a pending download progress event via the internal event sender
-        pending = DownloadPending(
-            node_id=node_id,
-            shard_metadata=shard,
-            model_directory="/tmp/model",
-        )
-        await coord.event_sender.send(NodeDownloadProgress(download_progress=pending))
-        # Allow the forwarder to process the event
-        await anyio.sleep(0.1)
-
-        # There should be exactly one entry awaiting ACK
-        assert len(coord._out_for_delivery) == 1  # pyright: ignore[reportPrivateUsage]
-        # Retrieve the stored LocalForwarderEvent
-        stored_fe: LocalForwarderEvent = next(iter(coord._out_for_delivery.values()))  # pyright: ignore[reportPrivateUsage]
-        # Simulate receiving a global ack for this event
-        ack = GlobalForwarderEvent(
-            origin_idx=0,
-            origin=node_id,
-            session=session_id,
-            event=stored_fe.event,
-        )
-        await global_sender.send(ack)
-        # Give the clear‑ofd task a moment to process the ack
-        await anyio.sleep(0.1)
-        # The out‑for‑delivery map should now be empty
-        assert len(coord._out_for_delivery) == 0  # pyright: ignore[reportPrivateUsage]
-        # Cancel background tasks
-        tg.cancel_scope.cancel()

src/exo/main.py

@@ -15,6 +15,7 @@ from exo.download.coordinator import DownloadCoordinator
 from exo.download.impl_shard_downloader import exo_shard_downloader
 from exo.master.api import API  # TODO: should API be in master?
 from exo.master.main import Master
+from exo.routing.event_router import EventRouter
 from exo.routing.router import Router, get_node_id_keypair
 from exo.shared.constants import EXO_LOG
 from exo.shared.election import Election, ElectionResult
@@ -29,6 +30,7 @@ from exo.worker.main import Worker
 @dataclass
 class Node:
     router: Router
+    event_router: EventRouter
     download_coordinator: DownloadCoordinator | None
     worker: Worker | None
     election: Election  # Every node participates in election, as we do want a node to become master even if it isn't a master candidate if no master candidates are present.
@@ -52,6 +54,12 @@ class Node:
         await router.register_topic(topics.ELECTION_MESSAGES)
         await router.register_topic(topics.CONNECTION_MESSAGES)
         await router.register_topic(topics.DOWNLOAD_COMMANDS)
+        event_router = EventRouter(
+            session_id,
+            command_sender=router.sender(topics.COMMANDS),
+            external_outbound=router.sender(topics.LOCAL_EVENTS),
+            external_inbound=router.receiver(topics.GLOBAL_EVENTS),
+        )
 
         logger.info(f"Starting node {node_id}")
 
@@ -59,13 +67,10 @@ class Node:
         if not args.no_downloads:
             download_coordinator = DownloadCoordinator(
                 node_id,
-                session_id,
                 exo_shard_downloader(offline=args.offline),
+                event_sender=event_router.sender(),
                 download_command_receiver=router.receiver(topics.DOWNLOAD_COMMANDS),
-                local_event_sender=router.sender(topics.LOCAL_EVENTS),
                 offline=args.offline,
-                # TODO(evan): remove
-                _global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
             )
         else:
             download_coordinator = None
@@ -73,9 +78,8 @@ class Node:
         if args.spawn_api:
             api = API(
                 node_id,
-                session_id,
                 port=args.api_port,
-                global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
+                event_receiver=event_router.receiver(),
                 command_sender=router.sender(topics.COMMANDS),
                 download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
                 election_receiver=router.receiver(topics.ELECTION_MESSAGES),
@@ -86,9 +90,8 @@ class Node:
         if not args.no_worker:
             worker = Worker(
                 node_id,
-                session_id,
-                global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
-                local_event_sender=router.sender(topics.LOCAL_EVENTS),
+                event_receiver=event_router.receiver(),
+                event_sender=event_router.sender(),
                 command_sender=router.sender(topics.COMMANDS),
                 download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
             )
@@ -99,6 +102,7 @@ class Node:
         master = Master(
             node_id,
             session_id,
+            event_sender=event_router.sender(),
             global_event_sender=router.sender(topics.GLOBAL_EVENTS),
             local_event_receiver=router.receiver(topics.LOCAL_EVENTS),
             command_receiver=router.receiver(topics.COMMANDS),
@@ -121,6 +125,7 @@ class Node:
 
         return cls(
             router,
+            event_router,
             download_coordinator,
             worker,
             election,
@@ -136,6 +141,7 @@ class Node:
             signal.signal(signal.SIGINT, lambda _, __: self.shutdown())
             signal.signal(signal.SIGTERM, lambda _, __: self.shutdown())
             tg.start_soon(self.router.run)
+            tg.start_soon(self.event_router.run)
             tg.start_soon(self.election.run)
             if self.download_coordinator:
                 tg.start_soon(self.download_coordinator.run)
@@ -170,6 +176,17 @@ class Node:
                 # - Shutdown and re-create the worker
                 # - Shut down and re-create the API
 
+                if result.is_new_master:
+                    await anyio.sleep(0)
+                    self.event_router.shutdown()
+                    self.event_router = EventRouter(
+                        result.session_id,
+                        self.router.sender(topics.COMMANDS),
+                        self.router.receiver(topics.GLOBAL_EVENTS),
+                        self.router.sender(topics.LOCAL_EVENTS),
+                    )
+                    self._tg.start_soon(self.event_router.run)
+
                 if (
                     result.session_id.master_node_id == self.node_id
                     and self.master is not None
@@ -183,6 +200,7 @@ class Node:
                     self.master = Master(
                         self.node_id,
                         result.session_id,
+                        event_sender=self.event_router.sender(),
                         global_event_sender=self.router.sender(topics.GLOBAL_EVENTS),
                         local_event_receiver=self.router.receiver(topics.LOCAL_EVENTS),
                         command_receiver=self.router.receiver(topics.COMMANDS),
@@ -205,22 +223,16 @@ class Node:
                         f"Node {result.session_id.master_node_id} elected master"
                     )
                 if result.is_new_master:
-                    await anyio.sleep(0)
                     if self.download_coordinator:
                         self.download_coordinator.shutdown()
                         self.download_coordinator = DownloadCoordinator(
                             self.node_id,
-                            result.session_id,
                             exo_shard_downloader(offline=self.offline),
+                            event_sender=self.event_router.sender(),
                             download_command_receiver=self.router.receiver(
                                 topics.DOWNLOAD_COMMANDS
                             ),
-                            local_event_sender=self.router.sender(topics.LOCAL_EVENTS),
                             offline=self.offline,
-                            # TODO(evan): remove
-                            _global_event_receiver=self.router.receiver(
-                                topics.GLOBAL_EVENTS
-                            ),
                         )
                         self._tg.start_soon(self.download_coordinator.run)
                     if self.worker:
@@ -228,11 +240,8 @@ class Node:
                         # TODO: add profiling etc to resource monitor
                         self.worker = Worker(
                             self.node_id,
-                            result.session_id,
-                            global_event_receiver=self.router.receiver(
-                                topics.GLOBAL_EVENTS
-                            ),
-                            local_event_sender=self.router.sender(topics.LOCAL_EVENTS),
+                            event_receiver=self.event_router.receiver(),
+                            event_sender=self.event_router.sender(),
                             command_sender=self.router.sender(topics.COMMANDS),
                             download_command_sender=self.router.sender(
                                 topics.DOWNLOAD_COMMANDS
@@ -240,7 +249,7 @@ class Node:
                         )
                         self._tg.start_soon(self.worker.run)
                     if self.api:
-                        self.api.reset(result.session_id, result.won_clock)
+                        self.api.reset(result.won_clock, self.event_router.receiver())
                 else:
                     if self.api:
                         self.api.unpause(result.won_clock)

src/exo/master/adapters/chat_completions.py

@@ -26,7 +26,11 @@ from exo.shared.types.chunks import (
     ToolCallChunk,
 )
 from exo.shared.types.common import CommandId
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo.shared.types.text_generation import (
+    InputMessage,
+    TextGenerationTaskParams,
+    resolve_reasoning_params,
+)
 
 
 def chat_request_to_text_generation(
@@ -75,6 +79,10 @@ def chat_request_to_text_generation(
             dumped: dict[str, Any] = msg_copy.model_dump(exclude_none=True)
             chat_template_messages.append(dumped)
 
+    resolved_effort, resolved_thinking = resolve_reasoning_params(
+        request.reasoning_effort, request.enable_thinking
+    )
+
     return TextGenerationTaskParams(
         model=request.model,
         input=input_messages
@@ -89,7 +97,8 @@ def chat_request_to_text_generation(
         seed=request.seed,
         stream=request.stream,
         tools=request.tools,
-        enable_thinking=request.enable_thinking,
+        reasoning_effort=resolved_effort,
+        enable_thinking=resolved_thinking,
         chat_template_messages=chat_template_messages
         if chat_template_messages
         else None,

src/exo/master/adapters/responses.py

@@ -42,7 +42,11 @@ from exo.shared.types.openai_responses import (
     ResponseTextDoneEvent,
     ResponseUsage,
 )
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo.shared.types.text_generation import (
+    InputMessage,
+    TextGenerationTaskParams,
+    resolve_reasoning_params,
+)
 
 
 def _format_sse(event: ResponsesStreamEvent) -> str:
@@ -119,6 +123,11 @@ def responses_request_to_text_generation(
         )
         built_chat_template = chat_template_messages if chat_template_messages else None
 
+    effort_from_reasoning = request.reasoning.effort if request.reasoning else None
+    resolved_effort, resolved_thinking = resolve_reasoning_params(
+        effort_from_reasoning, request.enable_thinking
+    )
+
     return TextGenerationTaskParams(
         model=request.model,
         input=input_value,
@@ -132,6 +141,8 @@ def responses_request_to_text_generation(
         stop=request.stop,
         seed=request.seed,
         chat_template_messages=built_chat_template or request.chat_template_messages,
+        reasoning_effort=resolved_effort,
+        enable_thinking=resolved_thinking,
     )
 
 

src/exo/master/api.py

@@ -73,6 +73,7 @@ from exo.shared.types.api import (
     BenchChatCompletionResponse,
     BenchImageGenerationResponse,
     BenchImageGenerationTaskParams,
+    CancelCommandResponse,
     ChatCompletionChoice,
     ChatCompletionMessage,
     ChatCompletionRequest,
@@ -81,6 +82,8 @@ from exo.shared.types.api import (
     CreateInstanceResponse,
     DeleteDownloadResponse,
     DeleteInstanceResponse,
+    DeleteTracesRequest,
+    DeleteTracesResponse,
     ErrorInfo,
     ErrorResponse,
     FinishReason,
@@ -140,11 +143,10 @@ from exo.shared.types.commands import (
     TaskFinished,
     TextGeneration,
 )
-from exo.shared.types.common import CommandId, Id, NodeId, SessionId, SystemId
+from exo.shared.types.common import CommandId, Id, NodeId, SystemId
 from exo.shared.types.events import (
     ChunkGenerated,
     Event,
-    GlobalForwarderEvent,
     IndexedEvent,
     TracesMerged,
 )
@@ -172,7 +174,6 @@ from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding
 from exo.utils.banner import print_startup_banner
 from exo.utils.channels import Receiver, Sender, channel
-from exo.utils.event_buffer import OrderedBuffer
 from exo.utils.task_group import TaskGroup
 
 _API_EVENT_LOG_DIR = EXO_EVENT_LOG_DIR / "api"
@@ -196,10 +197,9 @@ class API:
     def __init__(
         self,
         node_id: NodeId,
-        session_id: SessionId,
         *,
         port: int,
-        global_event_receiver: Receiver[GlobalForwarderEvent],
+        event_receiver: Receiver[IndexedEvent],
         command_sender: Sender[ForwarderCommand],
         download_command_sender: Sender[ForwarderDownloadCommand],
         # This lets us pause the API if an election is running
@@ -210,11 +210,9 @@ class API:
         self._system_id = SystemId()
         self.command_sender = command_sender
         self.download_command_sender = download_command_sender
-        self.global_event_receiver = global_event_receiver
+        self.event_receiver = event_receiver
         self.election_receiver = election_receiver
-        self.event_buffer: OrderedBuffer[Event] = OrderedBuffer[Event]()
         self.node_id: NodeId = node_id
-        self.session_id: SessionId = session_id
         self.last_completed_election: int = 0
         self.port = port
 
@@ -254,17 +252,18 @@ class API:
         self._image_store = ImageStore(EXO_IMAGE_CACHE_DIR)
         self._tg: TaskGroup = TaskGroup()
 
-    def reset(self, new_session_id: SessionId, result_clock: int):
+    def reset(self, result_clock: int, event_receiver: Receiver[IndexedEvent]):
         logger.info("Resetting API State")
         self._event_log.close()
         self._event_log = DiskEventLog(_API_EVENT_LOG_DIR)
         self.state = State()
         self._system_id = SystemId()
-        self.session_id = new_session_id
-        self.event_buffer = OrderedBuffer[Event]()
         self._text_generation_queues = {}
         self._image_generation_queues = {}
         self.unpause(result_clock)
+        self.event_receiver.close()
+        self.event_receiver = event_receiver
+        self._tg.start_soon(self._apply_state)
 
     def unpause(self, result_clock: int):
         logger.info("Unpausing API")
@@ -324,6 +323,7 @@ class API:
         self.app.get("/images/{image_id}")(self.get_image)
         self.app.post("/v1/messages", response_model=None)(self.claude_messages)
         self.app.post("/v1/responses", response_model=None)(self.openai_responses)
+        self.app.post("/v1/cancel/{command_id}")(self.cancel_command)
 
         # Ollama API
         self.app.head("/ollama/")(self.ollama_version)
@@ -344,6 +344,7 @@ class API:
         self.app.post("/download/start")(self.start_download)
         self.app.delete("/download/{node_id}/{model_id:path}")(self.delete_download)
         self.app.get("/v1/traces")(self.list_traces)
+        self.app.post("/v1/traces/delete")(self.delete_traces)
         self.app.get("/v1/traces/{task_id}")(self.get_trace)
         self.app.get("/v1/traces/{task_id}/stats")(self.get_trace_stats)
         self.app.get("/v1/traces/{task_id}/raw")(self.get_trace_raw)
@@ -568,6 +569,25 @@ class API:
             instance_id=instance_id,
         )
 
+    async def cancel_command(self, command_id: CommandId) -> CancelCommandResponse:
+        """Cancel an active command by closing its stream and notifying workers."""
+        sender = self._text_generation_queues.get(
+            command_id
+        ) or self._image_generation_queues.get(command_id)
+        if sender is None:
+            raise HTTPException(
+                status_code=404,
+                detail="Command not found or already completed",
+            )
+
+        await self._send(TaskCancelled(cancelled_command_id=command_id))
+        sender.close()
+
+        return CancelCommandResponse(
+            message="Command cancelled.",
+            command_id=command_id,
+        )
+
     async def _token_chunk_stream(
         self, command_id: CommandId
     ) -> AsyncGenerator[
@@ -1606,7 +1626,7 @@ class API:
         finally:
             self._event_log.close()
             self.command_sender.close()
-            self.global_event_receiver.close()
+            self.event_receiver.close()
 
     async def run_api(self, ev: anyio.Event):
         cfg = Config()
@@ -1623,38 +1643,31 @@ class API:
             )
 
     async def _apply_state(self):
-        with self.global_event_receiver as events:
-            async for f_event in events:
-                if f_event.session != self.session_id:
-                    continue
-                if f_event.origin != self.session_id.master_node_id:
-                    continue
-                self.event_buffer.ingest(f_event.origin_idx, f_event.event)
-                for idx, event in self.event_buffer.drain_indexed():
-                    self._event_log.append(event)
-                    self.state = apply(self.state, IndexedEvent(event=event, idx=idx))
+        with self.event_receiver as events:
+            async for i_event in events:
+                self._event_log.append(i_event.event)
+                self.state = apply(self.state, i_event)
+                event = i_event.event
 
-                    if isinstance(event, ChunkGenerated):
-                        if queue := self._image_generation_queues.get(
-                            event.command_id, None
-                        ):
-                            assert isinstance(event.chunk, ImageChunk)
-                            try:
-                                await queue.send(event.chunk)
-                            except BrokenResourceError:
-                                self._image_generation_queues.pop(
-                                    event.command_id, None
-                                )
-                        if queue := self._text_generation_queues.get(
-                            event.command_id, None
-                        ):
-                            assert not isinstance(event.chunk, ImageChunk)
-                            try:
-                                await queue.send(event.chunk)
-                            except BrokenResourceError:
-                                self._text_generation_queues.pop(event.command_id, None)
-                    if isinstance(event, TracesMerged):
-                        self._save_merged_trace(event)
+                if isinstance(event, ChunkGenerated):
+                    if queue := self._image_generation_queues.get(
+                        event.command_id, None
+                    ):
+                        assert isinstance(event.chunk, ImageChunk)
+                        try:
+                            await queue.send(event.chunk)
+                        except BrokenResourceError:
+                            self._image_generation_queues.pop(event.command_id, None)
+                    if queue := self._text_generation_queues.get(
+                        event.command_id, None
+                    ):
+                        assert not isinstance(event.chunk, ImageChunk)
+                        try:
+                            await queue.send(event.chunk)
+                        except BrokenResourceError:
+                            self._text_generation_queues.pop(event.command_id, None)
+                if isinstance(event, TracesMerged):
+                    self._save_merged_trace(event)
 
     def _save_merged_trace(self, event: TracesMerged) -> None:
         traces = [
@@ -1718,8 +1731,12 @@ class API:
         await self._send_download(command)
         return DeleteDownloadResponse(command_id=command.command_id)
 
-    def _get_trace_path(self, task_id: str) -> Path:
-        return EXO_TRACING_CACHE_DIR / f"trace_{task_id}.json"
+    @staticmethod
+    def _get_trace_path(task_id: str) -> Path:
+        trace_path = EXO_TRACING_CACHE_DIR / f"trace_{task_id}.json"
+        if not trace_path.resolve().is_relative_to(EXO_TRACING_CACHE_DIR.resolve()):
+            raise HTTPException(status_code=400, detail=f"Invalid task ID: {task_id}")
+        return trace_path
 
     async def list_traces(self) -> TraceListResponse:
         traces: list[TraceListItem] = []
@@ -1818,6 +1835,18 @@ class API:
             filename=f"trace_{task_id}.json",
         )
 
+    async def delete_traces(self, request: DeleteTracesRequest) -> DeleteTracesResponse:
+        deleted: list[str] = []
+        not_found: list[str] = []
+        for task_id in request.task_ids:
+            trace_path = self._get_trace_path(task_id)
+            if trace_path.exists():
+                trace_path.unlink()
+                deleted.append(task_id)
+            else:
+                not_found.append(task_id)
+        return DeleteTracesResponse(deleted=deleted, not_found=not_found)
+
     async def get_onboarding(self) -> JSONResponse:
         return JSONResponse({"completed": ONBOARDING_COMPLETE_FILE.exists()})
 

src/exo/master/main.py

@@ -60,7 +60,7 @@ from exo.shared.types.tasks import (
     TextGeneration as TextGenerationTask,
 )
 from exo.shared.types.worker.instances import InstanceId
-from exo.utils.channels import Receiver, Sender, channel
+from exo.utils.channels import Receiver, Sender
 from exo.utils.event_buffer import MultiSourceBuffer
 from exo.utils.task_group import TaskGroup
 
@@ -72,25 +72,21 @@ class Master:
         session_id: SessionId,
         *,
         command_receiver: Receiver[ForwarderCommand],
+        event_sender: Sender[Event],
         local_event_receiver: Receiver[LocalForwarderEvent],
         global_event_sender: Sender[GlobalForwarderEvent],
         download_command_sender: Sender[ForwarderDownloadCommand],
     ):
-        self.state = State()
-        self._tg: TaskGroup = TaskGroup()
         self.node_id = node_id
         self.session_id = session_id
+        self.state = State()
+        self._tg: TaskGroup = TaskGroup()
         self.command_task_mapping: dict[CommandId, TaskId] = {}
         self.command_receiver = command_receiver
         self.local_event_receiver = local_event_receiver
         self.global_event_sender = global_event_sender
         self.download_command_sender = download_command_sender
-        send, recv = channel[Event]()
-        self.event_sender: Sender[Event] = send
-        self._loopback_event_receiver: Receiver[Event] = recv
-        self._loopback_event_sender: Sender[LocalForwarderEvent] = (
-            local_event_receiver.clone_sender()
-        )
+        self.event_sender = event_sender
         self._system_id = SystemId()
         self._multi_buffer = MultiSourceBuffer[SystemId, Event]()
         self._event_log = DiskEventLog(EXO_EVENT_LOG_DIR / "master")
@@ -104,15 +100,12 @@ class Master:
             async with self._tg as tg:
                 tg.start_soon(self._event_processor)
                 tg.start_soon(self._command_processor)
-                tg.start_soon(self._loopback_processor)
                 tg.start_soon(self._plan)
         finally:
             self._event_log.close()
             self.global_event_sender.close()
             self.local_event_receiver.close()
             self.command_receiver.close()
-            self._loopback_event_sender.close()
-            self._loopback_event_receiver.close()
 
     async def shutdown(self):
         logger.info("Stopping Master")
@@ -335,17 +328,22 @@ class Master:
                                         task_id=task_id,
                                     )
                                 )
-                        case TaskFinished():
-                            generated_events.append(
-                                TaskDeleted(
-                                    task_id=self.command_task_mapping[
-                                        command.finished_command_id
-                                    ]
+                            else:
+                                logger.warning(
+                                    f"Nonexistent command {command.cancelled_command_id} cancelled"
                                 )
-                            )
-                            self.command_task_mapping.pop(
-                                command.finished_command_id, None
-                            )
+                        case TaskFinished():
+                            if (
+                                task_id := self.command_task_mapping.pop(
+                                    command.finished_command_id, None
+                                )
+                            ) is not None:
+                                generated_events.append(TaskDeleted(task_id=task_id))
+                            else:
+                                logger.warning(
+                                    f"Finished command {command.finished_command_id} finished"
+                                )
+
                         case RequestEventLog():
                             # We should just be able to send everything, since other buffers will ignore old messages
                             # rate limit to 1000 at a time
@@ -409,22 +407,6 @@ class Master:
                     self._event_log.append(event)
                     await self._send_event(indexed)
 
-    async def _loopback_processor(self) -> None:
-        # this would ideally not be necessary.
-        # this is WAY less hacky than how I was working around this before
-        local_index = 0
-        with self._loopback_event_receiver as events:
-            async for event in events:
-                await self._loopback_event_sender.send(
-                    LocalForwarderEvent(
-                        origin=self._system_id,
-                        origin_idx=local_index,
-                        session=self.session_id,
-                        event=event,
-                    )
-                )
-                local_index += 1
-
     # This function is re-entrant, take care!
     async def _send_event(self, event: IndexedEvent):
         # Convenience method since this line is ugly

src/exo/master/tests/test_cancel_command.py

@@ -0,0 +1,77 @@
+# pyright: reportUnusedFunction=false, reportAny=false
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock
+
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from exo.shared.types.common import CommandId
+
+
+def _make_api() -> Any:
+    """Create a minimal API instance with cancel route and error handler."""
+    from exo.master.api import API
+
+    app = FastAPI()
+    api = object.__new__(API)
+    api.app = app
+    api._text_generation_queues = {}  # pyright: ignore[reportPrivateUsage]
+    api._image_generation_queues = {}  # pyright: ignore[reportPrivateUsage]
+    api._send = AsyncMock()  # pyright: ignore[reportPrivateUsage]
+    api._setup_exception_handlers()  # pyright: ignore[reportPrivateUsage]
+    app.post("/v1/cancel/{command_id}")(api.cancel_command)
+    return api
+
+
+def test_cancel_nonexistent_command_returns_404() -> None:
+    """Cancel for an unknown command_id returns 404 in OpenAI error format."""
+    api = _make_api()
+    client = TestClient(api.app)
+
+    response = client.post("/v1/cancel/nonexistent-id")
+    assert response.status_code == 404
+    data: dict[str, Any] = response.json()
+    assert "error" in data
+    assert data["error"]["message"] == "Command not found or already completed"
+    assert data["error"]["type"] == "Not Found"
+    assert data["error"]["code"] == 404
+
+
+def test_cancel_active_text_generation() -> None:
+    """Cancel an active text generation command: returns 200, sender.close() called."""
+    api = _make_api()
+    client = TestClient(api.app)
+
+    cid = CommandId("text-cmd-123")
+    sender = MagicMock()
+    api._text_generation_queues[cid] = sender
+
+    response = client.post(f"/v1/cancel/{cid}")
+    assert response.status_code == 200
+    data: dict[str, Any] = response.json()
+    assert data["message"] == "Command cancelled."
+    assert data["command_id"] == str(cid)
+    sender.close.assert_called_once()
+    api._send.assert_called_once()
+    task_cancelled = api._send.call_args[0][0]
+    assert task_cancelled.cancelled_command_id == cid
+
+
+def test_cancel_active_image_generation() -> None:
+    """Cancel an active image generation command: returns 200, sender.close() called."""
+    api = _make_api()
+    client = TestClient(api.app)
+
+    cid = CommandId("img-cmd-456")
+    sender = MagicMock()
+    api._image_generation_queues[cid] = sender
+
+    response = client.post(f"/v1/cancel/{cid}")
+    assert response.status_code == 200
+    data: dict[str, Any] = response.json()
+    assert data["message"] == "Command cancelled."
+    assert data["command_id"] == str(cid)
+    sender.close.assert_called_once()
+    api._send.assert_called_once()
+    task_cancelled = api._send.call_args[0][0]
+    assert task_cancelled.cancelled_command_id == cid

src/exo/master/tests/test_master.py

@@ -17,6 +17,7 @@ from exo.shared.types.commands import (
 )
 from exo.shared.types.common import ModelId, NodeId, SessionId, SystemId
 from exo.shared.types.events import (
+    Event,
     GlobalForwarderEvent,
     IndexedEvent,
     InstanceCreated,
@@ -50,6 +51,22 @@ async def test_master():
     command_sender, co_receiver = channel[ForwarderCommand]()
     local_event_sender, le_receiver = channel[LocalForwarderEvent]()
     fcds, _fcdr = channel[ForwarderDownloadCommand]()
+    ev_send, ev_recv = channel[Event]()
+
+    async def mock_event_router():
+        idx = 0
+        sid = SystemId()
+        with ev_recv as master_events:
+            async for event in master_events:
+                await local_event_sender.send(
+                    LocalForwarderEvent(
+                        origin=sid,
+                        origin_idx=idx,
+                        session=session_id,
+                        event=event,
+                    )
+                )
+                idx += 1
 
     all_events: list[IndexedEvent] = []
 
@@ -67,6 +84,7 @@ async def test_master():
     master = Master(
         node_id,
         session_id,
+        event_sender=ev_send,
         global_event_sender=ge_sender,
         local_event_receiver=le_receiver,
         command_receiver=co_receiver,
@@ -75,6 +93,7 @@ async def test_master():
     logger.info("run the master")
     async with anyio.create_task_group() as tg:
         tg.start_soon(master.run)
+        tg.start_soon(mock_event_router)
 
         # inject a NodeGatheredInfo event
         logger.info("inject a NodeGatheredInfo event")
@@ -197,4 +216,5 @@ async def test_master():
             input=[InputMessage(role="user", content="Hello, how are you?")],
         )
 
+        ev_send.close()
         await master.shutdown()

src/exo/routing/connection_message.py

@@ -1,6 +1,4 @@
-from enum import Enum
-
-from exo_pyo3_bindings import ConnectionUpdate, ConnectionUpdateType
+from exo_pyo3_bindings import PyFromSwarm
 
 from exo.shared.types.common import NodeId
 from exo.utils.pydantic_ext import CamelCaseModel
@@ -8,30 +6,10 @@ from exo.utils.pydantic_ext import CamelCaseModel
 """Serialisable types for Connection Updates/Messages"""
 
 
-class ConnectionMessageType(Enum):
-    Connected = 0
-    Disconnected = 1
-
-    @staticmethod
-    def from_update_type(update_type: ConnectionUpdateType):
-        match update_type:
-            case ConnectionUpdateType.Connected:
-                return ConnectionMessageType.Connected
-            case ConnectionUpdateType.Disconnected:
-                return ConnectionMessageType.Disconnected
-
-
 class ConnectionMessage(CamelCaseModel):
     node_id: NodeId
-    connection_type: ConnectionMessageType
-    remote_ipv4: str
-    remote_tcp_port: int
+    connected: bool
 
     @classmethod
-    def from_update(cls, update: ConnectionUpdate) -> "ConnectionMessage":
-        return cls(
-            node_id=NodeId(update.peer_id),
-            connection_type=ConnectionMessageType.from_update_type(update.update_type),
-            remote_ipv4=update.remote_ipv4,
-            remote_tcp_port=update.remote_tcp_port,
-        )
+    def from_update(cls, update: PyFromSwarm.Connection) -> "ConnectionMessage":
+        return cls(node_id=NodeId(update.peer_id), connected=update.connected)

src/exo/routing/event_router.py

@@ -0,0 +1,161 @@
+from dataclasses import dataclass, field
+from random import random
+
+import anyio
+from anyio import BrokenResourceError, ClosedResourceError
+from anyio.abc import CancelScope
+from loguru import logger
+
+from exo.shared.types.commands import ForwarderCommand, RequestEventLog
+from exo.shared.types.common import SessionId, SystemId
+from exo.shared.types.events import (
+    Event,
+    EventId,
+    GlobalForwarderEvent,
+    IndexedEvent,
+    LocalForwarderEvent,
+)
+from exo.utils.channels import Receiver, Sender, channel
+from exo.utils.event_buffer import OrderedBuffer
+from exo.utils.task_group import TaskGroup
+
+
+@dataclass
+class EventRouter:
+    session_id: SessionId
+    command_sender: Sender[ForwarderCommand]
+    external_inbound: Receiver[GlobalForwarderEvent]
+    external_outbound: Sender[LocalForwarderEvent]
+    _system_id: SystemId = field(init=False, default_factory=SystemId)
+    internal_outbound: list[Sender[IndexedEvent]] = field(
+        init=False, default_factory=list
+    )
+    event_buffer: OrderedBuffer[Event] = field(
+        init=False, default_factory=OrderedBuffer
+    )
+    out_for_delivery: dict[EventId, tuple[float, LocalForwarderEvent]] = field(
+        init=False, default_factory=dict
+    )
+    _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
+
+    _nack_cancel_scope: CancelScope | None = field(init=False, default=None)
+    _nack_attempts: int = field(init=False, default=0)
+    _nack_base_seconds: float = field(init=False, default=0.5)
+    _nack_cap_seconds: float = field(init=False, default=10.0)
+
+    async def run(self):
+        try:
+            async with self._tg as tg:
+                tg.start_soon(self._run_ext_in)
+                tg.start_soon(self._simple_retry)
+        finally:
+            self.external_outbound.close()
+            for send in self.internal_outbound:
+                send.close()
+
+    # can make this better in future
+    async def _simple_retry(self):
+        while True:
+            await anyio.sleep(1 + random())
+            # list here is a shallow clone for shared mutation
+            for e_id, (time, event) in list(self.out_for_delivery.items()):
+                if anyio.current_time() > time + 5:
+                    self.out_for_delivery[e_id] = (anyio.current_time(), event)
+                    await self.external_outbound.send(event)
+
+    def sender(self) -> Sender[Event]:
+        send, recv = channel[Event]()
+        if self._tg.is_running():
+            self._tg.start_soon(self._ingest, SystemId(), recv)
+        else:
+            self._tg.queue(self._ingest, SystemId(), recv)
+        return send
+
+    def receiver(self) -> Receiver[IndexedEvent]:
+        send, recv = channel[IndexedEvent]()
+        self.internal_outbound.append(send)
+        return recv
+
+    def shutdown(self) -> None:
+        self._tg.cancel_tasks()
+
+    async def _ingest(self, system_id: SystemId, recv: Receiver[Event]):
+        idx = 0
+        with recv as events:
+            async for event in events:
+                f_ev = LocalForwarderEvent(
+                    origin_idx=idx,
+                    origin=system_id,
+                    session=self.session_id,
+                    event=event,
+                )
+                idx += 1
+                await self.external_outbound.send(f_ev)
+                self.out_for_delivery[event.event_id] = (anyio.current_time(), f_ev)
+
+    async def _run_ext_in(self):
+        buf = OrderedBuffer[Event]()
+        with self.external_inbound as events:
+            async for event in events:
+                if event.session != self.session_id:
+                    continue
+                if event.origin != self.session_id.master_node_id:
+                    continue
+
+                buf.ingest(event.origin_idx, event.event)
+                event_id = event.event.event_id
+                if event_id in self.out_for_delivery:
+                    self.out_for_delivery.pop(event_id)
+
+                drained = buf.drain_indexed()
+                if drained:
+                    self._nack_attempts = 0
+                    if self._nack_cancel_scope:
+                        self._nack_cancel_scope.cancel()
+
+                if not drained and (
+                    self._nack_cancel_scope is None
+                    or self._nack_cancel_scope.cancel_called
+                ):
+                    # Request the next index.
+                    self._tg.start_soon(self._nack_request, buf.next_idx_to_release)
+                    continue
+
+                for idx, event in drained:
+                    to_clear = set[int]()
+                    for i, sender in enumerate(self.internal_outbound):
+                        try:
+                            await sender.send(IndexedEvent(idx=idx, event=event))
+                        except (ClosedResourceError, BrokenResourceError):
+                            to_clear.add(i)
+                    for i in sorted(to_clear, reverse=True):
+                        self.internal_outbound.pop(i)
+
+    async def _nack_request(self, since_idx: int) -> None:
+        # We request all events after (and including) the missing index.
+        # This function is started whenever we receive an event that is out of sequence.
+        # It is cancelled as soon as we receiver an event that is in sequence.
+
+        if since_idx < 0:
+            logger.warning(f"Negative value encountered for nack request {since_idx=}")
+            since_idx = 0
+
+        with CancelScope() as scope:
+            self._nack_cancel_scope = scope
+            delay: float = self._nack_base_seconds * (2.0**self._nack_attempts)
+            delay = min(self._nack_cap_seconds, delay)
+            self._nack_attempts += 1
+            try:
+                await anyio.sleep(delay)
+                logger.info(
+                    f"Nack attempt {self._nack_attempts}: Requesting Event Log from {since_idx}"
+                )
+                await self.command_sender.send(
+                    ForwarderCommand(
+                        origin=self._system_id,
+                        command=RequestEventLog(since_idx=since_idx),
+                    )
+                )
+            finally:
+                if self._nack_cancel_scope is scope:
+                    self._nack_cancel_scope = None

src/exo/routing/router.py

@@ -17,6 +17,7 @@ from exo_pyo3_bindings import (
     MessageTooLargeError,
     NetworkingHandle,
     NoPeersSubscribedToTopicError,
+    PyFromSwarm,
 )
 from filelock import FileLock
 from loguru import logger
@@ -121,7 +122,8 @@ class Router:
             send = self.networking_receiver.clone_sender()
         router = TopicRouter[T](topic, send)
         self.topic_routers[topic.topic] = cast(TopicRouter[CamelCaseModel], router)
-        await self._networking_subscribe(str(topic.topic))
+        if self._tg.is_running():
+            await self._networking_subscribe(topic.topic)
 
     def sender[T: CamelCaseModel](self, topic: TypedTopic[T]) -> Sender[T]:
         router = self.topic_routers.get(topic.topic, None)
@@ -152,8 +154,10 @@ class Router:
                     router = self.topic_routers[topic]
                     tg.start_soon(router.run)
                 tg.start_soon(self._networking_recv)
-                tg.start_soon(self._networking_recv_connection_messages)
                 tg.start_soon(self._networking_publish)
+                # subscribe to pending topics
+                for topic in self.topic_routers:
+                    await self._networking_subscribe(topic)
                 # Router only shuts down if you cancel it.
                 await sleep_forever()
         finally:
@@ -176,41 +180,40 @@ class Router:
     async def _networking_recv(self):
         try:
             while True:
-                topic, data = await self._net.gossipsub_recv()
-                logger.trace(f"Received message on {topic} with payload {data}")
-                if topic not in self.topic_routers:
-                    logger.warning(
-                        f"Received message on unknown or inactive topic {topic}"
-                    )
-                    continue
-
-                router = self.topic_routers[topic]
-                await router.publish_bytes(data)
+                from_swarm = await self._net.recv()
+                logger.debug(from_swarm)
+                match from_swarm:
+                    case PyFromSwarm.Message(origin, topic, data):
+                        logger.trace(
+                            f"Received message on {topic} from {origin} with payload {data}"
+                        )
+                        if topic not in self.topic_routers:
+                            logger.warning(
+                                f"Received message on unknown or inactive topic {topic}"
+                            )
+                            continue
+                        router = self.topic_routers[topic]
+                        await router.publish_bytes(data)
+                    case PyFromSwarm.Connection():
+                        message = ConnectionMessage.from_update(from_swarm)
+                        logger.trace(
+                            f"Received message on connection_messages with payload {message}"
+                        )
+                        if CONNECTION_MESSAGES.topic in self.topic_routers:
+                            router = self.topic_routers[CONNECTION_MESSAGES.topic]
+                            assert router.topic.model_type == ConnectionMessage
+                            router = cast(TopicRouter[ConnectionMessage], router)
+                            await router.publish(message)
+                    case _:
+                        logger.critical(
+                            "failed to exhaustively check FromSwarm messages - logic error"
+                        )
         except Exception as exception:
             logger.opt(exception=exception).error(
                 "Gossipsub receive loop terminated unexpectedly"
             )
             raise
 
-    async def _networking_recv_connection_messages(self):
-        try:
-            while True:
-                update = await self._net.connection_update_recv()
-                message = ConnectionMessage.from_update(update)
-                logger.trace(
-                    f"Received message on connection_messages with payload {message}"
-                )
-                if CONNECTION_MESSAGES.topic in self.topic_routers:
-                    router = self.topic_routers[CONNECTION_MESSAGES.topic]
-                    assert router.topic.model_type == ConnectionMessage
-                    router = cast(TopicRouter[ConnectionMessage], router)
-                    await router.publish(message)
-        except Exception as exception:
-            logger.opt(exception=exception).error(
-                "Connection update receive loop terminated unexpectedly"
-            )
-            raise
-
     async def _networking_publish(self):
         with self.networking_receiver as networked_items:
             async for topic, data in networked_items:
@@ -255,6 +258,6 @@ def get_node_id_keypair(
 
         # if no valid credentials, create new ones and persist
         with open(path, "w+b") as f:
-            keypair = Keypair.generate_ed25519()
+            keypair = Keypair.generate()
             f.write(keypair.to_bytes())
             return keypair

src/exo/shared/models/model_cards.py

@@ -190,6 +190,8 @@ class ConfigData(BaseModel):
             ["DeepseekV3ForCausalLM"],
             ["Qwen3NextForCausalLM"],
             ["Qwen3MoeForCausalLM"],
+            ["Qwen3_5MoeForConditionalGeneration"],
+            ["Qwen3_5ForConditionalGeneration"],
             ["MiniMaxM2ForCausalLM"],
             ["LlamaForCausalLM"],
             ["GptOssForCausalLM"],

src/exo/shared/tests/test_election.py

@@ -1,7 +1,7 @@
 import pytest
 from anyio import create_task_group, fail_after, move_on_after
 
-from exo.routing.connection_message import ConnectionMessage, ConnectionMessageType
+from exo.routing.connection_message import ConnectionMessage
 from exo.shared.election import Election, ElectionMessage, ElectionResult
 from exo.shared.types.commands import ForwarderCommand, TestCommand
 from exo.shared.types.common import NodeId, SessionId, SystemId
@@ -327,14 +327,7 @@ async def test_connection_message_triggers_new_round_broadcast() -> None:
             tg.start_soon(election.run)
 
             # Send any connection message object; we close quickly to cancel before result creation
-            await cm_tx.send(
-                ConnectionMessage(
-                    node_id=NodeId(),
-                    connection_type=ConnectionMessageType.Connected,
-                    remote_ipv4="",
-                    remote_tcp_port=0,
-                )
-            )
+            await cm_tx.send(ConnectionMessage(node_id=NodeId(), connected=True))
 
             # Expect a broadcast for the new round at clock=1
             while True:

src/exo/shared/tests/test_resolve_reasoning_params.py

@@ -0,0 +1,30 @@
+import pytest
+
+from exo.shared.types.text_generation import resolve_reasoning_params
+
+
+def test_both_none_returns_none_none() -> None:
+    assert resolve_reasoning_params(None, None) == (None, None)
+
+
+def test_both_set_passes_through_unchanged() -> None:
+    assert resolve_reasoning_params("high", True) == ("high", True)
+    assert resolve_reasoning_params("none", True) == ("none", True)
+    assert resolve_reasoning_params("low", False) == ("low", False)
+
+
+def test_enable_thinking_true_derives_medium() -> None:
+    assert resolve_reasoning_params(None, True) == ("medium", True)
+
+
+def test_enable_thinking_false_derives_none() -> None:
+    assert resolve_reasoning_params(None, False) == ("none", False)
+
+
+def test_reasoning_effort_none_derives_thinking_false() -> None:
+    assert resolve_reasoning_params("none", None) == ("none", False)
+
+
+@pytest.mark.parametrize("effort", ["minimal", "low", "medium", "high", "xhigh"])
+def test_non_none_effort_derives_thinking_true(effort: str) -> None:
+    assert resolve_reasoning_params(effort, None) == (effort, True)  # pyright: ignore[reportArgumentType]

src/exo/shared/types/api.py

@@ -8,6 +8,7 @@ from pydantic import BaseModel, Field, field_validator
 from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.memory import Memory
+from exo.shared.types.text_generation import ReasoningEffort
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding, ShardMetadata
 from exo.utils.pydantic_ext import CamelCaseModel
@@ -198,6 +199,7 @@ class ChatCompletionRequest(BaseModel):
     top_p: float | None = None
     top_k: int | None = None
     tools: list[dict[str, Any]] | None = None
+    reasoning_effort: ReasoningEffort | None = None
     enable_thinking: bool | None = None
     tool_choice: str | dict[str, Any] | None = None
     parallel_tool_calls: bool | None = None
@@ -262,6 +264,11 @@ class DeleteInstanceResponse(BaseModel):
     instance_id: InstanceId
 
 
+class CancelCommandResponse(BaseModel):
+    message: str
+    command_id: CommandId
+
+
 ImageSize = Literal[
     "auto",
     "512x512",
@@ -437,3 +444,12 @@ class TraceListItem(CamelCaseModel):
 
 class TraceListResponse(CamelCaseModel):
     traces: list[TraceListItem]
+
+
+class DeleteTracesRequest(CamelCaseModel):
+    task_ids: list[str]
+
+
+class DeleteTracesResponse(CamelCaseModel):
+    deleted: list[str]
+    not_found: list[str]

src/exo/shared/types/mlx.py

@@ -2,6 +2,8 @@
 
 from collections.abc import Sequence
 
+from mlx import core as mx
+from mlx import nn as nn
 from mlx_lm.models.cache import (
     ArraysCache,
     CacheList,
@@ -14,3 +16,16 @@ from mlx_lm.models.cache import (
 KVCacheType = Sequence[
     KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache | CacheList
 ]
+
+
+# Model is a wrapper function to fix the fact that mlx is not strongly typed in the same way that EXO is.
+# For example - MLX has no guarantee of the interface that nn.Module will expose. But we need a guarantee that it has a __call__() function
+class Model(nn.Module):
+    layers: list[nn.Module]
+
+    def __call__(
+        self,
+        x: mx.array,
+        cache: KVCacheType | None,
+        input_embeddings: mx.array | None = None,
+    ) -> mx.array: ...

src/exo/shared/types/openai_responses.py

@@ -12,6 +12,7 @@ from typing import Any, Literal
 from pydantic import BaseModel, Field
 
 from exo.shared.types.common import ModelId
+from exo.shared.types.text_generation import ReasoningEffort
 
 # Type aliases
 ResponseStatus = Literal["completed", "failed", "in_progress", "incomplete"]
@@ -71,6 +72,13 @@ ResponseInputItem = (
 )
 
 
+class Reasoning(BaseModel, frozen=True):
+    """Reasoning configuration for OpenAI Responses API."""
+
+    effort: ReasoningEffort | None = None
+    summary: Literal["auto", "concise", "detailed"] | None = None
+
+
 class ResponsesRequest(BaseModel, frozen=True):
     """Request body for OpenAI Responses API.
 
@@ -89,8 +97,15 @@ class ResponsesRequest(BaseModel, frozen=True):
     stream: bool = False
     tools: list[dict[str, Any]] | None = None
     metadata: dict[str, str] | None = None
+    reasoning: Reasoning | None = None
 
     # --- exo extensions (not in OpenAI Responses API spec) ---
+    enable_thinking: bool | None = Field(
+        default=None,
+        description="[exo extension] Boolean thinking toggle. Not part of the OpenAI Responses API.",
+        json_schema_extra={"x-exo-extension": True},
+    )
+
     top_k: int | None = Field(
         default=None,
         description="[exo extension] Top-k sampling parameter. Not part of the OpenAI Responses API.",

src/exo/shared/types/text_generation.py

@@ -11,6 +11,29 @@ from pydantic import BaseModel
 from exo.shared.types.common import ModelId
 
 MessageRole = Literal["user", "assistant", "system", "developer"]
+ReasoningEffort = Literal["none", "minimal", "low", "medium", "high", "xhigh"]
+
+
+def resolve_reasoning_params(
+    reasoning_effort: ReasoningEffort | None,
+    enable_thinking: bool | None,
+) -> tuple[ReasoningEffort | None, bool | None]:
+    """
+    enable_thinking=True  -> reasoning_effort="medium"
+    enable_thinking=False -> reasoning_effort="none"
+    reasoning_effort="none" -> enable_thinking=False
+    reasoning_effort=<anything else> -> enable_thinking=True
+    """
+    resolved_effort: ReasoningEffort | None = reasoning_effort
+    resolved_thinking: bool | None = enable_thinking
+
+    if reasoning_effort is None and enable_thinking is not None:
+        resolved_effort = "medium" if enable_thinking else "none"
+
+    if enable_thinking is None and reasoning_effort is not None:
+        resolved_thinking = reasoning_effort != "none"
+
+    return resolved_effort, resolved_thinking
 
 
 class InputMessage(BaseModel, frozen=True):
@@ -40,6 +63,7 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
     stop: str | list[str] | None = None
     seed: int | None = None
     chat_template_messages: list[dict[str, Any]] | None = None
+    reasoning_effort: ReasoningEffort | None = None
     enable_thinking: bool | None = None
     logprobs: bool = False
     top_logprobs: int | None = None

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

@@ -1,17 +0,0 @@
-import mlx.core as mx
-import mlx.nn as nn
-from mlx_lm.models.cache import KVCache
-
-# These are wrapper functions to fix the fact that mlx is not strongly typed in the same way that EXO is.
-# For example - MLX has no guarantee of the interface that nn.Module will expose. But we need a guarantee that it has a __call__() function
-
-
-class Model(nn.Module):
-    layers: list[nn.Module]
-
-    def __call__(
-        self,
-        x: mx.array,
-        cache: list[KVCache] | None,
-        input_embeddings: mx.array | None = None,
-    ) -> mx.array: ...

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

@@ -16,6 +16,7 @@ from mlx.nn.layers.distributed import (
 from mlx_lm.models.base import (
     scaled_dot_product_attention,  # pyright: ignore[reportUnknownVariableType]
 )
+from mlx_lm.models.cache import ArraysCache, KVCache
 from mlx_lm.models.deepseek_v3 import DeepseekV3MLP
 from mlx_lm.models.deepseek_v3 import Model as DeepseekV3Model
 from mlx_lm.models.deepseek_v32 import DeepseekV32MLP
@@ -31,10 +32,19 @@ from mlx_lm.models.llama import Model as LlamaModel
 from mlx_lm.models.minimax import MiniMaxAttention
 from mlx_lm.models.minimax import Model as MiniMaxModel
 from mlx_lm.models.ministral3 import Model as Ministral3Model
+from mlx_lm.models.qwen3_5 import DecoderLayer as Qwen3_5DecoderLayer
+from mlx_lm.models.qwen3_5 import Model as Qwen3_5TextModel
+from mlx_lm.models.qwen3_5 import Qwen3_5TextModel as Qwen3_5TextModelInner
+from mlx_lm.models.qwen3_5 import SparseMoeBlock as Qwen3_5SparseMoeBlock
+from mlx_lm.models.qwen3_5_moe import Model as Qwen3_5MoeModel
 from mlx_lm.models.qwen3_moe import Model as Qwen3MoeModel
 from mlx_lm.models.qwen3_moe import Qwen3MoeDecoderLayer, Qwen3MoeSparseMoeBlock
 from mlx_lm.models.qwen3_next import Model as Qwen3NextModel
-from mlx_lm.models.qwen3_next import Qwen3NextDecoderLayer, Qwen3NextSparseMoeBlock
+from mlx_lm.models.qwen3_next import (
+    Qwen3NextDecoderLayer,
+    Qwen3NextGatedDeltaNet,
+    Qwen3NextSparseMoeBlock,
+)
 from mlx_lm.models.step3p5 import Model as Step35Model
 from mlx_lm.models.step3p5 import Step3p5MLP as Step35MLP
 from mlx_lm.models.step3p5 import Step3p5Model as Step35InnerModel
@@ -49,6 +59,21 @@ TimeoutCallback = Callable[[], None]
 LayerLoadedCallback = Callable[[int, int], None]  # (layers_loaded, total_layers)
 
 
+_pending_prefill_sends: list[tuple[mx.array, int, mx.distributed.Group]] = []
+
+
+def flush_prefill_sends() -> None:
+    for output, dst, group in _pending_prefill_sends:
+        sent = mx.distributed.send(output, dst, group=group)
+        mx.async_eval(sent)
+    _pending_prefill_sends.clear()
+
+
+def clear_prefill_sends() -> None:
+    # Discard pending sends (e.g. on cancellation).
+    _pending_prefill_sends.clear()
+
+
 def eval_with_timeout(
     mlx_item: Any,  # pyright: ignore[reportAny]
     timeout_seconds: float = 60.0,
@@ -150,6 +175,7 @@ class PipelineLastLayer(CustomMlxLayer):
         self.group = group
         self.original_layer_signature = signature(self.original_layer.__call__)
         self.is_prefill: bool = False
+        self.queue_sends: bool = False
 
     def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
         cache = self.original_layer_signature.bind_partial(
@@ -163,16 +189,22 @@ class PipelineLastLayer(CustomMlxLayer):
         mx.eval(output)
 
         if self.r != self.s - 1:
-            output = mx.distributed.send(
-                output, (self.r + 1) % self.s, group=self.group
-            )
+            if self.queue_sends:
+                _pending_prefill_sends.append(
+                    (output, (self.r + 1) % self.s, self.group)
+                )
+            else:
+                output = mx.distributed.send(
+                    output, (self.r + 1) % self.s, group=self.group
+                )
             if cache is not None:
                 # CacheList (used by MLA models like DeepSeekV32, GLM MoE DSA)
                 # doesn't have .keys directly; access via first sub-cache.
                 _cache = cache[0] if hasattr(cache, "caches") else cache  # type: ignore
-                _cache.keys = mx.depends(_cache.keys, output)  # type: ignore
+                if hasattr(_cache, "keys"):  # pyright: ignore[reportAny]
+                    _cache.keys = mx.depends(_cache.keys, output)  # type: ignore
             mx.eval(output)
-            if cache is not None:
+            if cache is not None and hasattr(_cache, "keys"):  # type: ignore
                 mx.eval(_cache.keys)  # type: ignore
 
         if not self.is_prefill:
@@ -190,6 +222,12 @@ def set_pipeline_prefill(model: nn.Module, is_prefill: bool) -> None:
             layer.is_prefill = is_prefill
 
 
+def set_pipeline_queue_sends(model: nn.Module, queue_sends: bool) -> None:
+    for layer in model.layers:  # type: ignore
+        if isinstance(layer, PipelineLastLayer):
+            layer.queue_sends = queue_sends
+
+
 def get_inner_model(model: nn.Module) -> nn.Module:
     inner = getattr(model, "model", None)
     if isinstance(inner, nn.Module):
@@ -221,6 +259,32 @@ def get_layers(inner_model_instance: nn.Module) -> list[_LayerCallable]:
     return layers
 
 
+def _patch_qwen35_cache(
+    model: Qwen3_5TextModel,
+    fa_idx: int,
+    has_full_attn: bool,
+    ssm_idx: int,
+    has_linear: bool,
+) -> None:
+    # Hacks to make make_mask happy.
+    original = model.make_cache
+
+    def patched() -> list[ArraysCache | KVCache]:
+        cache: list[ArraysCache | KVCache] = original()
+        if not has_full_attn:
+            entry = cache[fa_idx]
+            orig_make_mask = entry.make_mask
+            entry.make_mask = lambda n, **_kw: orig_make_mask(n)  # type: ignore
+        if not has_linear:
+            orig_ssm_make_mask = cache[ssm_idx].make_mask
+            cache[ssm_idx].make_mask = (  # type: ignore
+                lambda n, **kw: orig_ssm_make_mask(n, **kw) if kw else None  # type: ignore
+            )
+        return cache
+
+    model.make_cache = patched
+
+
 def pipeline_auto_parallel(
     model: nn.Module,
     group: mx.distributed.Group,
@@ -291,6 +355,24 @@ def pipeline_auto_parallel(
         inner_model_instance._swa_idx = 0 if not sliding_layers else sliding_layers[0]
         inner_model_instance._full_idx = 0 if not full_layers else full_layers[0]
 
+    if isinstance(inner_model_instance, Qwen3_5TextModelInner):
+        full_attn_layers = [
+            i for i, layer in enumerate(layers) if not getattr(layer, "is_linear", True)
+        ]
+        linear_layers = [
+            i for i, layer in enumerate(layers) if getattr(layer, "is_linear", False)
+        ]
+        inner_model_instance.fa_idx = full_attn_layers[0] if full_attn_layers else 0
+        inner_model_instance.ssm_idx = linear_layers[0] if linear_layers else 0
+        if not full_attn_layers or not linear_layers:
+            _patch_qwen35_cache(
+                cast(Qwen3_5TextModel, model),
+                fa_idx=inner_model_instance.fa_idx,
+                has_full_attn=bool(full_attn_layers),
+                ssm_idx=inner_model_instance.ssm_idx,
+                has_linear=bool(linear_layers),
+            )
+
     _set_layers(model, layers)
 
     assert isinstance(layers, list), (
@@ -320,7 +402,8 @@ def patch_pipeline_model[T](model: T, group: mx.distributed.Group) -> T:
         if cache is not None:
             last = cache[-1]  # type: ignore
             dep_cache = last[0] if hasattr(last, "caches") else last  # type: ignore
-            dep_cache.keys = mx.depends(dep_cache.keys, logits)  # type: ignore
+            if hasattr(dep_cache, "keys") and dep_cache.keys is not None:  # type: ignore
+                dep_cache.keys = mx.depends(dep_cache.keys, logits)  # type: ignore
 
         return logits
 
@@ -443,7 +526,9 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
         )
-    elif isinstance(model, (Qwen3MoeModel, Qwen3NextModel)):
+    elif isinstance(
+        model, (Qwen3MoeModel, Qwen3NextModel, Qwen3_5TextModel, Qwen3_5MoeModel)
+    ):
         tensor_parallel_sharding_strategy = QwenShardingStrategy(
             group,
             all_to_sharded_linear,
@@ -838,7 +923,9 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
         on_timeout: TimeoutCallback | None,
         on_layer_loaded: LayerLoadedCallback | None,
     ) -> nn.Module:
-        model = cast(Qwen3MoeModel | Qwen3NextModel, model)
+        model = cast(
+            Qwen3MoeModel | Qwen3NextModel | Qwen3_5TextModel | Qwen3_5MoeModel, model
+        )
         total = len(model.layers)
         for i, layer in enumerate(model.layers):
             eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
@@ -859,16 +946,39 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                 layer.self_attn.n_heads //= self.N
                 layer.self_attn.n_kv_heads //= self.N
             else:
-                assert isinstance(layer, Qwen3NextDecoderLayer)
+                assert isinstance(layer, (Qwen3NextDecoderLayer, Qwen3_5DecoderLayer))
                 if hasattr(layer, "linear_attn"):
                     linear_attn = layer.linear_attn
 
-                    linear_attn.in_proj_qkvz = self.all_to_sharded_linear(
-                        linear_attn.in_proj_qkvz
-                    )
-                    linear_attn.in_proj_ba = self.all_to_sharded_linear(
-                        linear_attn.in_proj_ba
-                    )
+                    if isinstance(linear_attn, Qwen3NextGatedDeltaNet):
+                        # Qwen3-Next: combined projections
+                        linear_attn.in_proj_qkvz = self.all_to_sharded_linear(
+                            linear_attn.in_proj_qkvz
+                        )
+                        linear_attn.in_proj_ba = self.all_to_sharded_linear(
+                            linear_attn.in_proj_ba
+                        )
+                    else:
+                        # Qwen3.5: separate projections
+                        # in_proj_qkv has sections [q(key_dim), k(key_dim), v(value_dim)]
+                        # that must be split section-aware, not as a contiguous block
+                        key_dim = linear_attn.key_dim
+                        value_dim = linear_attn.value_dim
+                        linear_attn.in_proj_qkv = shard_linear(
+                            linear_attn.in_proj_qkv,
+                            "all-to-sharded",
+                            segments=[key_dim, key_dim + key_dim],
+                            group=self.group,
+                        )
+                        linear_attn.in_proj_z = self.all_to_sharded_linear(
+                            linear_attn.in_proj_z
+                        )
+                        linear_attn.in_proj_b = self.all_to_sharded_linear(
+                            linear_attn.in_proj_b
+                        )
+                        linear_attn.in_proj_a = self.all_to_sharded_linear(
+                            linear_attn.in_proj_a
+                        )
                     linear_attn.out_proj = self.sharded_to_all_linear(
                         linear_attn.out_proj
                     )
@@ -930,11 +1040,20 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                     layer.self_attn.num_key_value_heads //= self.N
 
             # Shard the MoE.
-            if isinstance(layer.mlp, (Qwen3MoeSparseMoeBlock, Qwen3NextSparseMoeBlock)):
+            if isinstance(
+                layer.mlp,
+                (
+                    Qwen3MoeSparseMoeBlock,
+                    Qwen3NextSparseMoeBlock,
+                    Qwen3_5SparseMoeBlock,
+                ),
+            ):
                 self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj)
                 self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj)
                 self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj)
-                if isinstance(layer.mlp, Qwen3NextSparseMoeBlock):
+                if isinstance(
+                    layer.mlp, (Qwen3NextSparseMoeBlock, Qwen3_5SparseMoeBlock)
+                ):
                     self.all_to_sharded_linear_in_place(
                         layer.mlp.shared_expert.gate_proj
                     )

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

@@ -13,8 +13,7 @@ from mlx_lm.models.cache import (
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
 from exo.shared.types.memory import Memory
-from exo.shared.types.mlx import KVCacheType
-from exo.worker.engines.mlx import Model
+from exo.shared.types.mlx import KVCacheType, Model
 from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS
 from exo.worker.runner.bootstrap import logger
 
@@ -254,9 +253,9 @@ def trim_cache(
             if snapshot is not None and snapshot.states[i] is not None:
                 cache[i] = deepcopy(snapshot.states[i])  # type: ignore
             else:
-                c.state = [None] * len(c.state)  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]
+                c.state = [None] * len(c.state)
         else:
-            c.trim(num_tokens)  # pyright: ignore[reportUnknownMemberType]
+            c.trim(num_tokens)
 
 
 def encode_prompt(tokenizer: TokenizerWrapper, prompt: str) -> mx.array:

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

@@ -1,10 +1,14 @@
+import functools
 import math
 import time
 from copy import deepcopy
 from typing import Callable, Generator, cast, get_args
 
 import mlx.core as mx
-from mlx_lm.generate import stream_generate
+from mlx_lm.generate import (
+    maybe_quantize_kv_cache,
+    stream_generate,
+)
 from mlx_lm.models.cache import ArraysCache, RotatingKVCache
 from mlx_lm.sample_utils import make_sampler
 from mlx_lm.tokenizer_utils import TokenizerWrapper
@@ -19,13 +23,19 @@ from exo.shared.types.api import (
 )
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
-from exo.shared.types.mlx import KVCacheType
+from exo.shared.types.mlx import KVCacheType, Model
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 from exo.shared.types.worker.runner_response import (
     GenerationResponse,
 )
-from exo.worker.engines.mlx import Model
-from exo.worker.engines.mlx.auto_parallel import set_pipeline_prefill
+from exo.worker.engines.mlx.auto_parallel import (
+    PipelineFirstLayer,
+    PipelineLastLayer,
+    clear_prefill_sends,
+    flush_prefill_sends,
+    set_pipeline_prefill,
+    set_pipeline_queue_sends,
+)
 from exo.worker.engines.mlx.cache import (
     CacheSnapshot,
     KVPrefixCache,
@@ -56,6 +66,130 @@ class PrefillCancelled(BaseException):
     """Raised when prefill is cancelled via the progress callback."""
 
 
+def _has_pipeline_communication_layer(model: Model):
+    for layer in model.layers:
+        if isinstance(layer, (PipelineFirstLayer, PipelineLastLayer)):
+            return True
+    return False
+
+
+def pipeline_parallel_prefill(
+    model: Model,
+    prompt: mx.array,
+    prompt_cache: KVCacheType,
+    prefill_step_size: int,
+    kv_group_size: int | None,
+    kv_bits: int | None,
+    prompt_progress_callback: Callable[[int, int], None],
+    distributed_prompt_progress_callback: Callable[[], None] | None,
+    group: mx.distributed.Group,
+) -> None:
+    """Prefill the KV cache for pipeline parallel with overlapping stages.
+
+    Each rank processes the full prompt through its real cache, offset by leading
+    and trailing dummy iterations.
+
+    Total iterations per rank = N_real_chunks + world_size - 1:
+      - rank r leading dummies  (skip_pipeline_io, throwaway cache)
+      - N_real_chunks real      (pipeline IO active, real cache)
+      - (world_size-1-r) trailing dummies (skip_pipeline_io, throwaway cache)
+
+    e.g.
+    Timeline (2 ranks, 3 chunks of 10240 tokens @ step=4096):
+        iter 0: R0 real[0:4096]     R1 dummy
+        iter 1: R0 real[4096:8192]  R1 real[0:4096]
+        iter 2: R0 real[8192:10240] R1 real[4096:8192]
+        iter 3: R0 dummy            R1 real[8192:10240]
+
+    This function is designed to match mlx_lm's stream_generate exactly in terms of
+    side effects (given the same prefill step size)
+    """
+    prefill_step_size = prefill_step_size // min(4, group.size())
+
+    quantize_cache_fn: Callable[..., None] = functools.partial(
+        maybe_quantize_kv_cache,
+        quantized_kv_start=0,
+        kv_group_size=kv_group_size,
+        kv_bits=kv_bits,
+    )
+
+    _prompt_cache: KVCacheType = prompt_cache
+    rank = group.rank()
+    world_size = group.size()
+
+    # Build list of real prompt chunk sizes
+    total = len(prompt)
+    real_chunk_sizes: list[int] = []
+    remaining = total - 1
+    while remaining:
+        n = min(prefill_step_size, remaining)
+        real_chunk_sizes.append(n)
+        remaining -= n
+    n_real = len(real_chunk_sizes)
+
+    # Each rank does: [rank leading dummies] [N real chunks] [world_size-1-rank trailing dummies]
+    n_leading = rank
+    n_trailing = world_size - 1 - rank
+    n_total = n_leading + n_real + n_trailing
+
+    t_start = time.perf_counter()
+    processed = 0
+    logger.info(
+        f"[R{rank}] Pipeline prefill: {n_real} real + {n_leading} leading + {n_trailing} trailing = {n_total} iterations"
+    )
+    clear_prefill_sends()
+
+    # Initial callback matching generate_step
+    prompt_progress_callback(0, total)
+
+    try:
+        with mx.stream(generation_stream):
+            for _ in range(n_leading):
+                if distributed_prompt_progress_callback is not None:
+                    distributed_prompt_progress_callback()
+
+            for i in range(n_real):
+                chunk_size = real_chunk_sizes[i]
+                model(
+                    prompt[processed : processed + chunk_size][None],
+                    cache=_prompt_cache,
+                )
+                quantize_cache_fn(_prompt_cache)
+                processed += chunk_size
+
+                if distributed_prompt_progress_callback is not None:
+                    distributed_prompt_progress_callback()
+
+                flush_prefill_sends()
+
+                prompt_progress_callback(processed, total)
+
+            for _ in range(n_trailing):
+                if distributed_prompt_progress_callback is not None:
+                    distributed_prompt_progress_callback()
+
+    finally:
+        clear_prefill_sends()
+
+    # Post-loop: process remaining 1 token + add +1 entry to match stream_generate.
+    for _ in range(2):
+        with mx.stream(generation_stream):
+            model(prompt[-1:][None], cache=_prompt_cache)
+            quantize_cache_fn(_prompt_cache)
+        flush_prefill_sends()
+
+    assert _prompt_cache is not None
+    mx.eval([c.state for c in _prompt_cache])  # type: ignore
+
+    # Final callback matching generate_step
+    prompt_progress_callback(total, total)
+
+    logger.info(
+        f"[R{rank}] Prefill: {n_real} real + {n_leading}+{n_trailing} dummy iterations, "
+        f"Processed {processed} tokens in {(time.perf_counter() - t_start) * 1000:.1f}ms"
+    )
+
+
 def prefill(
     model: Model,
     tokenizer: TokenizerWrapper,
@@ -64,6 +198,7 @@ def prefill(
     cache: KVCacheType,
     group: mx.distributed.Group | None,
     on_prefill_progress: Callable[[int, int], None] | None,
+    distributed_prompt_progress_callback: Callable[[], None] | None,
 ) -> tuple[float, int, list[CacheSnapshot]]:
     """Prefill the KV cache with prompt tokens.
 
@@ -95,31 +230,57 @@ def prefill(
         if on_prefill_progress is not None:
             on_prefill_progress(processed, total)
 
+    def combined_progress_callback(processed: int, total: int) -> None:
+        if distributed_prompt_progress_callback is not None:
+            distributed_prompt_progress_callback()
+        progress_callback(processed, total)
+
     set_pipeline_prefill(model, is_prefill=True)
 
     mx_barrier(group)
     logger.info("Starting prefill")
 
-    # Use max_tokens=1 because max_tokens=0 does not work.
-    # We just throw away the generated token - we only care about filling the cache
+    is_pipeline = _has_pipeline_communication_layer(model)
+
+    prefill_step_size = 4096
+
     try:
-        for _ in stream_generate(
-            model=model,
-            tokenizer=tokenizer,
-            prompt=prompt_tokens,
-            max_tokens=1,
-            sampler=sampler,
-            prompt_cache=cache,
-            prefill_step_size=4096,
-            kv_group_size=KV_GROUP_SIZE,
-            kv_bits=KV_BITS,
-            prompt_progress_callback=progress_callback,
-        ):
-            break  # Stop after first iteration - cache is now filled
+        if is_pipeline and num_tokens >= prefill_step_size:
+            set_pipeline_queue_sends(model, queue_sends=True)
+            assert group is not None, "Pipeline prefill requires a distributed group"
+            pipeline_parallel_prefill(
+                model=model,
+                prompt=prompt_tokens,
+                prompt_cache=cache,
+                prefill_step_size=prefill_step_size,
+                kv_group_size=KV_GROUP_SIZE,
+                kv_bits=KV_BITS,
+                prompt_progress_callback=progress_callback,
+                distributed_prompt_progress_callback=distributed_prompt_progress_callback,
+                group=group,
+            )
+        else:
+            # Use max_tokens=1 because max_tokens=0 does not work.
+            # We just throw away the generated token - we only care about filling the cache
+            for _ in stream_generate(
+                model=model,
+                tokenizer=tokenizer,
+                prompt=prompt_tokens,
+                max_tokens=1,
+                sampler=sampler,
+                prompt_cache=cache,
+                prefill_step_size=prefill_step_size,
+                kv_group_size=KV_GROUP_SIZE,
+                kv_bits=KV_BITS,
+                prompt_progress_callback=combined_progress_callback,
+            ):
+                break  # Stop after first iteration - cache is now filled
     except PrefillCancelled:
+        set_pipeline_queue_sends(model, queue_sends=False)
         set_pipeline_prefill(model, is_prefill=False)
         raise
 
+    set_pipeline_queue_sends(model, queue_sends=False)
     set_pipeline_prefill(model, is_prefill=False)
 
     # stream_generate added 1 extra generated token to the cache, so we should trim it.
@@ -132,7 +293,7 @@ def prefill(
                 cache[i] = deepcopy(pre_gen.states[i])  # type: ignore
         else:
             assert not isinstance(c, (ArraysCache, RotatingKVCache))
-            c.trim(2)  # pyright: ignore[reportUnknownMemberType]
+            c.trim(2)
 
     elapsed = time.perf_counter() - start_time
     tokens_per_sec = num_tokens / elapsed if elapsed > 0 else 0.0
@@ -275,6 +436,8 @@ def mlx_generate(
     kv_prefix_cache: KVPrefixCache | None,
     group: mx.distributed.Group | None,
     on_prefill_progress: Callable[[int, int], None] | None = None,
+    distributed_prompt_progress_callback: Callable[[], None] | None = None,
+    on_generation_token: Callable[[], None] | None = None,
 ) -> Generator[GenerationResponse]:
     # Ensure that generation stats only contains peak memory for this generation
     mx.reset_peak_memory()
@@ -336,6 +499,7 @@ def mlx_generate(
         caches,
         group,
         on_prefill_progress,
+        distributed_prompt_progress_callback,
     )
     cache_snapshots: list[CacheSnapshot] | None = ssm_snapshots_list or None
 
@@ -481,6 +645,9 @@ def mlx_generate(
                         full_prompt_tokens, caches, cache_snapshots
                     )
 
+        if on_generation_token is not None:
+            on_generation_token()
+
         yield GenerationResponse(
             text=text,
             token=out.token,

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

@@ -40,6 +40,7 @@ from pydantic import RootModel
 from exo.download.download_utils import build_model_path
 from exo.shared.types.common import Host
 from exo.shared.types.memory import Memory
+from exo.shared.types.mlx import Model
 from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.instances import (
     BoundInstance,
@@ -52,7 +53,6 @@ from exo.shared.types.worker.shards import (
     ShardMetadata,
     TensorShardMetadata,
 )
-from exo.worker.engines.mlx import Model
 from exo.worker.engines.mlx.auto_parallel import (
     LayerLoadedCallback,
     TimeoutCallback,
@@ -318,6 +318,9 @@ def get_eos_token_ids_for_model(model_id: ModelId) -> list[int] | None:
         return [151336, 151329, 151338]
     elif "gpt-oss" in model_id_lower:
         return [200002, 200012]
+    elif "qwen3.5" in model_id_lower or "qwen-3.5" in model_id_lower:
+        # For Qwen3.5: 248046 (<|im_end|>), 248044 (<|endoftext|>)
+        return [248046, 248044]
     return None
 
 
@@ -551,6 +554,8 @@ def apply_chat_template(
         # Jinja ignores unknown variables, so passing both is safe.
         extra_kwargs["enable_thinking"] = task_params.enable_thinking
         extra_kwargs["thinking"] = task_params.enable_thinking
+    if task_params.reasoning_effort is not None:
+        extra_kwargs["reasoning_effort"] = task_params.reasoning_effort
 
     patched_template: str | None = None
     if task_params.tools:

src/exo/worker/main.py

@@ -1,9 +1,8 @@
 from collections import defaultdict
 from datetime import datetime, timezone
-from random import random
 
 import anyio
-from anyio import CancelScope, fail_after
+from anyio import fail_after
 from loguru import logger
 
 from exo.download.download_utils import resolve_model_in_path
@@ -13,17 +12,13 @@ from exo.shared.types.api import ImageEditsTaskParams
 from exo.shared.types.commands import (
     ForwarderCommand,
     ForwarderDownloadCommand,
-    RequestEventLog,
     StartDownload,
 )
-from exo.shared.types.common import CommandId, NodeId, SessionId, SystemId
+from exo.shared.types.common import CommandId, NodeId, SystemId
 from exo.shared.types.events import (
     Event,
-    EventId,
-    GlobalForwarderEvent,
     IndexedEvent,
     InputChunkReceived,
-    LocalForwarderEvent,
     NodeDownloadProgress,
     NodeGatheredInfo,
     TaskCreated,
@@ -46,7 +41,6 @@ from exo.shared.types.topology import Connection, SocketConnection
 from exo.shared.types.worker.downloads import DownloadCompleted
 from exo.shared.types.worker.runners import RunnerId
 from exo.utils.channels import Receiver, Sender, channel
-from exo.utils.event_buffer import OrderedBuffer
 from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer
 from exo.utils.info_gatherer.net_profile import check_reachable
 from exo.utils.keyed_backoff import KeyedBackoff
@@ -59,38 +53,26 @@ class Worker:
     def __init__(
         self,
         node_id: NodeId,
-        session_id: SessionId,
         *,
-        global_event_receiver: Receiver[GlobalForwarderEvent],
-        local_event_sender: Sender[LocalForwarderEvent],
+        event_receiver: Receiver[IndexedEvent],
+        event_sender: Sender[Event],
         # This is for requesting updates. It doesn't need to be a general command sender right now,
         # but I think it's the correct way to be thinking about commands
         command_sender: Sender[ForwarderCommand],
         download_command_sender: Sender[ForwarderDownloadCommand],
     ):
         self.node_id: NodeId = node_id
-        self.session_id: SessionId = session_id
-
-        self.global_event_receiver = global_event_receiver
-        self.local_event_sender = local_event_sender
+        self.event_receiver = event_receiver
+        self.event_sender = event_sender
         self.command_sender = command_sender
         self.download_command_sender = download_command_sender
-        self.event_buffer = OrderedBuffer[Event]()
-        self.out_for_delivery: dict[EventId, LocalForwarderEvent] = {}
 
         self.state: State = State()
         self.runners: dict[RunnerId, RunnerSupervisor] = {}
         self._tg: TaskGroup = TaskGroup()
 
-        self._nack_cancel_scope: CancelScope | None = None
-        self._nack_attempts: int = 0
-        self._nack_base_seconds: float = 0.5
-        self._nack_cap_seconds: float = 10.0
-
         self._system_id = SystemId()
 
-        self.event_sender, self.event_receiver = channel[Event]()
-
         # Buffer for input image chunks (for image editing)
         self.input_chunk_buffer: dict[CommandId, dict[int, str]] = {}
         self.input_chunk_counts: dict[CommandId, int] = {}
@@ -108,14 +90,12 @@ class Worker:
                 tg.start_soon(info_gatherer.run)
                 tg.start_soon(self._forward_info, info_recv)
                 tg.start_soon(self.plan_step)
-                tg.start_soon(self._resend_out_for_delivery)
                 tg.start_soon(self._event_applier)
-                tg.start_soon(self._forward_events)
                 tg.start_soon(self._poll_connection_updates)
         finally:
             # Actual shutdown code - waits for all tasks to complete before executing.
             logger.info("Stopping Worker")
-            self.local_event_sender.close()
+            self.event_sender.close()
             self.command_sender.close()
             self.download_command_sender.close()
             for runner in self.runners.values():
@@ -133,47 +113,22 @@ class Worker:
                 )
 
     async def _event_applier(self):
-        with self.global_event_receiver as events:
-            async for f_event in events:
-                if f_event.session != self.session_id:
-                    continue
-                if f_event.origin != self.session_id.master_node_id:
-                    continue
-                self.event_buffer.ingest(f_event.origin_idx, f_event.event)
-                event_id = f_event.event.event_id
-                if event_id in self.out_for_delivery:
-                    del self.out_for_delivery[event_id]
-
+        with self.event_receiver as events:
+            async for event in events:
                 # 2. for each event, apply it to the state
-                indexed_events = self.event_buffer.drain_indexed()
-                if indexed_events:
-                    self._nack_attempts = 0
+                self.state = apply(self.state, event=event)
+                event = event.event
 
-                if not indexed_events and (
-                    self._nack_cancel_scope is None
-                    or self._nack_cancel_scope.cancel_called
-                ):
-                    # Request the next index.
-                    self._tg.start_soon(
-                        self._nack_request, self.state.last_event_applied_idx + 1
+                # Buffer input image chunks for image editing
+                if isinstance(event, InputChunkReceived):
+                    cmd_id = event.command_id
+                    if cmd_id not in self.input_chunk_buffer:
+                        self.input_chunk_buffer[cmd_id] = {}
+                        self.input_chunk_counts[cmd_id] = event.chunk.total_chunks
+
+                    self.input_chunk_buffer[cmd_id][event.chunk.chunk_index] = (
+                        event.chunk.data
                     )
-                    continue
-                elif indexed_events and self._nack_cancel_scope:
-                    self._nack_cancel_scope.cancel()
-
-                for idx, event in indexed_events:
-                    self.state = apply(self.state, IndexedEvent(idx=idx, event=event))
-
-                    # Buffer input image chunks for image editing
-                    if isinstance(event, InputChunkReceived):
-                        cmd_id = event.command_id
-                        if cmd_id not in self.input_chunk_buffer:
-                            self.input_chunk_buffer[cmd_id] = {}
-                            self.input_chunk_counts[cmd_id] = event.chunk.total_chunks
-
-                        self.input_chunk_buffer[cmd_id][event.chunk.chunk_index] = (
-                            event.chunk.data
-                        )
 
     async def plan_step(self):
         while True:
@@ -325,43 +280,6 @@ class Worker:
                 instance.shard_assignments.node_to_runner[self.node_id]
             ].start_task(task)
 
-    async def _nack_request(self, since_idx: int) -> None:
-        # We request all events after (and including) the missing index.
-        # This function is started whenever we receive an event that is out of sequence.
-        # It is cancelled as soon as we receiver an event that is in sequence.
-
-        if since_idx < 0:
-            logger.warning(f"Negative value encountered for nack request {since_idx=}")
-            since_idx = 0
-
-        with CancelScope() as scope:
-            self._nack_cancel_scope = scope
-            delay: float = self._nack_base_seconds * (2.0**self._nack_attempts)
-            delay = min(self._nack_cap_seconds, delay)
-            self._nack_attempts += 1
-            try:
-                await anyio.sleep(delay)
-                logger.info(
-                    f"Nack attempt {self._nack_attempts}: Requesting Event Log from {since_idx}"
-                )
-                await self.command_sender.send(
-                    ForwarderCommand(
-                        origin=self._system_id,
-                        command=RequestEventLog(since_idx=since_idx),
-                    )
-                )
-            finally:
-                if self._nack_cancel_scope is scope:
-                    self._nack_cancel_scope = None
-
-    async def _resend_out_for_delivery(self) -> None:
-        # This can also be massively tightened, we should check events are at least a certain age before resending.
-        # Exponential backoff would also certainly help here.
-        while True:
-            await anyio.sleep(1 + random())
-            for event in self.out_for_delivery.copy().values():
-                await self.local_event_sender.send(event)
-
     def _create_supervisor(self, task: CreateRunner) -> RunnerSupervisor:
         """Creates and stores a new AssignedRunner with initial downloading status."""
         runner = RunnerSupervisor.create(
@@ -372,21 +290,6 @@ class Worker:
         self._tg.start_soon(runner.run)
         return runner
 
-    async def _forward_events(self) -> None:
-        idx = 0
-        with self.event_receiver as events:
-            async for event in events:
-                fe = LocalForwarderEvent(
-                    origin_idx=idx,
-                    origin=self._system_id,
-                    session=self.session_id,
-                    event=event,
-                )
-                idx += 1
-                logger.debug(f"Worker published event {idx}: {str(event)[:100]}")
-                await self.local_event_sender.send(fe)
-                self.out_for_delivery[event.event_id] = fe
-
     async def _poll_connection_updates(self):
         while True:
             edges = set(

src/exo/worker/plan.py

@@ -297,10 +297,10 @@ def _pending_tasks(
             # the task status _should_ be set to completed by the LAST runner
             # it is currently set by the first
             # this is definitely a hack
-            if task.task_id in runner.completed:
+            if task.task_id in runner.completed or task.task_id in runner.in_progress:
                 continue
 
-            if isinstance(runner.status, RunnerReady) and all(
+            if isinstance(runner.status, (RunnerReady, RunnerRunning)) and all(
                 isinstance(all_runners[global_runner_id], (RunnerReady, RunnerRunning))
                 for global_runner_id in runner.bound_instance.instance.shard_assignments.runner_to_shard
             ):

src/exo/worker/runner/bootstrap.py

@@ -32,11 +32,19 @@ def entrypoint(
     # Import main after setting global logger - this lets us just import logger from this module
     try:
         if bound_instance.is_image_model:
-            from exo.worker.runner.image_models.runner import main
-        else:
-            from exo.worker.runner.llm_inference.runner import main
+            from exo.worker.runner.image_models.runner import Runner as ImageRunner
 
-        main(bound_instance, event_sender, task_receiver, cancel_receiver)
+            runner = ImageRunner(
+                bound_instance, event_sender, task_receiver, cancel_receiver
+            )
+            runner.main()
+        else:
+            from exo.worker.runner.llm_inference.runner import Runner
+
+            runner = Runner(
+                bound_instance, event_sender, task_receiver, cancel_receiver
+            )
+            runner.main()
 
     except ClosedResourceError:
         logger.warning("Runner communication closed unexpectedly")

src/exo/worker/runner/image_models/runner.py

@@ -182,272 +182,266 @@ def _send_image_chunk(
         )
 
 
-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))
+class Runner:
+    def __init__(
+        self,
+        bound_instance: BoundInstance,
+        event_sender: MpSender[Event],
+        task_receiver: MpReceiver[Task],
+        cancel_receiver: MpReceiver[TaskId],
+    ):
+        self.event_sender = event_sender
+        self.task_receiver = task_receiver
+        self.cancel_receiver = cancel_receiver
+        self.bound_instance = bound_instance
 
-    instance, runner_id, shard_metadata = (
-        bound_instance.instance,
-        bound_instance.bound_runner_id,
-        bound_instance.bound_shard,
-    )
-    device_rank = shard_metadata.device_rank
-    logger.info("hello from the runner")
-    if getattr(shard_metadata, "immediate_exception", False):
-        raise Exception("Fake exception - runner failed to spin up.")
-    if timeout := getattr(shard_metadata, "should_timeout", 0):
-        time.sleep(timeout)
+        soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
+        resource.setrlimit(resource.RLIMIT_NOFILE, (min(max(soft, 2048), hard), hard))
 
-    setup_start_time = time.time()
-    cancelled_tasks = set[TaskId]()
+        self.instance, self.runner_id, self.shard_metadata = (
+            bound_instance.instance,
+            bound_instance.bound_runner_id,
+            bound_instance.bound_shard,
+        )
+        self.device_rank = self.shard_metadata.device_rank
 
-    image_model: DistributedImageModel | None = None
-    group = None
+        logger.info("hello from the runner")
+        if getattr(self.shard_metadata, "immediate_exception", False):
+            raise Exception("Fake exception - runner failed to spin up.")
+        if timeout := getattr(self.shard_metadata, "should_timeout", 0):
+            time.sleep(timeout)
 
-    current_status: RunnerStatus = RunnerIdle()
-    logger.info("runner created")
-    event_sender.send(
-        RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
-    )
-    seen = set[TaskId]()
-    with task_receiver as tasks:
-        for task in tasks:
-            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)
+        self.setup_start_time = time.time()
+        self.cancelled_tasks = set[TaskId]()
+
+        self.image_model: DistributedImageModel | None = None
+        self.group = None
+
+        self.current_status: RunnerStatus = RunnerIdle()
+        logger.info("runner created")
+        self.update_status(RunnerIdle())
+        self.seen = set[TaskId]()
+
+    def update_status(self, status: RunnerStatus):
+        self.current_status = status
+        self.event_sender.send(
+            RunnerStatusUpdated(
+                runner_id=self.runner_id, runner_status=self.current_status
             )
-            match task:
-                case ConnectToGroup() if isinstance(
-                    current_status, (RunnerIdle, RunnerFailed)
-                ):
-                    logger.info("runner connecting")
-                    current_status = RunnerConnecting()
-                    event_sender.send(
-                        RunnerStatusUpdated(
-                            runner_id=runner_id, runner_status=current_status
-                        )
-                    )
-                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
-                    group = initialize_mlx(bound_instance)
+        )
 
-                    logger.info("runner connected")
-                    current_status = RunnerConnected()
+    def send_task_status(self, task: Task, status: TaskStatus):
+        self.event_sender.send(
+            TaskStatusUpdated(task_id=task.task_id, task_status=status)
+        )
 
-                # we load the model if it's connected with a group, or idle without a group. we should never tell a model to connect if it doesn't need to
-                case LoadModel() if (
-                    isinstance(current_status, RunnerConnected) and group is not None
-                ) or (isinstance(current_status, RunnerIdle) and group is None):
-                    current_status = RunnerLoading()
-                    logger.info("runner loading")
-                    event_sender.send(
-                        RunnerStatusUpdated(
-                            runner_id=runner_id, runner_status=current_status
-                        )
-                    )
-                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
+    def acknowledge_task(self, task: Task):
+        self.event_sender.send(TaskAcknowledged(task_id=task.task_id))
 
-                    assert (
-                        ModelTask.TextToImage in shard_metadata.model_card.tasks
-                        or ModelTask.ImageToImage in shard_metadata.model_card.tasks
-                    ), f"Incorrect model task(s): {shard_metadata.model_card.tasks}"
-
-                    image_model = initialize_image_model(bound_instance)
-                    current_status = RunnerLoaded()
-                    logger.info("runner loaded")
-
-                case StartWarmup() if isinstance(current_status, RunnerLoaded):
-                    current_status = RunnerWarmingUp()
-                    logger.info("runner warming up")
-                    event_sender.send(
-                        RunnerStatusUpdated(
-                            runner_id=runner_id, runner_status=current_status
-                        )
-                    )
-                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
-
-                    logger.info(f"warming up inference for instance: {instance}")
-
-                    assert image_model
-                    image = warmup_image_generator(model=image_model)
-                    if image is not None:
-                        logger.info(f"warmed up by generating {image.size} image")
-                    else:
-                        logger.info("warmup completed (non-primary node)")
-
-                    logger.info(
-                        f"runner initialized in {time.time() - setup_start_time} seconds"
-                    )
-
-                    current_status = RunnerReady()
-                    logger.info("runner ready")
-
-                case ImageGeneration(
-                    task_params=task_params, command_id=command_id
-                ) if isinstance(current_status, RunnerReady):
-                    assert image_model
-                    logger.info(f"received image generation request: {str(task)[:500]}")
-                    current_status = RunnerRunning()
-                    logger.info("runner running")
-                    event_sender.send(
-                        RunnerStatusUpdated(
-                            runner_id=runner_id, runner_status=current_status
-                        )
-                    )
-                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
-
-                    try:
-                        image_index = 0
-                        for response in generate_image(
-                            model=image_model, task=task_params
-                        ):
-                            is_primary_output = _is_primary_output_node(shard_metadata)
-
-                            if is_primary_output:
-                                match response:
-                                    case PartialImageResponse():
-                                        logger.info(
-                                            f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
-                                        )
-                                        _process_image_response(
-                                            response,
-                                            command_id,
-                                            shard_metadata,
-                                            event_sender,
-                                            image_index,
-                                        )
-                                    case ImageGenerationResponse():
-                                        logger.info("sending final ImageChunk")
-                                        _process_image_response(
-                                            response,
-                                            command_id,
-                                            shard_metadata,
-                                            event_sender,
-                                            image_index,
-                                        )
-                                        image_index += 1
-                    # can we make this more explicit?
-                    except Exception as e:
-                        if _is_primary_output_node(shard_metadata):
-                            event_sender.send(
-                                ChunkGenerated(
-                                    command_id=command_id,
-                                    chunk=ErrorChunk(
-                                        model=shard_metadata.model_card.model_id,
-                                        finish_reason="error",
-                                        error_message=str(e),
-                                    ),
-                                )
-                            )
-                        raise
-                    finally:
-                        _send_traces_if_enabled(event_sender, task.task_id, device_rank)
-
-                    current_status = RunnerReady()
-                    logger.info("runner ready")
-
-                case ImageEdits(task_params=task_params, command_id=command_id) if (
-                    isinstance(current_status, RunnerReady)
-                ):
-                    assert image_model
-                    logger.info(f"received image edits request: {str(task)[:500]}")
-                    current_status = RunnerRunning()
-                    logger.info("runner running")
-                    event_sender.send(
-                        RunnerStatusUpdated(
-                            runner_id=runner_id, runner_status=current_status
-                        )
-                    )
-                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
-
-                    try:
-                        image_index = 0
-                        for response in generate_image(
-                            model=image_model, task=task_params
-                        ):
-                            if _is_primary_output_node(shard_metadata):
-                                match response:
-                                    case PartialImageResponse():
-                                        logger.info(
-                                            f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
-                                        )
-                                        _process_image_response(
-                                            response,
-                                            command_id,
-                                            shard_metadata,
-                                            event_sender,
-                                            image_index,
-                                        )
-                                    case ImageGenerationResponse():
-                                        logger.info("sending final ImageChunk")
-                                        _process_image_response(
-                                            response,
-                                            command_id,
-                                            shard_metadata,
-                                            event_sender,
-                                            image_index,
-                                        )
-                                        image_index += 1
-                    except Exception as e:
-                        if _is_primary_output_node(shard_metadata):
-                            event_sender.send(
-                                ChunkGenerated(
-                                    command_id=command_id,
-                                    chunk=ErrorChunk(
-                                        model=shard_metadata.model_card.model_id,
-                                        finish_reason="error",
-                                        error_message=str(e),
-                                    ),
-                                )
-                            )
-                        raise
-                    finally:
-                        _send_traces_if_enabled(event_sender, task.task_id, device_rank)
-
-                    current_status = RunnerReady()
-                    logger.info("runner ready")
-
-                case Shutdown():
-                    current_status = RunnerShuttingDown()
-                    logger.info("runner shutting down")
-                    if not TYPE_CHECKING:
-                        del image_model, group
-                        mx.clear_cache()
-                        import gc
-
-                        gc.collect()
-
-                    event_sender.send(
-                        RunnerStatusUpdated(
-                            runner_id=runner_id, runner_status=current_status
-                        )
-                    )
-                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
-
-                    current_status = RunnerShutdown()
-                case _:
-                    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
-            )
-            if not was_cancelled:
-                event_sender.send(
-                    TaskStatusUpdated(
-                        task_id=task.task_id, task_status=TaskStatus.Complete
-                    )
+    def main(self):
+        with self.task_receiver as tasks:
+            for task in tasks:
+                if task.task_id in self.seen:
+                    logger.warning("repeat task - potential error")
+                self.seen.add(task.task_id)
+                self.cancelled_tasks.discard(TaskId("CANCEL_CURRENT_TASK"))
+                self.send_task_status(task, TaskStatus.Running)
+                self.handle_task(task)
+                was_cancelled = (task.task_id in self.cancelled_tasks) or (
+                    TaskId("CANCEL_CURRENT_TASK") in self.cancelled_tasks
                 )
-            event_sender.send(
-                RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
-            )
+                if not was_cancelled:
+                    self.send_task_status(task, TaskStatus.Complete)
+                self.update_status(self.current_status)
 
-            if isinstance(current_status, RunnerShutdown):
-                break
+                if isinstance(self.current_status, RunnerShutdown):
+                    break
+
+    def handle_task(self, task: Task):
+        match task:
+            case ConnectToGroup() if isinstance(
+                self.current_status, (RunnerIdle, RunnerFailed)
+            ):
+                logger.info("runner connecting")
+                self.update_status(RunnerConnecting())
+                self.acknowledge_task(task)
+                self.group = initialize_mlx(self.bound_instance)
+
+                logger.info("runner connected")
+                self.current_status = RunnerConnected()
+
+            # we load the model if it's connected with a group, or idle without a group. we should never tell a model to connect if it doesn't need to
+            case LoadModel() if (
+                isinstance(self.current_status, RunnerConnected)
+                and self.group is not None
+            ) or (isinstance(self.current_status, RunnerIdle) and self.group is None):
+                logger.info("runner loading")
+                self.update_status(RunnerLoading())
+                self.acknowledge_task(task)
+
+                assert (
+                    ModelTask.TextToImage in self.shard_metadata.model_card.tasks
+                    or ModelTask.ImageToImage in self.shard_metadata.model_card.tasks
+                ), f"Incorrect model task(s): {self.shard_metadata.model_card.tasks}"
+
+                self.image_model = initialize_image_model(self.bound_instance)
+                self.current_status = RunnerLoaded()
+                logger.info("runner loaded")
+
+            case StartWarmup() if isinstance(self.current_status, RunnerLoaded):
+                logger.info("runner warming up")
+                self.update_status(RunnerWarmingUp())
+                self.acknowledge_task(task)
+
+                logger.info(f"warming up inference for instance: {self.instance}")
+
+                assert self.image_model
+                image = warmup_image_generator(model=self.image_model)
+                if image is not None:
+                    logger.info(f"warmed up by generating {image.size} image")
+                else:
+                    logger.info("warmup completed (non-primary node)")
+
+                logger.info(
+                    f"runner initialized in {time.time() - self.setup_start_time} seconds"
+                )
+
+                self.current_status = RunnerReady()
+                logger.info("runner ready")
+
+            case ImageGeneration(task_params=task_params, command_id=command_id) if (
+                isinstance(self.current_status, RunnerReady)
+            ):
+                assert self.image_model
+                logger.info(f"received image generation request: {str(task)[:500]}")
+                logger.info("runner running")
+                self.update_status(RunnerRunning())
+                self.acknowledge_task(task)
+
+                try:
+                    image_index = 0
+                    for response in generate_image(
+                        model=self.image_model, task=task_params
+                    ):
+                        is_primary_output = _is_primary_output_node(self.shard_metadata)
+
+                        if is_primary_output:
+                            match response:
+                                case PartialImageResponse():
+                                    logger.info(
+                                        f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
+                                    )
+                                    _process_image_response(
+                                        response,
+                                        command_id,
+                                        self.shard_metadata,
+                                        self.event_sender,
+                                        image_index,
+                                    )
+                                case ImageGenerationResponse():
+                                    logger.info("sending final ImageChunk")
+                                    _process_image_response(
+                                        response,
+                                        command_id,
+                                        self.shard_metadata,
+                                        self.event_sender,
+                                        image_index,
+                                    )
+                                    image_index += 1
+                # can we make this more explicit?
+                except Exception as e:
+                    if _is_primary_output_node(self.shard_metadata):
+                        self.event_sender.send(
+                            ChunkGenerated(
+                                command_id=command_id,
+                                chunk=ErrorChunk(
+                                    model=self.shard_metadata.model_card.model_id,
+                                    finish_reason="error",
+                                    error_message=str(e),
+                                ),
+                            )
+                        )
+                    raise
+                finally:
+                    _send_traces_if_enabled(
+                        self.event_sender, task.task_id, self.device_rank
+                    )
+
+                self.current_status = RunnerReady()
+                logger.info("runner ready")
+
+            case ImageEdits(task_params=task_params, command_id=command_id) if (
+                isinstance(self.current_status, RunnerReady)
+            ):
+                assert self.image_model
+                logger.info(f"received image edits request: {str(task)[:500]}")
+                logger.info("runner running")
+                self.update_status(RunnerRunning())
+                self.acknowledge_task(task)
+
+                try:
+                    image_index = 0
+                    for response in generate_image(
+                        model=self.image_model, task=task_params
+                    ):
+                        if _is_primary_output_node(self.shard_metadata):
+                            match response:
+                                case PartialImageResponse():
+                                    logger.info(
+                                        f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
+                                    )
+                                    _process_image_response(
+                                        response,
+                                        command_id,
+                                        self.shard_metadata,
+                                        self.event_sender,
+                                        image_index,
+                                    )
+                                case ImageGenerationResponse():
+                                    logger.info("sending final ImageChunk")
+                                    _process_image_response(
+                                        response,
+                                        command_id,
+                                        self.shard_metadata,
+                                        self.event_sender,
+                                        image_index,
+                                    )
+                                    image_index += 1
+                except Exception as e:
+                    if _is_primary_output_node(self.shard_metadata):
+                        self.event_sender.send(
+                            ChunkGenerated(
+                                command_id=command_id,
+                                chunk=ErrorChunk(
+                                    model=self.shard_metadata.model_card.model_id,
+                                    finish_reason="error",
+                                    error_message=str(e),
+                                ),
+                            )
+                        )
+                    raise
+                finally:
+                    _send_traces_if_enabled(
+                        self.event_sender, task.task_id, self.device_rank
+                    )
+
+                self.current_status = RunnerReady()
+                logger.info("runner ready")
+
+            case Shutdown():
+                logger.info("runner shutting down")
+                if not TYPE_CHECKING:
+                    del self.image_model, self.group
+                    mx.clear_cache()
+                    import gc
+
+                    gc.collect()
+
+                self.update_status(RunnerShuttingDown())
+                self.acknowledge_task(task)
+
+                self.current_status = RunnerShutdown()
+            case _:
+                raise ValueError(
+                    f"Received {task.__class__.__name__} outside of state machine in {self.current_status=}"
+                )

src/exo/worker/runner/llm_inference/batch_generator.py

@@ -0,0 +1,346 @@
+import itertools
+import math
+import time
+from abc import ABC, abstractmethod
+from collections import deque
+from collections.abc import Generator, Iterable
+from dataclasses import dataclass, field
+from typing import cast
+
+import mlx.core as mx
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+
+from exo.shared.types.chunks import ErrorChunk, PrefillProgressChunk
+from exo.shared.types.common import ModelId
+from exo.shared.types.events import ChunkGenerated, Event
+from exo.shared.types.mlx import Model
+from exo.shared.types.tasks import TaskId, TextGeneration
+from exo.shared.types.text_generation import TextGenerationTaskParams
+from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse
+from exo.utils.channels import MpReceiver, MpSender
+from exo.worker.engines.mlx.cache import KVPrefixCache
+from exo.worker.engines.mlx.generator.generate import (
+    PrefillCancelled,
+    mlx_generate,
+    warmup_inference,
+)
+from exo.worker.engines.mlx.utils_mlx import (
+    apply_chat_template,
+    mx_any,
+)
+from exo.worker.runner.bootstrap import logger
+
+from .model_output_parsers import apply_all_parsers
+from .tool_parsers import ToolParser
+
+
+class Cancelled:
+    pass
+
+
+class Finished:
+    pass
+
+
+class GeneratorQueue[T]:
+    def __init__(self):
+        self._q = deque[T]()
+
+    def push(self, t: T):
+        self._q.append(t)
+
+    def gen(self) -> Generator[T | None]:
+        while True:
+            if len(self._q) == 0:
+                yield None
+            else:
+                yield self._q.popleft()
+
+
+class InferenceGenerator(ABC):
+    @abstractmethod
+    def warmup(self) -> None: ...
+
+    @abstractmethod
+    def submit(
+        self,
+        task: TextGeneration,
+    ) -> None: ...
+
+    @abstractmethod
+    def step(
+        self,
+    ) -> Iterable[
+        tuple[TaskId, ToolCallResponse | GenerationResponse | Cancelled | Finished]
+    ]: ...
+
+    @abstractmethod
+    def close(self) -> None: ...
+
+
+EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL"
+EXO_RUNNER_MUST_OOM = "EXO RUNNER MUST OOM"
+EXO_RUNNER_MUST_TIMEOUT = "EXO RUNNER MUST TIMEOUT"
+
+
+def _check_for_debug_prompts(task_params: TextGenerationTaskParams) -> None:
+    """Check for debug prompt triggers in the input."""
+    from exo.worker.engines.mlx.utils_mlx import mlx_force_oom
+
+    if len(task_params.input) == 0:
+        return
+    prompt = task_params.input[0].content
+    if not prompt:
+        return
+    if EXO_RUNNER_MUST_FAIL in prompt:
+        raise Exception("Artificial runner exception - for testing purposes only.")
+    if EXO_RUNNER_MUST_OOM in prompt:
+        mlx_force_oom()
+    if EXO_RUNNER_MUST_TIMEOUT in prompt:
+        time.sleep(100)
+
+
+@dataclass(eq=False)
+class SequentialGenerator(InferenceGenerator):
+    model: Model
+    tokenizer: TokenizerWrapper
+    group: mx.distributed.Group | None
+    kv_prefix_cache: KVPrefixCache | None
+    tool_parser: ToolParser | None
+    model_id: ModelId
+    device_rank: int
+    cancel_receiver: MpReceiver[TaskId]
+    event_sender: MpSender[Event]
+    check_for_cancel_every: int = 50
+
+    _cancelled_tasks: set[TaskId] = field(default_factory=set, init=False)
+    _maybe_queue: list[TextGeneration] = field(default_factory=list, init=False)
+    _queue: deque[TextGeneration] = field(default_factory=deque, init=False)
+    _active: (
+        tuple[
+            TextGeneration,
+            # mlx generator that does work
+            Generator[GenerationResponse],
+            # queue that the 1st generator should push to and 3rd generator should pull from
+            GeneratorQueue[GenerationResponse],
+            # generator to get parsed outputs
+            Generator[GenerationResponse | ToolCallResponse | None],
+        ]
+        | None
+    ) = field(default=None, init=False)
+
+    def warmup(self):
+        logger.info(f"warming up inference for instance: {self.model_id}")
+
+        t = time.monotonic()
+        toks = warmup_inference(
+            model=self.model,
+            tokenizer=self.tokenizer,
+            group=self.group,
+        )
+        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 self.group is not None:
+            self.check_for_cancel_every = int(
+                mx.max(
+                    mx.distributed.all_gather(
+                        mx.array([check_for_cancel_every]),
+                        group=self.group,
+                    )
+                ).item()
+            )
+
+        logger.info(
+            f"runner checking for cancellation every {check_for_cancel_every} tokens"
+        )
+
+    def submit(
+        self,
+        task: TextGeneration,
+    ) -> None:
+        self._cancelled_tasks.discard(TaskId("CANCEL_CURRENT_TASK"))
+        self._maybe_queue.append(task)
+
+    def agree_on_tasks(self) -> None:
+        """Agree between all ranks about the task ordering (some may have received in different order or not at all)."""
+        agreed, different = mx_all_gather_tasks(self._maybe_queue, self.group)
+        self._queue.extend(task for task in self._maybe_queue if task in agreed)
+        self._maybe_queue = [task for task in self._maybe_queue if task in different]
+
+    def step(
+        self,
+    ) -> Iterable[
+        tuple[TaskId, GenerationResponse | ToolCallResponse | Cancelled | Finished]
+    ]:
+        if self._active is None:
+            self.agree_on_tasks()
+
+            if self._queue:
+                self._start_next()
+            else:
+                return map(lambda task: (task, Cancelled()), self._cancelled_tasks)
+
+        assert self._active is not None
+
+        task, mlx_gen, queue, output_generator = self._active
+        response = None
+        try:
+            queue.push(next(mlx_gen))
+            response = next(output_generator)
+        except (StopIteration, PrefillCancelled):
+            response = Finished()
+            self._active = None
+            if self._queue:
+                self._start_next()
+        except Exception as e:
+            self._send_error(task, e)
+            self._active = None
+            raise
+        return itertools.chain(
+            [] if response is None else [(task.task_id, response)],
+            map(lambda task: (task, Cancelled()), self._cancelled_tasks),
+        )
+
+    def _start_next(self) -> None:
+        task = self._queue.popleft()
+        try:
+            mlx_gen = self._build_generator(task)
+        except Exception as e:
+            self._send_error(task, e)
+            raise
+        queue = GeneratorQueue[GenerationResponse]()
+        output_generator = apply_all_parsers(
+            queue.gen(),
+            apply_chat_template(self.tokenizer, task.task_params),
+            self.tool_parser,
+            self.tokenizer,
+            type(self.model),
+            self.model_id,
+            task.task_params.tools,
+        )
+        self._active = (task, mlx_gen, queue, output_generator)
+
+    def _send_error(self, task: TextGeneration, e: Exception) -> None:
+        if self.device_rank == 0:
+            self.event_sender.send(
+                ChunkGenerated(
+                    command_id=task.command_id,
+                    chunk=ErrorChunk(
+                        model=self.model_id,
+                        finish_reason="error",
+                        error_message=str(e),
+                    ),
+                )
+            )
+
+    def _build_generator(self, task: TextGeneration) -> Generator[GenerationResponse]:
+        _check_for_debug_prompts(task.task_params)
+        prompt = apply_chat_template(self.tokenizer, task.task_params)
+
+        def on_prefill_progress(processed: int, total: int) -> None:
+            if self.device_rank == 0:
+                self.event_sender.send(
+                    ChunkGenerated(
+                        command_id=task.command_id,
+                        chunk=PrefillProgressChunk(
+                            model=self.model_id,
+                            processed_tokens=processed,
+                            total_tokens=total,
+                        ),
+                    )
+                )
+
+        def distributed_prompt_progress_callback() -> None:
+            self._cancelled_tasks.update(self.cancel_receiver.collect())
+            want_to_cancel = (task.task_id in self._cancelled_tasks) or (
+                TaskId("CANCEL_CURRENT_TASK") in self._cancelled_tasks
+            )
+            if mx_any(want_to_cancel, self.group):
+                raise PrefillCancelled()
+
+            self.agree_on_tasks()
+
+        tokens_since_cancel_check = self.check_for_cancel_every
+
+        def on_generation_token() -> None:
+            nonlocal tokens_since_cancel_check
+            tokens_since_cancel_check += 1
+            if tokens_since_cancel_check >= self.check_for_cancel_every:
+                tokens_since_cancel_check = 0
+                self._cancelled_tasks.update(self.cancel_receiver.collect())
+                want_to_cancel = (task.task_id in self._cancelled_tasks) or (
+                    TaskId("CANCEL_CURRENT_TASK") in self._cancelled_tasks
+                )
+                if mx_any(want_to_cancel, self.group):
+                    raise PrefillCancelled()
+
+                self.agree_on_tasks()
+
+        return mlx_generate(
+            model=self.model,
+            tokenizer=self.tokenizer,
+            task=task.task_params,
+            prompt=prompt,
+            kv_prefix_cache=self.kv_prefix_cache,
+            on_prefill_progress=on_prefill_progress,
+            distributed_prompt_progress_callback=distributed_prompt_progress_callback,
+            on_generation_token=on_generation_token,
+            group=self.group,
+        )
+
+    def close(self) -> None:
+        del self.model, self.tokenizer, self.group
+
+
+def mx_all_gather_tasks(
+    tasks: list[TextGeneration],
+    group: mx.distributed.Group | None,
+) -> tuple[list[TextGeneration], list[TextGeneration]]:
+    def encode_task_id(task_id: TaskId) -> list[int]:
+        utf8_task_id = task_id.encode()
+        return [
+            int.from_bytes(utf8_task_id[i : i + 1]) for i in range(len(utf8_task_id))
+        ]
+
+    def decode_task_id(encoded_task_id: list[int]) -> TaskId:
+        return TaskId(
+            bytes.decode(b"".join((x).to_bytes(length=1) for x in encoded_task_id))
+        )
+
+    uuid_byte_length = 36
+
+    n_tasks = len(tasks)
+    all_counts = cast(
+        list[int],
+        mx.distributed.all_gather(mx.array([n_tasks]), group=group).tolist(),
+    )
+    max_tasks = max(all_counts)
+    world_size: int = 1 if group is None else group.size()
+
+    if max_tasks == 0:
+        return [], []
+
+    padded = [encode_task_id(task.task_id) for task in tasks] + [
+        [0] * uuid_byte_length
+    ] * (max_tasks - n_tasks)
+    gathered = cast(
+        list[list[list[int]]],
+        mx.distributed.all_gather(mx.array(padded), group=group)
+        .reshape(world_size, max_tasks, -1)
+        .tolist(),
+    )
+    all_task_ids: list[list[TaskId]] = [
+        [decode_task_id(encoded_task_id) for encoded_task_id in rank_tasks[:count]]
+        for rank_tasks, count in zip(gathered, all_counts, strict=True)
+    ]
+
+    agreed_ids: set[TaskId] = set(all_task_ids[0])
+    for rank_tasks in all_task_ids[1:]:
+        agreed_ids &= set(rank_tasks)
+
+    local_tasks = {task.task_id: task for task in tasks}
+    agreed = [local_tasks[tid] for tid in sorted(agreed_ids)]
+    different = [task for task in tasks if task.task_id not in agreed_ids]
+    return agreed, different

src/exo/worker/runner/llm_inference/model_output_parsers.py

@@ -0,0 +1,378 @@
+from collections.abc import Generator
+from functools import cache
+from typing import Any
+
+from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model
+from mlx_lm.models.gpt_oss import Model as GptOssModel
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+from openai_harmony import (  # pyright: ignore[reportMissingTypeStubs]
+    HarmonyEncodingName,
+    HarmonyError,  # pyright: ignore[reportUnknownVariableType]
+    Role,
+    StreamableParser,
+    load_harmony_encoding,
+)
+
+from exo.shared.types.api import ToolCallItem
+from exo.shared.types.common import ModelId
+from exo.shared.types.mlx import Model
+from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse
+from exo.worker.engines.mlx.utils_mlx import (
+    detect_thinking_prompt_suffix,
+)
+from exo.worker.runner.bootstrap import logger
+
+from .tool_parsers import ToolParser
+
+
+@cache
+def get_gpt_oss_encoding():
+    encoding = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS)
+    return encoding
+
+
+def apply_all_parsers(
+    receiver: Generator[GenerationResponse | None],
+    prompt: str,
+    tool_parser: ToolParser | None,
+    tokenizer: TokenizerWrapper,
+    model_type: type[Model],
+    model_id: ModelId,
+    tools: list[dict[str, Any]] | None,
+) -> Generator[GenerationResponse | ToolCallResponse | None]:
+    mlx_generator = receiver
+
+    if tokenizer.has_thinking:
+        mlx_generator = parse_thinking_models(
+            mlx_generator,
+            tokenizer,
+            starts_in_thinking=detect_thinking_prompt_suffix(prompt, tokenizer),
+        )
+
+    if issubclass(model_type, GptOssModel):
+        mlx_generator = parse_gpt_oss(mlx_generator)
+    elif (
+        issubclass(model_type, DeepseekV32Model)
+        and "deepseek" in model_id.normalize().lower()
+    ):
+        mlx_generator = parse_deepseek_v32(mlx_generator)
+    elif tool_parser:
+        mlx_generator = parse_tool_calls(mlx_generator, tool_parser, tools)
+
+    return mlx_generator
+
+
+def parse_gpt_oss(
+    responses: Generator[GenerationResponse | None],
+) -> Generator[GenerationResponse | ToolCallResponse | None]:
+    encoding = get_gpt_oss_encoding()
+    stream = StreamableParser(encoding, role=Role.ASSISTANT)
+    thinking = False
+    current_tool_name: str | None = None
+    tool_arg_parts: list[str] = []
+
+    for response in responses:
+        if response is None:
+            yield None
+            continue
+        try:
+            stream.process(response.token)
+        except HarmonyError:
+            logger.error("Encountered critical Harmony Error, returning early")
+            return
+
+        delta = stream.last_content_delta
+        ch = stream.current_channel
+        recipient = stream.current_recipient
+
+        # Debug: log every token with state
+        logger.debug(
+            f"parse_gpt_oss token={response.token} text={response.text!r} "
+            f"recipient={recipient!r} ch={ch!r} delta={delta!r} "
+            f"state={stream.state} current_tool={current_tool_name!r}"
+        )
+
+        if recipient != current_tool_name:
+            if current_tool_name is not None:
+                prefix = "functions."
+                if current_tool_name.startswith(prefix):
+                    current_tool_name = current_tool_name[len(prefix) :]
+                logger.info(
+                    f"parse_gpt_oss yielding tool call: name={current_tool_name!r}"
+                )
+                yield ToolCallResponse(
+                    tool_calls=[
+                        ToolCallItem(
+                            name=current_tool_name,
+                            arguments="".join(tool_arg_parts).strip(),
+                        )
+                    ],
+                    usage=response.usage,
+                )
+                tool_arg_parts = []
+            current_tool_name = recipient
+
+        # If inside a tool call, accumulate arguments
+        if current_tool_name is not None:
+            if delta:
+                tool_arg_parts.append(delta)
+            continue
+
+        if ch == "analysis" and not thinking:
+            thinking = True
+
+        if ch != "analysis" and thinking:
+            thinking = False
+
+        if delta:
+            yield response.model_copy(update={"text": delta, "is_thinking": thinking})
+
+        if response.finish_reason is not None:
+            yield response
+
+
+def parse_deepseek_v32(
+    responses: Generator[GenerationResponse | None],
+) -> Generator[GenerationResponse | ToolCallResponse | None]:
+    """Parse DeepSeek V3.2 DSML tool calls from the generation stream.
+
+    Uses accumulated-text matching (not per-token marker checks) because
+    DSML markers like <｜DSML｜function_calls> may span multiple tokens.
+    Also handles <think>...</think> blocks for thinking mode.
+    """
+    from exo.worker.engines.mlx.dsml_encoding import (
+        THINKING_END,
+        THINKING_START,
+        TOOL_CALLS_END,
+        TOOL_CALLS_START,
+        parse_dsml_output,
+    )
+
+    accumulated = ""
+    in_tool_call = False
+    thinking = False
+    # Tokens buffered while we detect the start of a DSML block
+    pending_buffer: list[GenerationResponse] = []
+    # Text accumulated during a tool call block
+    tool_call_text = ""
+
+    for response in responses:
+        if response is None:
+            yield None
+            continue
+
+        # ── Handle thinking tags ──
+        if not thinking and THINKING_START in response.text:
+            thinking = True
+            # Yield any text before the <think> tag
+            before = response.text[: response.text.index(THINKING_START)]
+            if before:
+                yield response.model_copy(update={"text": before})
+            continue
+
+        if thinking and THINKING_END in response.text:
+            thinking = False
+            # Yield any text after the </think> tag
+            after = response.text[
+                response.text.index(THINKING_END) + len(THINKING_END) :
+            ]
+            if after:
+                yield response.model_copy(update={"text": after, "is_thinking": False})
+            continue
+
+        if thinking:
+            yield response.model_copy(update={"is_thinking": True})
+            continue
+
+        # ── Handle tool call accumulation ──
+        if in_tool_call:
+            tool_call_text += response.text
+            if TOOL_CALLS_END in tool_call_text:
+                # Parse the accumulated DSML block
+                parsed = parse_dsml_output(tool_call_text)
+                if parsed is not None:
+                    logger.info(f"parsed DSML tool calls: {parsed}")
+                    yield ToolCallResponse(
+                        tool_calls=parsed,
+                        usage=response.usage,
+                        stats=response.stats,
+                    )
+                else:
+                    logger.warning(
+                        f"DSML tool call parsing failed for: {tool_call_text}"
+                    )
+                    yield response.model_copy(update={"text": tool_call_text})
+                in_tool_call = False
+                tool_call_text = ""
+                continue
+
+            # EOS reached before end marker — yield buffered text as-is
+            if response.finish_reason is not None:
+                logger.info("DSML tool call parsing interrupted by EOS")
+                yield response.model_copy(update={"text": tool_call_text})
+                in_tool_call = False
+                tool_call_text = ""
+            continue
+
+        # ── Detect start of tool call block ──
+        accumulated += response.text
+
+        if TOOL_CALLS_START in accumulated:
+            # The start marker might be split across pending_buffer + current token
+            start_idx = accumulated.index(TOOL_CALLS_START)
+            # Yield any pending tokens that are purely before the marker
+            pre_text = accumulated[:start_idx]
+            if pre_text:
+                # Flush pending buffer tokens that contributed text before the marker
+                for buf_resp in pending_buffer:
+                    if pre_text:
+                        chunk = buf_resp.text
+                        if len(chunk) <= len(pre_text):
+                            yield buf_resp
+                            pre_text = pre_text[len(chunk) :]
+                        else:
+                            yield buf_resp.model_copy(update={"text": pre_text})
+                            pre_text = ""
+            pending_buffer = []
+            tool_call_text = accumulated[start_idx:]
+            accumulated = ""
+
+            # Check if the end marker is already present (entire tool call in one token)
+            if TOOL_CALLS_END in tool_call_text:
+                parsed = parse_dsml_output(tool_call_text)
+                if parsed is not None:
+                    logger.info(f"parsed DSML tool calls: {parsed}")
+                    yield ToolCallResponse(
+                        tool_calls=parsed,
+                        usage=response.usage,
+                        stats=response.stats,
+                    )
+                else:
+                    logger.warning(
+                        f"DSML tool call parsing failed for: {tool_call_text}"
+                    )
+                    yield response.model_copy(update={"text": tool_call_text})
+                tool_call_text = ""
+            else:
+                in_tool_call = True
+            continue
+
+        # Check if accumulated text might be the start of a DSML marker
+        # Buffer tokens if we see a partial match at the end
+        if _could_be_dsml_prefix(accumulated):
+            pending_buffer.append(response)
+            continue
+
+        # No partial match — flush all pending tokens and the current one
+        for buf_resp in pending_buffer:
+            yield buf_resp
+        pending_buffer = []
+        accumulated = ""
+        yield response
+
+    # Flush any remaining pending buffer at generator end
+    for buf_resp in pending_buffer:
+        yield buf_resp
+
+
+def _could_be_dsml_prefix(text: str) -> bool:
+    """Check if the end of text could be the start of a DSML function_calls marker.
+
+    We look for suffixes of text that are prefixes of the TOOL_CALLS_START pattern.
+    This allows us to buffer tokens until we can determine if a tool call is starting.
+    """
+    from exo.worker.engines.mlx.dsml_encoding import TOOL_CALLS_START
+
+    # Only check the last portion of text that could overlap with the marker
+    max_check = len(TOOL_CALLS_START)
+    tail = text[-max_check:] if len(text) > max_check else text
+
+    # Check if any suffix of tail is a prefix of TOOL_CALLS_START
+    for i in range(len(tail)):
+        suffix = tail[i:]
+        if TOOL_CALLS_START.startswith(suffix):
+            return True
+    return False
+
+
+def parse_thinking_models(
+    responses: Generator[GenerationResponse | None],
+    tokenizer: TokenizerWrapper,
+    starts_in_thinking: bool = True,
+) -> Generator[GenerationResponse | None]:
+    """Route thinking tokens via is_thinking flag.
+
+    Swallows think tag tokens, sets is_thinking on all others.
+    Always yields tokens with finish_reason to avoid hanging the chunk stream.
+    """
+    in_thinking = starts_in_thinking
+    for response in responses:
+        if response is None:
+            yield None
+            continue
+
+        is_think_tag = (
+            tokenizer.think_end is not None and response.text == tokenizer.think_end
+        ) or (
+            tokenizer.think_start is not None and response.text == tokenizer.think_start
+        )
+
+        if is_think_tag:
+            in_thinking = response.text != tokenizer.think_end
+            # Never swallow finish_reason — the chunk stream needs it to terminate.
+            if response.finish_reason is not None:
+                yield response.model_copy(update={"text": "", "is_thinking": False})
+            continue
+        yield response.model_copy(update={"is_thinking": in_thinking})
+
+
+def parse_tool_calls(
+    responses: Generator[GenerationResponse | None],
+    tool_parser: ToolParser,
+    tools: list[dict[str, Any]] | None,
+) -> Generator[GenerationResponse | ToolCallResponse | None]:
+    in_tool_call = False
+    tool_call_text_parts: list[str] = []
+    for response in responses:
+        if response is None:
+            yield None
+            continue
+        if not in_tool_call and response.text.startswith(tool_parser.start_parsing):
+            in_tool_call = True
+
+        if in_tool_call:
+            tool_call_text_parts.append(response.text)
+            if response.text.endswith(tool_parser.end_parsing):
+                # parse the actual tool calls from the tool call text
+                parsed = tool_parser.parse("".join(tool_call_text_parts).strip(), tools)
+                logger.info(f"parsed {tool_call_text_parts=} into {parsed=}")
+                if parsed is not None:
+                    yield ToolCallResponse(
+                        tool_calls=parsed, usage=response.usage, stats=response.stats
+                    )
+                else:
+                    logger.warning(
+                        f"tool call parsing failed for text {''.join(tool_call_text_parts)}"
+                    )
+                    response.text = "".join(tool_call_text_parts)
+                    yield response
+
+                in_tool_call = False
+                tool_call_text_parts = []
+                continue
+
+            if response.finish_reason is not None:
+                logger.info(
+                    "tool call parsing interrupted, yield partial tool call as text"
+                )
+                response = response.model_copy(
+                    update={
+                        "text": "".join(tool_call_text_parts),
+                        "token": 0,
+                    }
+                )
+                yield response
+
+        else:
+            # fallthrough
+            yield response

src/exo/worker/runner/llm_inference/tool_parsers.py

@@ -1,4 +1,5 @@
 import json
+import math
 from dataclasses import dataclass
 from typing import Any, Callable
 
@@ -9,7 +10,177 @@ from exo.shared.types.api import ToolCallItem
 class ToolParser:
     start_parsing: str
     end_parsing: str
-    parse_tool_calls: Callable[[str], list[ToolCallItem] | None]
+    _inner_parser: Callable[[str], list[ToolCallItem] | None]
+
+    def parse(
+        self, text: str, tools: list[dict[str, Any]] | None
+    ) -> list[ToolCallItem] | None:
+        parsed = self._inner_parser(text)
+        if parsed is None:
+            return None
+        if tools is not None:
+            parsed = _coerce_tool_calls_to_schema(parsed, tools)
+        return parsed
+
+
+def _json_type_matches(value: Any, expected_type: str) -> bool:  # pyright: ignore[reportAny]
+    if expected_type == "object":
+        return isinstance(value, dict)
+    if expected_type == "array":
+        return isinstance(value, list)
+    if expected_type == "string":
+        return isinstance(value, str)
+    if expected_type == "integer":
+        return isinstance(value, int) and not isinstance(value, bool)
+    if expected_type == "number":
+        return (isinstance(value, int) and not isinstance(value, bool)) or isinstance(
+            value, float
+        )
+    if expected_type == "boolean":
+        return isinstance(value, bool)
+    if expected_type == "null":
+        return value is None
+    return False
+
+
+def _coerce_tool_arg_with_schema(value: Any, schema: dict[str, Any]) -> Any:  # pyright: ignore[reportAny]
+    schema_type = schema.get("type")
+
+    if isinstance(schema_type, list):
+        for candidate in schema_type:  # pyright: ignore[reportUnknownVariableType]
+            if not isinstance(candidate, str):
+                continue
+            if candidate == "null" and value is None:
+                return None
+            candidate_schema = {**schema, "type": candidate}
+            coerced = _coerce_tool_arg_with_schema(value, candidate_schema)  # pyright: ignore[reportAny]
+            if _json_type_matches(coerced, candidate):
+                return coerced  # pyright: ignore[reportAny]
+        return value  # pyright: ignore[reportAny]
+
+    if not isinstance(schema_type, str):
+        return value  # pyright: ignore[reportAny]
+
+    if schema_type == "object":
+        parsed = value  # pyright: ignore[reportAny]
+        if isinstance(parsed, str):
+            try:
+                parsed = json.loads(parsed)  # pyright: ignore[reportAny]
+            except Exception:
+                return value  # pyright: ignore[reportAny]
+        if not isinstance(parsed, dict):
+            return value  # pyright: ignore[reportAny]
+        properties = schema.get("properties")
+        if not isinstance(properties, dict):
+            return parsed  # pyright: ignore[reportUnknownVariableType]
+        return {
+            key: (
+                _coerce_tool_arg_with_schema(prop_value, prop_schema)  # pyright: ignore[reportUnknownArgumentType]
+                if isinstance(prop_schema, dict)
+                else prop_value
+            )
+            for key, prop_value in parsed.items()  # pyright: ignore[reportUnknownVariableType]
+            for prop_schema in [properties.get(key)]  # type: ignore
+        }
+
+    if schema_type == "array":
+        parsed = value  # pyright: ignore[reportAny]
+        if isinstance(parsed, str):
+            try:
+                parsed = json.loads(parsed)  # pyright: ignore[reportAny]
+            except Exception:
+                return value  # pyright: ignore[reportAny]
+        if not isinstance(parsed, list):
+            return value  # pyright: ignore[reportAny]
+        item_schema = schema.get("items")
+        if not isinstance(item_schema, dict):
+            return parsed  # pyright: ignore[reportUnknownVariableType]
+        return [_coerce_tool_arg_with_schema(item, item_schema) for item in parsed]  # type: ignore
+
+    if schema_type == "integer":
+        if isinstance(value, bool):
+            return value
+        if isinstance(value, int):
+            return value
+        if isinstance(value, float) and value.is_integer():
+            return int(value)
+        if isinstance(value, str):
+            try:
+                return int(value.strip())
+            except ValueError:
+                return value
+        return value
+
+    if schema_type == "number":
+        if isinstance(value, bool):
+            return value
+        if isinstance(value, (int, float)):
+            return value
+        if isinstance(value, str):
+            try:
+                num = float(value.strip())
+                if math.isfinite(num):
+                    return num
+            except ValueError:
+                return value
+        return value
+
+    if schema_type == "boolean":
+        if isinstance(value, bool):
+            return value
+        if isinstance(value, str):
+            lowered = value.strip().lower()
+            if lowered == "true":
+                return True
+            if lowered == "false":
+                return False
+        return value
+
+    return value  # pyright: ignore[reportAny]
+
+
+def _coerce_tool_calls_to_schema(
+    tool_calls: list[ToolCallItem], tools: list[dict[str, Any]]
+) -> list[ToolCallItem]:
+    schema_by_name: dict[str, dict[str, Any]] = {}
+    for tool in tools:
+        function = tool.get("function")
+        if not isinstance(function, dict):
+            continue
+        name = function.get("name")  # type: ignore
+        parameters = function.get("parameters")  # type: ignore
+        if isinstance(name, str) and isinstance(parameters, dict):
+            schema_by_name[name] = parameters
+
+    if not schema_by_name:
+        return tool_calls
+
+    coerced_calls: list[ToolCallItem] = []
+    for tool_call in tool_calls:
+        schema = schema_by_name.get(tool_call.name)
+        if schema is None:
+            coerced_calls.append(tool_call)
+            continue
+
+        try:
+            parsed_args = json.loads(tool_call.arguments)  # pyright: ignore[reportAny]
+        except Exception:
+            coerced_calls.append(tool_call)
+            continue
+
+        if not isinstance(parsed_args, dict):
+            coerced_calls.append(tool_call)
+            continue
+
+        coerced_args = _coerce_tool_arg_with_schema(parsed_args, schema)  # pyright: ignore[reportAny]
+        if not isinstance(coerced_args, dict):
+            coerced_calls.append(tool_call)
+            continue
+
+        coerced_calls.append(
+            tool_call.model_copy(update={"arguments": json.dumps(coerced_args)})
+        )
+    return coerced_calls
 
 
 def make_mlx_parser(
@@ -33,7 +204,7 @@ def make_mlx_parser(
     return ToolParser(
         start_parsing=tool_call_start,
         end_parsing=tool_call_end,
-        parse_tool_calls=parse_tool_calls,
+        _inner_parser=parse_tool_calls,
     )
 
 
@@ -62,7 +233,7 @@ def make_json_parser() -> ToolParser:
     return ToolParser(
         start_parsing="<tool_call>",
         end_parsing="</tool_call>",
-        parse_tool_calls=_parse_json_calls,
+        _inner_parser=_parse_json_calls,
     )
 
 

src/exo/worker/runner/runner_supervisor.py

@@ -12,13 +12,22 @@ from anyio import (
 )
 from loguru import logger
 
+from exo.shared.types.chunks import ErrorChunk
 from exo.shared.types.events import (
+    ChunkGenerated,
     Event,
     RunnerStatusUpdated,
     TaskAcknowledged,
     TaskStatusUpdated,
 )
-from exo.shared.types.tasks import Task, TaskId, TaskStatus
+from exo.shared.types.tasks import (
+    ImageEdits,
+    ImageGeneration,
+    Task,
+    TaskId,
+    TaskStatus,
+    TextGeneration,
+)
 from exo.shared.types.worker.instances import BoundInstance
 from exo.shared.types.worker.runners import (
     RunnerConnecting,
@@ -52,6 +61,7 @@ class RunnerSupervisor:
     _tg: TaskGroup = field(default_factory=TaskGroup, init=False)
     status: RunnerStatus = field(default_factory=RunnerIdle, init=False)
     pending: dict[TaskId, anyio.Event] = field(default_factory=dict, init=False)
+    in_progress: dict[TaskId, Task] = field(default_factory=dict, init=False)
     completed: set[TaskId] = field(default_factory=set, init=False)
     cancelled: set[TaskId] = field(default_factory=set, init=False)
     _cancel_watch_runner: anyio.CancelScope = field(
@@ -147,9 +157,11 @@ class RunnerSupervisor:
         logger.info(f"Starting task {task}")
         event = anyio.Event()
         self.pending[task.task_id] = event
+        self.in_progress[task.task_id] = task
         try:
             await self._task_sender.send_async(task)
         except ClosedResourceError:
+            self.in_progress.pop(task.task_id, None)
             logger.warning(f"Task {task} dropped, runner closed communication.")
             return
         await event.wait()
@@ -157,10 +169,17 @@ class RunnerSupervisor:
     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")
+            self.cancelled.add(task_id)
             return
         self.cancelled.add(task_id)
         with anyio.move_on_after(0.5) as scope:
-            await self._cancel_sender.send_async(task_id)
+            try:
+                await self._cancel_sender.send_async(task_id)
+            except ClosedResourceError:
+                # typically occurs when trying to shut down a failed instance
+                logger.warning(
+                    f"Cancelling task {task_id} failed, runner closed communication"
+                )
         if scope.cancel_called:
             logger.error("RunnerSupervisor cancel pipe blocked")
             await self._check_runner(TimeoutError("cancel pipe blocked"))
@@ -189,6 +208,7 @@ class RunnerSupervisor:
                                 RunnerShuttingDown,
                             ),
                         )
+                        self.in_progress.pop(event.task_id, None)
                         self.completed.add(event.task_id)
                     await self._event_sender.send(event)
         except (ClosedResourceError, BrokenResourceError) as e:
@@ -233,6 +253,22 @@ class RunnerSupervisor:
 
         logger.opt(exception=e).error(f"Runner terminated with {cause}")
 
+        for task in self.in_progress.values():
+            if isinstance(task, (TextGeneration, ImageGeneration, ImageEdits)):
+                with anyio.CancelScope(shield=True):
+                    await self._event_sender.send(
+                        ChunkGenerated(
+                            command_id=task.command_id,
+                            chunk=ErrorChunk(
+                                model=self.shard_metadata.model_card.model_id,
+                                error_message=(
+                                    "Runner shutdown before completing command "
+                                    f"({cause})"
+                                ),
+                            ),
+                        )
+                    )
+
         try:
             self.status = RunnerFailed(error_message=f"Terminated ({cause})")
             with anyio.CancelScope(shield=True):

src/exo/worker/tests/unittests/conftest.py

@@ -20,6 +20,8 @@ class FakeRunnerSupervisor:
     bound_instance: BoundInstance
     status: RunnerStatus
     completed: set[TaskId] = field(default_factory=set)
+    in_progress: set[TaskId] = field(default_factory=set)
+    pending: dict[TaskId, object] = field(default_factory=dict)
 
 
 class OtherTask(BaseTask):

src/exo/worker/tests/unittests/test_mlx/conftest.py

@@ -14,9 +14,9 @@ from exo.shared.constants import EXO_MODELS_DIR
 from exo.shared.models.model_cards import ModelCard, ModelTask
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
+from exo.shared.types.mlx import Model
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 from exo.shared.types.worker.shards import PipelineShardMetadata, TensorShardMetadata
-from exo.worker.engines.mlx import Model
 from exo.worker.engines.mlx.generator.generate import mlx_generate
 from exo.worker.engines.mlx.utils_mlx import apply_chat_template, shard_and_load
 

src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py

@@ -9,8 +9,8 @@ from mlx_lm.models.cache import KVCache
 from mlx_lm.sample_utils import make_sampler
 
 from exo.shared.types.common import ModelId
+from exo.shared.types.mlx import Model
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
-from exo.worker.engines.mlx import Model
 from exo.worker.engines.mlx.cache import (
     KVPrefixCache,
     cache_length,
@@ -143,7 +143,14 @@ class TestKVPrefixCacheWithModel:
         cache = make_kv_cache(model)
 
         _, _, snapshots = prefill(
-            model, tokenizer, make_sampler(0.0), tokens, cache, group=None
+            model,
+            tokenizer,
+            make_sampler(0.0),
+            tokens,
+            cache,
+            group=None,
+            on_prefill_progress=None,
+            distributed_prompt_progress_callback=None,
         )
 
         # Cache should now hold the prompt tokens minus one
@@ -164,7 +171,14 @@ class TestKVPrefixCacheWithModel:
         cache = make_kv_cache(model)
 
         _, _, snapshots = prefill(
-            model, tokenizer, make_sampler(0.0), tokens, cache, group=None
+            model,
+            tokenizer,
+            make_sampler(0.0),
+            tokens,
+            cache,
+            group=None,
+            on_prefill_progress=None,
+            distributed_prompt_progress_callback=None,
         )
 
         kv_prefix_cache = KVPrefixCache(None)
@@ -200,7 +214,14 @@ class TestKVPrefixCacheWithModel:
         cache = make_kv_cache(model)
 
         _, _, snapshots = prefill(
-            model, tokenizer, make_sampler(0.0), short_tokens, cache, group=None
+            model,
+            tokenizer,
+            make_sampler(0.0),
+            short_tokens,
+            cache,
+            group=None,
+            on_prefill_progress=None,
+            distributed_prompt_progress_callback=None,
         )
 
         kv_prefix_cache = KVPrefixCache(None)
@@ -245,7 +266,14 @@ class TestKVPrefixCacheWithModel:
         cache = make_kv_cache(model)
 
         _, _, snapshots = prefill(
-            model, tokenizer, make_sampler(0.0), tokens, cache, group=None
+            model,
+            tokenizer,
+            make_sampler(0.0),
+            tokens,
+            cache,
+            group=None,
+            on_prefill_progress=None,
+            distributed_prompt_progress_callback=None,
         )
 
         kv_prefix_cache = KVPrefixCache(None)
@@ -285,7 +313,14 @@ class TestKVPrefixCacheWithModel:
         cache = make_kv_cache(model)
 
         _, _, snapshots = prefill(
-            model, tokenizer, make_sampler(0.0), tokens, cache, group=None
+            model,
+            tokenizer,
+            make_sampler(0.0),
+            tokens,
+            cache,
+            group=None,
+            on_prefill_progress=None,
+            distributed_prompt_progress_callback=None,
         )
 
         kv_prefix_cache = KVPrefixCache(None)
@@ -513,7 +548,16 @@ class TestKVPrefixCacheWithModel:
             prompt = apply_chat_template(tokenizer, task)
             tokens = encode_prompt(tokenizer, prompt)
             cache = make_kv_cache(model)
-            prefill(model, tokenizer, make_sampler(0.0), tokens, cache, group=None)
+            prefill(
+                model,
+                tokenizer,
+                make_sampler(0.0),
+                tokens,
+                cache,
+                group=None,
+                on_prefill_progress=None,
+                distributed_prompt_progress_callback=None,
+            )
             kv_prefix_cache.add_kv_cache(tokens, cache)
             # Stagger _last_used so LRU order is deterministic
             kv_prefix_cache._last_used[i] = float(i)
@@ -538,7 +582,16 @@ class TestKVPrefixCacheWithModel:
             prompt = apply_chat_template(tokenizer, task)
             tokens = encode_prompt(tokenizer, prompt)
             cache = make_kv_cache(model)
-            prefill(model, tokenizer, make_sampler(0.0), tokens, cache, group=None)
+            prefill(
+                model,
+                tokenizer,
+                make_sampler(0.0),
+                tokens,
+                cache,
+                group=None,
+                on_prefill_progress=None,
+                distributed_prompt_progress_callback=None,
+            )
             kv_prefix_cache.add_kv_cache(tokens, cache)
 
         # LRU entries should have been evicted (entries 0, 1, 2 in order of _last_used)

src/exo/worker/tests/unittests/test_mlx/test_pipeline_prefill_callbacks.py

@@ -0,0 +1,512 @@
+# type: ignore
+"""Test that pipeline prefill callbacks and output exactly match stream_generate.
+
+Spins up a single-device (non-pipeline) run and a distributed pipeline run,
+then verifies that the prompt_progress_callback sequences are identical
+and that generated text matches.
+"""
+
+import json
+import multiprocessing as mp
+import os
+import tempfile
+import traceback
+from typing import Any, cast
+
+import pytest
+
+from exo.shared.constants import EXO_MODELS_DIR
+from exo.shared.models.model_cards import ModelCard, ModelTask
+from exo.shared.types.common import ModelId
+from exo.shared.types.memory import Memory
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+
+MODEL_ID = "mlx-community/gpt-oss-20b-MXFP4-Q8"
+MODEL_PATH = EXO_MODELS_DIR / "mlx-community--gpt-oss-20b-MXFP4-Q8"
+TOTAL_LAYERS = 24
+MAX_TOKENS = 10
+SEED = 42
+TEMPERATURE = 0.0
+
+
+def _model_card() -> ModelCard:
+    return ModelCard(
+        model_id=ModelId(MODEL_ID),
+        storage_size=Memory.from_gb(12),
+        n_layers=TOTAL_LAYERS,
+        hidden_size=2880,
+        supports_tensor=False,
+        tasks=[ModelTask.TextGeneration],
+    )
+
+
+def _build_prompt(tokenizer: Any, prompt_tokens: int) -> tuple[str, Any]:
+    """Build a prompt with the given number of user-content tokens, return (chat_prompt, task)."""
+    from exo.worker.engines.mlx.utils_mlx import apply_chat_template
+
+    base_text = "The quick brown fox jumps over the lazy dog. "
+    base_toks = tokenizer.encode(base_text)
+    repeats = (prompt_tokens // len(base_toks)) + 2
+    long_text = base_text * repeats
+    tokens = tokenizer.encode(long_text)[:prompt_tokens]
+    prompt_text = tokenizer.decode(tokens)
+
+    task = TextGenerationTaskParams(
+        model=MODEL_ID,
+        input=[InputMessage(role="user", content=prompt_text)],
+        max_output_tokens=MAX_TOKENS,
+        temperature=TEMPERATURE,
+        seed=SEED,
+    )
+
+    prompt = apply_chat_template(tokenizer, task)
+    return prompt, task
+
+
+# ---------------------------------------------------------------------------
+# Single-device process: uses stream_generate path (no pipeline layers)
+# ---------------------------------------------------------------------------
+def _run_single_device(
+    prompt_tokens: int,
+    result_queue: Any,
+) -> None:
+    """Load full model without pipeline sharding, run mlx_generate, record callbacks."""
+    try:
+        import mlx.core as mx
+        from mlx_lm.utils import load_model
+
+        from exo.shared.types.worker.shards import PipelineShardMetadata
+        from exo.worker.engines.mlx.cache import encode_prompt
+        from exo.worker.engines.mlx.generator.generate import mlx_generate
+        from exo.worker.engines.mlx.utils_mlx import (
+            build_model_path,
+            get_tokenizer,
+        )
+
+        model_path = build_model_path(ModelId(MODEL_ID))
+        model, _ = load_model(model_path, lazy=True, strict=False)
+        mx.eval(model)
+
+        # Use PipelineShardMetadata just for get_tokenizer (needs model_card), but
+        # do NOT apply pipeline sharding — the model keeps all layers unwrapped.
+        dummy_meta = PipelineShardMetadata(
+            model_card=_model_card(),
+            device_rank=0,
+            world_size=1,
+            start_layer=0,
+            end_layer=TOTAL_LAYERS,
+            n_layers=TOTAL_LAYERS,
+        )
+        tokenizer = get_tokenizer(model_path, dummy_meta)
+
+        prompt, task = _build_prompt(tokenizer, prompt_tokens)
+
+        callbacks: list[tuple[int, int]] = []
+
+        def on_progress(processed: int, total: int) -> None:
+            callbacks.append((processed, total))
+
+        generated_text = ""
+        for response in mlx_generate(
+            model=model,
+            tokenizer=tokenizer,
+            task=task,
+            prompt=prompt,
+            kv_prefix_cache=None,
+            group=None,
+            on_prefill_progress=on_progress,
+        ):
+            generated_text += response.text
+            if response.finish_reason is not None:
+                break
+
+        # Also record the token count that prefill() received (prompt_tokens[:-1])
+        all_tokens = encode_prompt(tokenizer, prompt)
+        prefill_token_count = len(all_tokens) - 1
+
+        result_queue.put(
+            (
+                True,
+                {
+                    "callbacks": callbacks,
+                    "text": generated_text,
+                    "prefill_token_count": prefill_token_count,
+                },
+            )
+        )
+
+    except Exception as e:
+        result_queue.put((False, f"{e}\n{traceback.format_exc()}"))
+
+
+# ---------------------------------------------------------------------------
+# Pipeline device process: uses _pipeline_prefill_cache path
+# ---------------------------------------------------------------------------
+def _run_pipeline_device(
+    rank: int,
+    world_size: int,
+    hostfile_path: str,
+    layer_splits: list[tuple[int, int]],
+    prompt_tokens: int,
+    result_queue: Any,
+) -> None:
+    """Load model with pipeline sharding, run mlx_generate, record callbacks."""
+    os.environ["MLX_HOSTFILE"] = hostfile_path
+    os.environ["MLX_RANK"] = str(rank)
+
+    try:
+        import mlx.core as mx
+
+        from exo.shared.types.worker.shards import PipelineShardMetadata
+        from exo.worker.engines.mlx.cache import encode_prompt
+        from exo.worker.engines.mlx.generator.generate import mlx_generate
+        from exo.worker.engines.mlx.utils_mlx import shard_and_load
+
+        group = mx.distributed.init(backend="ring", strict=True)
+
+        start_layer, end_layer = layer_splits[rank]
+        shard_meta = PipelineShardMetadata(
+            model_card=_model_card(),
+            device_rank=rank,
+            world_size=world_size,
+            start_layer=start_layer,
+            end_layer=end_layer,
+            n_layers=TOTAL_LAYERS,
+        )
+
+        model, tokenizer = shard_and_load(
+            shard_meta, group, on_timeout=None, on_layer_loaded=None
+        )
+        model = cast(Any, model)
+
+        prompt, task = _build_prompt(tokenizer, prompt_tokens)
+
+        callbacks: list[tuple[int, int]] = []
+
+        def on_progress(processed: int, total: int) -> None:
+            callbacks.append((processed, total))
+
+        def distributed_prompt_progress_callback(_group: Any = group) -> None:
+            from exo.worker.engines.mlx.utils_mlx import mx_any
+
+            mx_any(False, _group)
+
+        generated_text = ""
+        for response in mlx_generate(
+            model=model,
+            tokenizer=tokenizer,
+            task=task,
+            prompt=prompt,
+            kv_prefix_cache=None,
+            group=group,
+            on_prefill_progress=on_progress,
+            distributed_prompt_progress_callback=distributed_prompt_progress_callback,
+        ):
+            generated_text += response.text
+            if response.finish_reason is not None:
+                break
+
+        all_tokens = encode_prompt(tokenizer, prompt)
+        prefill_token_count = len(all_tokens) - 1
+
+        result_queue.put(
+            (
+                rank,
+                True,
+                {
+                    "callbacks": callbacks,
+                    "text": generated_text,
+                    "prefill_token_count": prefill_token_count,
+                },
+            )
+        )
+
+    except Exception as e:
+        result_queue.put((rank, False, f"{e}\n{traceback.format_exc()}"))
+
+
+# ---------------------------------------------------------------------------
+# Test helpers
+# ---------------------------------------------------------------------------
+def _create_hostfile(world_size: int, base_port: int) -> str:
+    hosts = [f"127.0.0.1:{base_port + i}" for i in range(world_size)]
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+        json.dump(hosts, f)
+        return f.name
+
+
+def _run_single_device_test(prompt_tokens: int, timeout: int = 120) -> dict[str, Any]:
+    """Run single-device (stream_generate) prefill and return results."""
+    ctx = mp.get_context("spawn")
+    result_queue: Any = ctx.Queue()
+
+    p = ctx.Process(target=_run_single_device, args=(prompt_tokens, result_queue))
+    p.start()
+    p.join(timeout=timeout)
+
+    if p.is_alive():
+        p.terminate()
+        p.join(timeout=5)
+        pytest.fail("Single-device process timed out")
+
+    assert not result_queue.empty(), "Single-device process produced no result"
+    success, data = result_queue.get()
+    assert success, f"Single-device process failed:\n{data}"
+    return data
+
+
+def _run_pipeline_test(
+    layer_splits: list[tuple[int, int]],
+    prompt_tokens: int,
+    base_port: int,
+    timeout: int = 120,
+) -> dict[int, dict[str, Any]]:
+    """Run pipeline prefill across ranks and return per-rank results."""
+    world_size = len(layer_splits)
+    hostfile_path = _create_hostfile(world_size, base_port)
+    ctx = mp.get_context("spawn")
+    result_queue: Any = ctx.Queue()
+
+    try:
+        processes: list[Any] = []
+        for rank in range(world_size):
+            p = ctx.Process(
+                target=_run_pipeline_device,
+                args=(
+                    rank,
+                    world_size,
+                    hostfile_path,
+                    layer_splits,
+                    prompt_tokens,
+                    result_queue,
+                ),
+            )
+            p.start()
+            processes.append(p)
+
+        for p in processes:
+            p.join(timeout=timeout)
+
+        timed_out = any(p.is_alive() for p in processes)
+        for p in processes:
+            if p.is_alive():
+                p.terminate()
+                p.join(timeout=5)
+
+        assert not timed_out, "Pipeline processes timed out"
+
+        results: dict[int, dict[str, Any]] = {}
+        while not result_queue.empty():
+            rank, success, data = result_queue.get()
+            assert success, f"Pipeline rank {rank} failed:\n{data}"
+            results[rank] = data
+
+        assert len(results) == world_size, (
+            f"Expected {world_size} results, got {len(results)}: missing ranks {set(range(world_size)) - results.keys()}"
+        )
+        return results
+
+    finally:
+        os.unlink(hostfile_path)
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+pytestmark = [
+    pytest.mark.slow,
+    pytest.mark.skipif(
+        not MODEL_PATH.exists(),
+        reason=f"GPT-OSS model not found at {MODEL_PATH}",
+    ),
+]
+
+LAYER_SPLITS_4WAY: list[tuple[int, int]] = [(0, 6), (6, 12), (12, 18), (18, 24)]
+LAYER_SPLITS_2WAY: list[tuple[int, int]] = [(0, 12), (12, 24)]
+
+
+class TestPipelineNoDeadlock:
+    """Pipeline prefill must not deadlock at any rank count or prompt length."""
+
+    @pytest.mark.parametrize(
+        "layer_splits,prompt_tokens",
+        [
+            (LAYER_SPLITS_2WAY, 128),
+            (LAYER_SPLITS_2WAY, 4096),
+            (LAYER_SPLITS_2WAY, 8192),
+            (LAYER_SPLITS_2WAY, 16384),
+            (LAYER_SPLITS_4WAY, 128),
+            (LAYER_SPLITS_4WAY, 4096),
+            (LAYER_SPLITS_4WAY, 8192),
+            (LAYER_SPLITS_4WAY, 16384),
+        ],
+        ids=[
+            "2rank_128tok",
+            "2rank_4096tok",
+            "2rank_8192tok",
+            "2rank_16384tok",
+            "4rank_128tok",
+            "4rank_4096tok",
+            "4rank_8192tok",
+            "4rank_16384tok",
+        ],
+    )
+    def test_no_deadlock(
+        self,
+        layer_splits: list[tuple[int, int]],
+        prompt_tokens: int,
+    ) -> None:
+        """Pipeline must complete without deadlock at various prompt lengths."""
+        pipeline_results = _run_pipeline_test(
+            layer_splits=layer_splits,
+            prompt_tokens=prompt_tokens,
+            base_port=29650,
+            timeout=60,
+        )
+        # If we get here, no deadlock. Verify all ranks produced output.
+        for rank, pipe_data in sorted(pipeline_results.items()):
+            assert pipe_data["text"], f"Rank {rank} produced no output text"
+
+
+class TestPipelinePrefillCallbacks:
+    """Verify that pipeline prefill callbacks exactly match stream_generate callbacks."""
+
+    @pytest.mark.parametrize(
+        "prompt_tokens",
+        [50, 500, 5000],
+        ids=["short_50", "medium_500", "long_5000"],
+    )
+    def test_callbacks_match(self, prompt_tokens: int) -> None:
+        """All pipeline ranks must produce identical callback sequences."""
+        # Run 4-rank pipeline
+        pipeline_results = _run_pipeline_test(
+            layer_splits=LAYER_SPLITS_4WAY,
+            prompt_tokens=prompt_tokens,
+            base_port=29700,
+            timeout=180,
+        )
+
+        # All ranks must agree on prefill token count and callback sequence
+        rank0_data = pipeline_results[0]
+        rank0_callbacks = rank0_data["callbacks"]
+        prefill_count = rank0_data["prefill_token_count"]
+
+        for rank, pipe_data in sorted(pipeline_results.items()):
+            pipe_callbacks = pipe_data["callbacks"]
+
+            assert pipe_data["prefill_token_count"] == prefill_count, (
+                f"Rank {rank} prefill token count mismatch: "
+                f"{pipe_data['prefill_token_count']} vs {prefill_count}"
+            )
+
+            assert pipe_callbacks == rank0_callbacks, (
+                f"Rank {rank} callback mismatch for {prompt_tokens} prompt tokens "
+                f"(prefill M={prefill_count}):\n"
+                f"  pipeline R0 ({len(rank0_callbacks)} callbacks): {rank0_callbacks}\n"
+                f"  pipeline R{rank} ({len(pipe_callbacks)} callbacks): {pipe_callbacks}"
+            )
+
+        # Structural checks: starts with (0, M), ends with (M, M), monotonically increasing
+        assert rank0_callbacks[0] == (0, prefill_count), (
+            f"First callback should be (0, {prefill_count}), got {rank0_callbacks[0]}"
+        )
+        assert rank0_callbacks[-1] == (prefill_count, prefill_count), (
+            f"Last callback should be ({prefill_count}, {prefill_count}), got {rank0_callbacks[-1]}"
+        )
+        for i in range(1, len(rank0_callbacks)):
+            assert rank0_callbacks[i][0] >= rank0_callbacks[i - 1][0], (
+                f"Callbacks not monotonically increasing at index {i}: {rank0_callbacks}"
+            )
+
+    @pytest.mark.parametrize(
+        "prompt_tokens",
+        [50, 500],
+        ids=["short_50", "medium_500"],
+    )
+    def test_output_matches(self, prompt_tokens: int) -> None:
+        """Pipeline-generated text must match single-device output."""
+        single = _run_single_device_test(prompt_tokens, timeout=180)
+
+        pipeline_results = _run_pipeline_test(
+            layer_splits=LAYER_SPLITS_4WAY,
+            prompt_tokens=prompt_tokens,
+            base_port=29800,
+            timeout=180,
+        )
+
+        single_text = single["text"]
+
+        # The last rank produces the final logits, so its output should match.
+        # Due to SDPA tiling non-determinism, allow minor differences in text.
+        last_rank = max(pipeline_results.keys())
+        pipe_text = pipeline_results[last_rank]["text"]
+
+        # For deterministic sampling (temp=0.0), outputs should match exactly
+        # or be very close. Log both for debugging even if they match.
+        if single_text != pipe_text:
+            # Find first divergence point
+            min_len = min(len(single_text), len(pipe_text))
+            diverge_idx = next(
+                (i for i in range(min_len) if single_text[i] != pipe_text[i]),
+                min_len,
+            )
+            pytest.fail(
+                f"Output text diverged at character {diverge_idx} for {prompt_tokens} prompt tokens:\n"
+                f"  single-device: {single_text!r}\n"
+                f"  pipeline R{last_rank}: {pipe_text!r}"
+            )
+
+
+class TestPipelineCallbacksStructure:
+    """Verify structural properties of callbacks independent of model output."""
+
+    def test_callback_structure_matches_generate_step(self) -> None:
+        """Verify callbacks follow generate_step's pattern: (0,M), chunks up to M-1, (M,M)."""
+        prompt_tokens = 200
+        pipeline_results = _run_pipeline_test(
+            layer_splits=LAYER_SPLITS_4WAY,
+            prompt_tokens=prompt_tokens,
+            base_port=29900,
+            timeout=180,
+        )
+
+        for rank, pipe_data in sorted(pipeline_results.items()):
+            callbacks = pipe_data["callbacks"]
+            m = pipe_data["prefill_token_count"]
+            assert m > 0, f"Rank {rank}: prefill token count is 0"
+
+            assert callbacks[0] == (0, m), (
+                f"Rank {rank}: first callback should be (0, {m}), got {callbacks[0]}"
+            )
+
+            assert callbacks[-1] == (m, m), (
+                f"Rank {rank}: last callback should be ({m}, {m}), got {callbacks[-1]}"
+            )
+
+            if len(callbacks) > 2:
+                second_to_last = callbacks[-2]
+                assert second_to_last[0] < m, (
+                    f"Rank {rank}: second-to-last callback should report < {m}, "
+                    f"got {second_to_last}"
+                )
+
+            # All callbacks must have total == M
+            for i, (_, total) in enumerate(callbacks):
+                assert total == m, (
+                    f"Rank {rank}: callback {i} has total={total}, expected {m}"
+                )
+
+            # processed values must be non-decreasing
+            processed_vals = [p for p, _ in callbacks]
+            for i in range(1, len(processed_vals)):
+                assert processed_vals[i] >= processed_vals[i - 1], (
+                    f"Rank {rank}: callbacks not non-decreasing at index {i}: "
+                    f"{processed_vals}"
+                )
+
+            # No duplicate consecutive callbacks (pipeline dummies must not emit callbacks)
+            for i in range(1, len(callbacks)):
+                assert callbacks[i] != callbacks[i - 1], (
+                    f"Rank {rank}: duplicate consecutive callback at index {i}: "
+                    f"{callbacks[i]} (this suggests dummy iterations are emitting callbacks)"
+                )

src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.py

@@ -15,8 +15,8 @@ from mlx.utils import tree_flatten, tree_unflatten
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
 from exo.shared.types.common import ModelId
+from exo.shared.types.mlx import Model
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
-from exo.worker.engines.mlx import Model
 from exo.worker.engines.mlx.cache import KVPrefixCache
 from exo.worker.engines.mlx.generator.generate import mlx_generate
 from exo.worker.engines.mlx.utils_mlx import (

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

@@ -19,7 +19,7 @@ from exo.worker.engines.mlx.dsml_encoding import (
     encode_messages,
     parse_dsml_output,
 )
-from exo.worker.runner.llm_inference.runner import parse_deepseek_v32
+from exo.worker.runner.llm_inference.model_output_parsers import parse_deepseek_v32
 
 # ── Shared fixtures ──────────────────────────────────────────────
 

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

@@ -6,6 +6,8 @@ from typing import Callable
 import mlx.core as mx
 import pytest
 
+import exo.worker.runner.llm_inference.batch_generator as mlx_batch_generator
+import exo.worker.runner.llm_inference.model_output_parsers as mlx_model_output_parsers
 import exo.worker.runner.llm_inference.runner as mlx_runner
 from exo.shared.types.chunks import TokenChunk
 from exo.shared.types.events import (
@@ -114,27 +116,41 @@ def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
     # initialize_mlx returns a mock group
     monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(MockGroup()))
     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))
+    monkeypatch.setattr(mlx_batch_generator, "warmup_inference", make_nothin(1))
+    monkeypatch.setattr(mlx_batch_generator, "_check_for_debug_prompts", nothin)
+    monkeypatch.setattr(mlx_batch_generator, "mx_any", make_nothin(False))
+
+    def fake_all_gather(
+        tasks: list[TextGeneration], group: object
+    ) -> tuple[list[TextGeneration], list[TextGeneration]]:
+        return (tasks, [])
+
+    monkeypatch.setattr(mlx_batch_generator, "mx_all_gather_tasks", fake_all_gather)
     # 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"))
-    monkeypatch.setattr(mlx_runner, "detect_thinking_prompt_suffix", make_nothin(False))
+    monkeypatch.setattr(
+        mlx_batch_generator, "apply_chat_template", make_nothin("test prompt")
+    )
+    monkeypatch.setattr(
+        mlx_model_output_parsers, "detect_thinking_prompt_suffix", make_nothin(False)
+    )
 
     def fake_generate(*_1: object, **_2: object):
         yield GenerationResponse(token=0, text="hi", finish_reason="stop", usage=None)
 
-    monkeypatch.setattr(mlx_runner, "mlx_generate", fake_generate)
+    monkeypatch.setattr(mlx_batch_generator, "mlx_generate", fake_generate)
 
 
 # Use a fake event_sender to remove test flakiness.
 class EventCollector:
-    def __init__(self) -> None:
+    def __init__(self, on_event: Callable[[Event], None] | None = None) -> None:
         self.events: list[Event] = []
+        self._on_event = on_event
 
     def send(self, event: Event) -> None:
         self.events.append(event)
+        if self._on_event:
+            self._on_event(event)
 
     def close(self) -> None:
         pass
@@ -159,7 +175,7 @@ class MockGroup:
         return 1
 
 
-def _run(tasks: Iterable[Task]):
+def _run(tasks: Iterable[Task], send_after_ready: list[Task] | None = None):
     bound_instance = get_bound_mlx_ring_instance(
         instance_id=INSTANCE_1_ID,
         model_id=MODEL_A_ID,
@@ -169,7 +185,23 @@ def _run(tasks: Iterable[Task]):
 
     task_sender, task_receiver = mp_channel[Task]()
     _cancel_sender, cancel_receiver = mp_channel[TaskId]()
-    event_sender = EventCollector()
+
+    on_event: Callable[[Event], None] | None = None
+    if send_after_ready:
+        _saw_running = False
+
+        def _on_event(event: Event) -> None:
+            nonlocal _saw_running
+            if isinstance(event, RunnerStatusUpdated):
+                if isinstance(event.runner_status, RunnerRunning):
+                    _saw_running = True
+                elif _saw_running and isinstance(event.runner_status, RunnerReady):
+                    for t in send_after_ready:
+                        task_sender.send(t)
+
+        on_event = _on_event
+
+    event_sender = EventCollector(on_event=on_event)
 
     with task_sender:
         for t in tasks:
@@ -183,18 +215,22 @@ def _run(tasks: Iterable[Task]):
             "exo.worker.runner.llm_inference.runner.mx.distributed.all_gather",
             make_nothin(mx.array([1])),
         ):
-            mlx_runner.main(
+            runner = mlx_runner.Runner(
                 bound_instance,
                 event_sender,  # pyright: ignore[reportArgumentType]
                 task_receiver,
                 cancel_receiver,
             )
+            runner.main()
 
         return event_sender.events
 
 
 def test_events_processed_in_correct_order(patch_out_mlx: pytest.MonkeyPatch):
-    events = _run([INIT_TASK, LOAD_TASK, WARMUP_TASK, CHAT_TASK, SHUTDOWN_TASK])
+    events = _run(
+        [INIT_TASK, LOAD_TASK, WARMUP_TASK, CHAT_TASK],
+        send_after_ready=[SHUTDOWN_TASK],
+    )
 
     expected_chunk = ChunkGenerated(
         command_id=COMMAND_1_ID,

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

@@ -4,7 +4,7 @@ from exo.shared.types.worker.runner_response import (
     GenerationResponse,
     ToolCallResponse,
 )
-from exo.worker.runner.llm_inference.runner import parse_gpt_oss
+from exo.worker.runner.llm_inference.model_output_parsers import parse_gpt_oss
 
 # Token IDs from mlx-community/gpt-oss-20b-MXFP4-Q8 tokenizer.
 # These are stable since they come from the model's vocabulary.
@@ -107,7 +107,7 @@ def _collect(
     def _gen() -> Generator[GenerationResponse, None, None]:
         yield from _make_gen_responses(tokens)
 
-    return list(parse_gpt_oss(_gen()))
+    return list(x for x in parse_gpt_oss(_gen()) if x is not None)
 
 
 def _get_tool_call(

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

@@ -1,10 +1,11 @@
 """Tests for parse_tool_calls generator, especially unclosed tool call handling."""
 
+import json
 from collections.abc import Generator
 from typing import Any
 
 from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse
-from exo.worker.runner.llm_inference.runner import parse_tool_calls
+from exo.worker.runner.llm_inference.model_output_parsers import parse_tool_calls
 from exo.worker.runner.llm_inference.tool_parsers import make_mlx_parser
 
 
@@ -40,6 +41,7 @@ class TestParseToolCalls:
             parse_tool_calls(
                 _make_responses(texts, finish_on_last=False),
                 _dummy_parser,
+                tools=None,
             )
         )
 
@@ -53,6 +55,7 @@ class TestParseToolCalls:
             parse_tool_calls(
                 _make_responses(texts),
                 _dummy_parser,
+                tools=None,
             )
         )
 
@@ -77,9 +80,101 @@ class TestParseToolCalls:
             parse_tool_calls(
                 _make_responses(texts, finish_on_last=False),
                 make_mlx_parser("<tool_call>", "</tool_call>", _failing_parser),
+                tools=None,
             )
         )
 
         assert len(results) == 1
         assert isinstance(results[0], GenerationResponse)
         assert results[0].text == "<tool_call>bad content</tool_call>"
+
+    def test_tool_schema_coerces_string_arguments_to_expected_types(self):
+        """Tool argument values should be coerced using provided JSON schema."""
+
+        def _parser_with_string_args(_text: str) -> dict[str, Any]:
+            return {
+                "name": "process",
+                "arguments": {
+                    "action": "output",
+                    "id": "0",
+                    "verbose": "true",
+                    "temperature": "0.75",
+                },
+            }
+
+        tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": "process",
+                    "description": "Manage background processes",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "action": {"type": "string"},
+                            "id": {"type": "integer"},
+                            "verbose": {"type": "boolean"},
+                            "temperature": {"type": "number"},
+                        },
+                        "required": ["action"],
+                    },
+                },
+            }
+        ]
+
+        results = list(
+            parse_tool_calls(
+                _make_responses(["<tool_call>", "process", "</tool_call>"]),
+                make_mlx_parser(
+                    "<tool_call>", "</tool_call>", _parser_with_string_args
+                ),
+                tools,
+            )
+        )
+
+        assert len(results) == 1
+        assert isinstance(results[0], ToolCallResponse)
+
+        args = json.loads(results[0].tool_calls[0].arguments)  # pyright: ignore[reportAny]
+        assert args == {
+            "action": "output",
+            "id": 0,
+            "verbose": True,
+            "temperature": 0.75,
+        }
+
+    def test_schema_coercion_skips_unknown_tools(self):
+        """If no matching tool schema exists, arguments should remain unchanged."""
+
+        def _parser_with_string_id(_text: str) -> dict[str, Any]:
+            return {
+                "name": "process",
+                "arguments": {"action": "output", "id": "0"},
+            }
+
+        tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": "different_tool",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {"id": {"type": "integer"}},
+                    },
+                },
+            }
+        ]
+
+        results = list(
+            parse_tool_calls(
+                _make_responses(["<tool_call>", "process", "</tool_call>"]),
+                make_mlx_parser("<tool_call>", "</tool_call>", _parser_with_string_id),
+                tools,
+            )
+        )
+
+        assert len(results) == 1
+        assert isinstance(results[0], ToolCallResponse)
+
+        args = json.loads(results[0].tool_calls[0].arguments)  # pyright: ignore[reportAny]
+        assert args == {"action": "output", "id": "0"}

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

@@ -1 +1,93 @@
-# TODO:
+import multiprocessing as mp
+from typing import cast
+
+import anyio
+import pytest
+
+from exo.shared.models.model_cards import ModelId
+from exo.shared.types.chunks import ErrorChunk
+from exo.shared.types.common import CommandId, NodeId
+from exo.shared.types.events import ChunkGenerated, Event, RunnerStatusUpdated
+from exo.shared.types.tasks import Task, TaskId, TextGeneration
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo.shared.types.worker.instances import BoundInstance, InstanceId
+from exo.shared.types.worker.runners import RunnerFailed, RunnerId
+from exo.utils.channels import channel, mp_channel
+from exo.worker.runner.runner_supervisor import RunnerSupervisor
+from exo.worker.tests.unittests.conftest import get_bound_mlx_ring_instance
+
+
+class _DeadProcess:
+    exitcode = -6
+
+    def start(self) -> None:
+        return None
+
+    def is_alive(self) -> bool:
+        return False
+
+    def join(self, _timeout: float | None = None) -> None:
+        return None
+
+    def terminate(self) -> None:
+        return None
+
+    def kill(self) -> None:
+        return None
+
+
+@pytest.mark.asyncio
+async def test_check_runner_emits_error_chunk_for_inflight_text_generation() -> None:
+    event_sender, event_receiver = channel[Event]()
+    task_sender, _ = mp_channel[Task]()
+    cancel_sender, _ = mp_channel[TaskId]()
+    _, ev_recv = mp_channel[Event]()
+
+    bound_instance: BoundInstance = get_bound_mlx_ring_instance(
+        instance_id=InstanceId("instance-a"),
+        model_id=ModelId("mlx-community/Llama-3.2-1B-Instruct-4bit"),
+        runner_id=RunnerId("runner-a"),
+        node_id=NodeId("node-a"),
+    )
+
+    supervisor = RunnerSupervisor(
+        shard_metadata=bound_instance.bound_shard,
+        bound_instance=bound_instance,
+        runner_process=cast("mp.Process", cast(object, _DeadProcess())),
+        initialize_timeout=400,
+        _ev_recv=ev_recv,
+        _task_sender=task_sender,
+        _event_sender=event_sender,
+        _cancel_sender=cancel_sender,
+    )
+
+    command_id = CommandId("cmd-a")
+    task = TextGeneration(
+        task_id=TaskId("task-a"),
+        instance_id=bound_instance.instance.instance_id,
+        command_id=command_id,
+        task_params=TextGenerationTaskParams(
+            model=bound_instance.bound_shard.model_card.model_id,
+            input=[InputMessage(role="user", content="hi")],
+            stream=True,
+        ),
+    )
+    supervisor.in_progress[task.task_id] = task
+    supervisor.shutdown = lambda: None
+
+    await supervisor._check_runner(RuntimeError("boom"))  # pyright: ignore[reportPrivateUsage]
+
+    got_chunk = await event_receiver.receive()
+    got_status = await event_receiver.receive()
+
+    assert isinstance(got_chunk, ChunkGenerated)
+    assert got_chunk.command_id == command_id
+    assert isinstance(got_chunk.chunk, ErrorChunk)
+    assert "Runner shutdown before completing command" in got_chunk.chunk.error_message
+
+    assert isinstance(got_status, RunnerStatusUpdated)
+    assert isinstance(got_status.runner_status, RunnerFailed)
+
+    event_sender.close()
+    with anyio.move_on_after(0.1):
+        await event_receiver.aclose()

tmp/gen_card.py

@@ -0,0 +1,33 @@
+"""
+Generates inference model cards for EXO.
+Usage:
+    uv run tmp/gen_card.py mlx-community/my_cool_model-8bit [repo-id/model-id-2] [...]
+
+Model Cards require cleanup for family & quantization data
+"""
+
+import sys
+
+import anyio
+
+from exo.shared.models.model_cards import ModelCard, ModelId
+
+
+async def main():
+    if len(sys.argv) == 1:
+        print(f"USAGE: {sys.argv[0]} repo-id/model-id-1 [repo-id/model-id-2] [...]")
+        quit(1)
+    print("Remember! Model Cards require cleanup for family & quantization data")
+    for arg in sys.argv[1:]:
+        mid = ModelId(arg)
+        mc = await ModelCard.fetch_from_hf(mid)
+        await mc.save(
+            anyio.Path(__file__).parent.parent
+            / "resources"
+            / "inference_model_cards"
+            / (mid.normalize() + ".toml")
+        )
+
+
+if __name__ == "__main__":
+    anyio.run(main)

uv.lock

@@ -418,7 +418,7 @@ requires-dist = [
     { name = "mflux", specifier = "==0.15.5" },
     { name = "mlx", marker = "sys_platform == 'darwin'", git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks" },
     { name = "mlx", extras = ["cpu"], marker = "sys_platform == 'linux'", specifier = "==0.30.6" },
-    { name = "mlx-lm", specifier = "==0.30.7" },
+    { name = "mlx-lm", git = "https://github.com/ml-explore/mlx-lm?rev=834fac934c4e04de9b3d723e2b9287a2c60cfd4a" },
     { name = "msgspec", specifier = ">=0.19.0" },
     { name = "openai-harmony", specifier = ">=0.0.8" },
     { name = "pillow", specifier = ">=11.0,<12.0" },
@@ -469,7 +469,7 @@ requires-dist = [
 
 [[package]]
 name = "exo-pyo3-bindings"
-version = "0.1.0"
+version = "0.2.0"
 source = { editable = "rust/exo_pyo3_bindings" }
 
 [package.dev-dependencies]
@@ -1104,8 +1104,8 @@ wheels = [
 
 [[package]]
 name = "mlx-lm"
-version = "0.30.7"
-source = { registry = "https://pypi.org/simple" }
+version = "0.30.8"
+source = { git = "https://github.com/ml-explore/mlx-lm?rev=834fac934c4e04de9b3d723e2b9287a2c60cfd4a#834fac934c4e04de9b3d723e2b9287a2c60cfd4a" }
 dependencies = [
     { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mlx", version = "0.30.7.dev20260225+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
@@ -1115,10 +1115,6 @@ dependencies = [
     { name = "sentencepiece", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "transformers", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/66/0d/56542e2ae13ec6f542d3977d7cff89a205d4f6c5122e0ce23f33265f61c9/mlx_lm-0.30.7.tar.gz", hash = "sha256:e5f31ac58d9f2381f28e1ba639ff903e64f7cff1bdc245c0bc97f72264be329c", size = 275764, upload-time = "2026-02-12T18:41:11.86Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/1e/17/a41c798a3d9cbdc47f39c6db5bba4c2cd199203ead26bf911cb03b644070/mlx_lm-0.30.7-py3-none-any.whl", hash = "sha256:17442a4bf01c4c2d3bca1e647712fe44f19890c3f1eadc8589d389e57b44b9bf", size = 386591, upload-time = "2026-02-12T18:41:10.236Z" },
-]
 
 [[package]]
 name = "more-itertools"