feat: auto-discover models from HuggingFace Hub cache and custom paths

Searches ~/.cache/huggingface/hub by default so models downloaded by
other tools (HF CLI, transformers, LM Studio, etc.) are used without
re-downloading. Adds EXO_MODELS_PATH env var support and a Model Search
Paths setting in the macOS app.

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

.mlx_typings/mlx/core/__init__.pyi

@@ -215,22 +215,6 @@ class StreamContext:
         traceback: object | None = ...,
     ) -> None: ...
 
-def device_info() -> dict[str, str | int]:
-    """
-    Get information about the GPU device and system settings.
-
-    Currently returns:
-
-    * ``architecture``
-    * ``max_buffer_size``
-    * ``max_recommended_working_set_size``
-    * ``memory_size``
-    * ``resource_limit``
-
-    Returns:
-        dict: A dictionary with string keys and string or integer values.
-    """
-
 def abs(a: array, /, *, stream: Stream | Device | None = ...) -> array:
     """
     Element-wise absolute value.

.mlx_typings/mlx_lm/generate.pyi

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

.mlx_typings/mlx_lm/models/qwen3_5.pyi

@@ -1,153 +0,0 @@
-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

@@ -1,19 +0,0 @@
-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,15 +7,6 @@ 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,28 +216,6 @@ 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"
@@ -2781,7 +2759,6 @@ dependencies = [
 name = "networking"
 version = "0.0.1"
 dependencies = [
- "async-stream",
  "delegate",
  "either",
  "extend",
@@ -2790,7 +2767,6 @@ dependencies = [
  "keccak-const",
  "libp2p",
  "log",
- "pin-project",
  "tokio",
  "tracing-subscriber",
  "util",

Cargo.toml

@@ -1,6 +1,10 @@
 [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"
@@ -30,7 +34,6 @@ 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"

app/EXO/EXO/ContentView.swift

@@ -21,21 +21,6 @@ struct ContentView: View {
     @State private var showAllNodes = false
     @State private var showAllInstances = false
     @State private var baseURLCopied = false
-    @State private var showAdvanced = false
-    @State private var showDebugInfo = false
-    private enum BugReportPhase: Equatable {
-        case idle
-        case prompting
-        case sending(String)
-        case success(String)
-        case failure(String)
-    }
-    @State private var bugReportPhase: BugReportPhase = .idle
-    @State private var bugReportUserDescription: String = ""
-    @State private var uninstallInProgress = false
-    @State private var pendingNamespace: String = ""
-    @State private var pendingHFToken: String = ""
-    @State private var pendingEnableImageModels = false
 
     var body: some View {
         VStack(alignment: .leading, spacing: 12) {
@@ -394,283 +379,6 @@ struct ContentView: View {
         }
     }
 
-    private var thunderboltStatusText: String {
-        switch networkStatusService.status.thunderboltBridgeState {
-        case .some(.disabled):
-            return "Thunderbolt Bridge: Disabled"
-        case .some(.deleted):
-            return "Thunderbolt Bridge: Deleted"
-        case .some(.enabled):
-            return "Thunderbolt Bridge: Enabled"
-        case nil:
-            return "Thunderbolt Bridge: Unknown"
-        }
-    }
-
-    private var thunderboltStatusColor: Color {
-        switch networkStatusService.status.thunderboltBridgeState {
-        case .some(.disabled), .some(.deleted):
-            return .green
-        case .some(.enabled):
-            return .red
-        case nil:
-            return .secondary
-        }
-    }
-
-    /// Shows TB bridge status for all nodes from exo cluster state
-    private var clusterThunderboltBridgeView: some View {
-        let bridgeStatuses = stateService.latestSnapshot?.nodeThunderboltBridge ?? [:]
-        let localNodeId = stateService.localNodeId
-        let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:]
-
-        return VStack(alignment: .leading, spacing: 1) {
-            if bridgeStatuses.isEmpty {
-                Text("Cluster TB Bridge: No data")
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-            } else {
-                Text("Cluster TB Bridge Status:")
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-                ForEach(Array(bridgeStatuses.keys.sorted()), id: \.self) { nodeId in
-                    if let status = bridgeStatuses[nodeId] {
-                        let nodeName =
-                            nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8))
-                        let isLocal = nodeId == localNodeId
-                        let prefix = isLocal ? "  \(nodeName) (local):" : "  \(nodeName):"
-                        let statusText =
-                            !status.exists
-                            ? "N/A"
-                            : (status.enabled ? "Enabled" : "Disabled")
-                        let color: Color =
-                            !status.exists
-                            ? .secondary
-                            : (status.enabled ? .red : .green)
-                        Text("\(prefix) \(statusText)")
-                            .font(.caption2)
-                            .foregroundColor(color)
-                    }
-                }
-            }
-        }
-    }
-
-    private var interfaceIpList: some View {
-        let statuses = networkStatusService.status.interfaceStatuses
-        return VStack(alignment: .leading, spacing: 1) {
-            Text("Interfaces (en0–en7):")
-                .font(.caption2)
-                .foregroundColor(.secondary)
-            if statuses.isEmpty {
-                Text("  Unknown")
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-            } else {
-                ForEach(statuses, id: \.interfaceName) { status in
-                    let ipText = status.ipAddress ?? "No IP"
-                    Text("  \(status.interfaceName): \(ipText)")
-                        .font(.caption2)
-                        .foregroundColor(status.ipAddress == nil ? .red : .green)
-                }
-            }
-        }
-    }
-
-    private var debugSection: some View {
-        VStack(alignment: .leading, spacing: 4) {
-            HoverButton(
-                title: "Debug Info",
-                tint: .primary,
-                trailingSystemImage: showDebugInfo ? "chevron.up" : "chevron.down",
-                small: true
-            ) {
-                showDebugInfo.toggle()
-            }
-            if showDebugInfo {
-                VStack(alignment: .leading, spacing: 4) {
-                    Text("Version: \(buildTag)")
-                        .font(.caption2)
-                        .foregroundColor(.secondary)
-                    Text("Commit: \(buildCommit)")
-                        .font(.caption2)
-                        .foregroundColor(.secondary)
-                    Text(thunderboltStatusText)
-                        .font(.caption2)
-                        .foregroundColor(thunderboltStatusColor)
-                    clusterThunderboltBridgeView
-                    interfaceIpList
-                    rdmaStatusView
-                    sendBugReportButton
-                        .padding(.top, 6)
-                }
-                .padding(.leading, 8)
-                .transition(.opacity)
-            }
-        }
-        .animation(.easeInOut(duration: 0.25), value: showDebugInfo)
-    }
-
-    private var rdmaStatusView: some View {
-        let rdmaStatuses = stateService.latestSnapshot?.nodeRdmaCtl ?? [:]
-        let localNodeId = stateService.localNodeId
-        let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:]
-        let localDevices = networkStatusService.status.localRdmaDevices
-        let localPorts = networkStatusService.status.localRdmaActivePorts
-
-        return VStack(alignment: .leading, spacing: 1) {
-            if rdmaStatuses.isEmpty {
-                Text("Cluster RDMA: No data")
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-            } else {
-                Text("Cluster RDMA Status:")
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-                ForEach(Array(rdmaStatuses.keys.sorted()), id: \.self) { nodeId in
-                    if let status = rdmaStatuses[nodeId] {
-                        let nodeName =
-                            nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8))
-                        let isLocal = nodeId == localNodeId
-                        let prefix = isLocal ? "  \(nodeName) (local):" : "  \(nodeName):"
-                        let statusText = status.enabled ? "Enabled" : "Disabled"
-                        let color: Color = status.enabled ? .green : .orange
-                        Text("\(prefix) \(statusText)")
-                            .font(.caption2)
-                            .foregroundColor(color)
-                    }
-                }
-            }
-            if !localDevices.isEmpty {
-                Text("  Local Devices: \(localDevices.joined(separator: ", "))")
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-            }
-            if !localPorts.isEmpty {
-                Text("  Local Active Ports:")
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-                ForEach(localPorts, id: \.device) { port in
-                    Text("    \(port.device) port \(port.port): \(port.state)")
-                        .font(.caption2)
-                        .foregroundColor(.green)
-                }
-            }
-        }
-    }
-
-    private var sendBugReportButton: some View {
-        VStack(alignment: .leading, spacing: 6) {
-            switch bugReportPhase {
-            case .idle:
-                Button {
-                    bugReportPhase = .prompting
-                    bugReportUserDescription = ""
-                } label: {
-                    HStack {
-                        Text("Send Bug Report")
-                            .font(.caption)
-                            .fontWeight(.semibold)
-                        Spacer()
-                    }
-                    .padding(.vertical, 6)
-                    .padding(.horizontal, 8)
-                    .background(
-                        RoundedRectangle(cornerRadius: 6)
-                            .fill(Color.accentColor.opacity(0.12))
-                    )
-                }
-                .buttonStyle(.plain)
-
-            case .prompting:
-                VStack(alignment: .leading, spacing: 6) {
-                    Text("What's the issue? (optional)")
-                        .font(.caption2)
-                        .foregroundColor(.secondary)
-                    TextEditor(text: $bugReportUserDescription)
-                        .font(.caption2)
-                        .frame(height: 60)
-                        .overlay(
-                            RoundedRectangle(cornerRadius: 4)
-                                .stroke(Color.secondary.opacity(0.3), lineWidth: 1)
-                        )
-                    HStack(spacing: 8) {
-                        Button("Send") {
-                            Task {
-                                await sendBugReport()
-                            }
-                        }
-                        .font(.caption2)
-                        .buttonStyle(.borderedProminent)
-                        .controlSize(.small)
-                        Button("Cancel") {
-                            bugReportPhase = .idle
-                        }
-                        .font(.caption2)
-                        .buttonStyle(.bordered)
-                        .controlSize(.small)
-                    }
-                }
-                .padding(8)
-                .background(
-                    RoundedRectangle(cornerRadius: 6)
-                        .fill(Color.accentColor.opacity(0.06))
-                )
-
-            case .sending(let message):
-                HStack(spacing: 6) {
-                    ProgressView()
-                        .scaleEffect(0.6)
-                    Text(message)
-                        .font(.caption2)
-                        .foregroundColor(.secondary)
-                }
-
-            case .success(let message):
-                VStack(alignment: .leading, spacing: 6) {
-                    Text(message)
-                        .font(.caption2)
-                        .foregroundColor(.secondary)
-                        .fixedSize(horizontal: false, vertical: true)
-                    Button {
-                        openGitHubIssue()
-                    } label: {
-                        HStack(spacing: 4) {
-                            Image(systemName: "arrow.up.right.square")
-                                .imageScale(.small)
-                            Text("Create GitHub Issue")
-                                .font(.caption2)
-                        }
-                    }
-                    .buttonStyle(.bordered)
-                    .controlSize(.small)
-                    Button("Done") {
-                        bugReportPhase = .idle
-                        bugReportUserDescription = ""
-                    }
-                    .font(.caption2)
-                    .buttonStyle(.plain)
-                    .foregroundColor(.secondary)
-                }
-
-            case .failure(let message):
-                VStack(alignment: .leading, spacing: 4) {
-                    Text(message)
-                        .font(.caption2)
-                        .foregroundColor(.red)
-                        .fixedSize(horizontal: false, vertical: true)
-                    Button("Dismiss") {
-                        bugReportPhase = .idle
-                    }
-                    .font(.caption2)
-                    .buttonStyle(.plain)
-                    .foregroundColor(.secondary)
-                }
-            }
-        }
-        .animation(.easeInOut(duration: 0.2), value: bugReportPhase)
-    }
-
     private var processToggleBinding: Binding<Bool> {
         Binding(
             get: {
@@ -711,143 +419,6 @@ struct ContentView: View {
         )
     }
 
-    private func sendBugReport() async {
-        bugReportPhase = .sending("Collecting logs...")
-        let service = BugReportService()
-        let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines)
-        do {
-            let outcome = try await service.sendReport(
-                isManual: true,
-                userDescription: description.isEmpty ? nil : description
-            )
-            if outcome.success {
-                bugReportPhase = .success(outcome.message)
-            } else {
-                bugReportPhase = .failure(outcome.message)
-            }
-        } catch {
-            bugReportPhase = .failure(error.localizedDescription)
-        }
-    }
-
-    private func openGitHubIssue() {
-        let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines)
-
-        var bodyParts: [String] = []
-        bodyParts.append("## Describe the bug")
-        bodyParts.append("")
-        if !description.isEmpty {
-            bodyParts.append(description)
-        } else {
-            bodyParts.append("A clear and concise description of what the bug is.")
-        }
-        bodyParts.append("")
-        bodyParts.append("## Environment")
-        bodyParts.append("")
-        bodyParts.append("- macOS Version: \(ProcessInfo.processInfo.operatingSystemVersionString)")
-        bodyParts.append("- EXO Version: \(buildTag) (\(buildCommit))")
-        bodyParts.append("")
-        bodyParts.append("## Additional context")
-        bodyParts.append("")
-        bodyParts.append("A bug report with diagnostic logs was submitted via the app.")
-
-        let body = bodyParts.joined(separator: "\n")
-
-        var components = URLComponents(string: "https://github.com/exo-explore/exo/issues/new")!
-        components.queryItems = [
-            URLQueryItem(name: "template", value: "bug_report.md"),
-            URLQueryItem(name: "title", value: "[BUG] "),
-            URLQueryItem(name: "body", value: body),
-            URLQueryItem(name: "labels", value: "bug"),
-        ]
-
-        if let url = components.url {
-            NSWorkspace.shared.open(url)
-        }
-    }
-
-    private func showUninstallConfirmationAlert() {
-        let alert = NSAlert()
-        alert.messageText = "Uninstall EXO"
-        alert.informativeText = """
-            This will remove EXO and all its system components:
-
-            • Network configuration daemon
-            • Launch at login registration
-            • EXO network location
-
-            The app will be moved to Trash.
-            """
-        alert.alertStyle = .warning
-        alert.addButton(withTitle: "Uninstall")
-        alert.addButton(withTitle: "Cancel")
-
-        // Style the Uninstall button as destructive
-        if let uninstallButton = alert.buttons.first {
-            uninstallButton.hasDestructiveAction = true
-        }
-
-        let response = alert.runModal()
-        if response == .alertFirstButtonReturn {
-            performUninstall()
-        }
-    }
-
-    private func performUninstall() {
-        uninstallInProgress = true
-
-        // Stop EXO process first
-        controller.cancelPendingLaunch()
-        controller.stop()
-        stateService.stopPolling()
-
-        // Run the privileged uninstall on a background thread
-        // Using .utility QoS to avoid priority inversion with NSAppleScript's subprocess
-        DispatchQueue.global(qos: .utility).async {
-            do {
-                // Remove network setup daemon and components (requires admin privileges)
-                try NetworkSetupHelper.uninstall()
-
-                DispatchQueue.main.async {
-                    // Unregister from launch at login
-                    LaunchAtLoginHelper.disable()
-
-                    // Move app to trash
-                    self.moveAppToTrash()
-
-                    // Quit the app
-                    DispatchQueue.main.asyncAfter(deadline: .now() + 0.5) {
-                        NSApplication.shared.terminate(nil)
-                    }
-                }
-            } catch {
-                DispatchQueue.main.async {
-                    self.showErrorAlert(message: error.localizedDescription)
-                    self.uninstallInProgress = false
-                }
-            }
-        }
-    }
-
-    private func showErrorAlert(message: String) {
-        let alert = NSAlert()
-        alert.messageText = "Uninstall Failed"
-        alert.informativeText = message
-        alert.alertStyle = .critical
-        alert.addButton(withTitle: "OK")
-        alert.runModal()
-    }
-
-    private func moveAppToTrash() {
-        guard let appURL = Bundle.main.bundleURL as URL? else { return }
-        do {
-            try FileManager.default.trashItem(at: appURL, resultingItemURL: nil)
-        } catch {
-            // If we can't trash the app, that's OK - user can do it manually
-            // The important system components have already been cleaned up
-        }
-    }
-
     private var buildTag: String {
         Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown"
     }

app/EXO/EXO/ExoProcessController.swift

@@ -5,8 +5,8 @@ import Foundation
 private let customNamespaceKey = "EXOCustomNamespace"
 private let hfTokenKey = "EXOHFToken"
 private let enableImageModelsKey = "EXOEnableImageModels"
-private let offlineModeKey = "EXOOfflineMode"
 private let onboardingCompletedKey = "EXOOnboardingCompleted"
+private let modelSearchPathsKey = "EXOModelSearchPaths"
 
 @MainActor
 final class ExoProcessController: ObservableObject {
@@ -61,12 +61,12 @@ final class ExoProcessController: ObservableObject {
             UserDefaults.standard.set(enableImageModels, forKey: enableImageModelsKey)
         }
     }
-    @Published var offlineMode: Bool = {
-        return UserDefaults.standard.bool(forKey: offlineModeKey)
+    @Published var modelSearchPaths: String = {
+        return UserDefaults.standard.string(forKey: modelSearchPathsKey) ?? ""
     }()
     {
         didSet {
-            UserDefaults.standard.set(offlineMode, forKey: offlineModeKey)
+            UserDefaults.standard.set(modelSearchPaths, forKey: modelSearchPathsKey)
         }
     }
 
@@ -276,8 +276,8 @@ final class ExoProcessController: ObservableObject {
         if enableImageModels {
             environment["EXO_ENABLE_IMAGE_MODELS"] = "true"
         }
-        if offlineMode {
-            environment["EXO_OFFLINE"] = "true"
+        if !modelSearchPaths.isEmpty {
+            environment["EXO_MODELS_PATH"] = modelSearchPaths
         }
 
         var paths: [String] = []

app/EXO/EXO/Services/BugReportService.swift

@@ -38,8 +38,7 @@ struct BugReportService {
     func sendReport(
         baseURL: URL = URL(string: "http://127.0.0.1:52415")!,
         now: Date = Date(),
-        isManual: Bool = false,
-        userDescription: String? = nil
+        isManual: Bool = false
     ) async throws -> BugReportOutcome {
         let timestamp = Self.runTimestampString(now)
         let dayPrefix = Self.dayPrefixString(now)
@@ -61,8 +60,7 @@ struct BugReportService {
             ifconfig: ifconfigText,
             debugInfo: debugInfo,
             isManual: isManual,
-            clusterTbBridgeStatus: clusterTbBridgeStatus,
-            userDescription: userDescription
+            clusterTbBridgeStatus: clusterTbBridgeStatus
         )
 
         let eventLogFiles = readAllEventLogs()
@@ -308,8 +306,7 @@ struct BugReportService {
         ifconfig: String,
         debugInfo: DebugInfo,
         isManual: Bool,
-        clusterTbBridgeStatus: [[String: Any]]?,
-        userDescription: String? = nil
+        clusterTbBridgeStatus: [[String: Any]]?
     ) -> Data? {
         let system = readSystemMetadata()
         let exo = readExoMetadata()
@@ -326,9 +323,6 @@ struct BugReportService {
         if let tbStatus = clusterTbBridgeStatus {
             payload["cluster_thunderbolt_bridge"] = tbStatus
         }
-        if let desc = userDescription, !desc.isEmpty {
-            payload["user_description"] = desc
-        }
         return try? JSONSerialization.data(withJSONObject: payload, options: [.prettyPrinted])
     }
 

app/EXO/EXO/Views/SettingsView.swift

@@ -13,7 +13,7 @@ struct SettingsView: View {
     @State private var pendingNamespace: String = ""
     @State private var pendingHFToken: String = ""
     @State private var pendingEnableImageModels = false
-    @State private var pendingOfflineMode = false
+    @State private var pendingModelSearchPaths: String = ""
     @State private var needsRestart = false
     @State private var bugReportInFlight = false
     @State private var bugReportMessage: String?
@@ -43,7 +43,7 @@ struct SettingsView: View {
             pendingNamespace = controller.customNamespace
             pendingHFToken = controller.hfToken
             pendingEnableImageModels = controller.enableImageModels
-            pendingOfflineMode = controller.offlineMode
+            pendingModelSearchPaths = controller.modelSearchPaths
             needsRestart = false
         }
     }
@@ -74,13 +74,6 @@ struct SettingsView: View {
                     .foregroundColor(.secondary)
             }
 
-            Section {
-                Toggle("Offline Mode", isOn: $pendingOfflineMode)
-                Text("Skip internet checks and use only locally available models.")
-                    .font(.caption)
-                    .foregroundColor(.secondary)
-            }
-
             Section {
                 HStack {
                     Spacer()
@@ -106,6 +99,19 @@ struct SettingsView: View {
                     .foregroundColor(.secondary)
             }
 
+            Section {
+                LabeledContent("Model Search Paths") {
+                    TextField("/path/one:/path/two", text: $pendingModelSearchPaths)
+                        .textFieldStyle(.roundedBorder)
+                        .frame(width: 200)
+                }
+                Text(
+                    "Extra directories to search for pre-downloaded models (colon-separated). HuggingFace cache (~/.cache/huggingface/hub) is checked automatically."
+                )
+                .font(.caption)
+                .foregroundColor(.secondary)
+            }
+
             Section {
                 HStack {
                     Spacer()
@@ -454,22 +460,22 @@ struct SettingsView: View {
 
     private var hasGeneralChanges: Bool {
         pendingNamespace != controller.customNamespace || pendingHFToken != controller.hfToken
-            || pendingOfflineMode != controller.offlineMode
     }
 
     private var hasModelChanges: Bool {
         pendingEnableImageModels != controller.enableImageModels
+            || pendingModelSearchPaths != controller.modelSearchPaths
     }
 
     private func applyGeneralSettings() {
         controller.customNamespace = pendingNamespace
         controller.hfToken = pendingHFToken
-        controller.offlineMode = pendingOfflineMode
         restartIfRunning()
     }
 
     private func applyModelSettings() {
         controller.enableImageModels = pendingEnableImageModels
+        controller.modelSearchPaths = pendingModelSearchPaths
         restartIfRunning()
     }
 

bench/bench.toml

@@ -2,4 +2,6 @@
 #
 # 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/exo_bench.py

@@ -75,7 +75,7 @@ def load_tokenizer_for_bench(model_id: str) -> Any:
         model_path = Path(
             snapshot_download(
                 model_id,
-                allow_patterns=["*.json", "*.py", "*.tiktoken", "*.model"],
+                allow_patterns=["*.json", "*.py", "*.tiktoken"],
             )
         )
 

bench/harness.py

@@ -371,7 +371,7 @@ def run_planning_phase(
                 unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
                     "modelId"
                 ],
-                p["DownloadCompleted"]["total"]["inBytes"],
+                p["DownloadCompleted"]["totalBytes"]["inBytes"],
             )
             for p in node_downloads
             if "DownloadCompleted" in p

bench/pyproject.toml

@@ -4,13 +4,12 @@ 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",
 ]
 
 [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/components/ChatForm.svelte

@@ -7,6 +7,8 @@
     editingImage,
     clearEditingImage,
     selectedChatModel,
+    setSelectedChatModel,
+    instances,
     ttftMs,
     tps,
     totalTokens,
@@ -28,18 +30,6 @@
     modelTasks?: Record<string, string[]>;
     modelCapabilities?: Record<string, string[]>;
     onSend?: () => void;
-    onAutoSend?: (
-      content: string,
-      files?: {
-        id: string;
-        name: string;
-        type: string;
-        textContent?: string;
-        preview?: string;
-      }[],
-    ) => void;
-    onOpenModelPicker?: () => void;
-    modelDisplayOverride?: string;
   }
 
   let {
@@ -51,9 +41,6 @@
     modelTasks = {},
     modelCapabilities = {},
     onSend,
-    onAutoSend,
-    onOpenModelPicker,
-    modelDisplayOverride,
   }: Props = $props();
 
   let message = $state("");
@@ -64,12 +51,27 @@
   const thinkingEnabled = $derived(thinkingEnabledStore());
   let loading = $derived(isLoading());
   const currentModel = $derived(selectedChatModel());
+  const instanceData = $derived(instances());
   const currentTtft = $derived(ttftMs());
   const currentTps = $derived(tps());
   const currentTokens = $derived(totalTokens());
   const currentEditingImage = $derived(editingImage());
   const isEditMode = $derived(currentEditingImage !== null);
 
+  // Custom dropdown state
+  let isModelDropdownOpen = $state(false);
+  let dropdownButtonRef: HTMLButtonElement | undefined = $state();
+  let dropdownPosition = $derived(() => {
+    if (!dropdownButtonRef || !isModelDropdownOpen)
+      return { top: 0, left: 0, width: 0 };
+    const rect = dropdownButtonRef.getBoundingClientRect();
+    return {
+      top: rect.top,
+      left: rect.left,
+      width: rect.width,
+    };
+  });
+
   // Accept all supported file types
   const acceptString = getAcceptString(["image", "text", "pdf"]);
 
@@ -122,14 +124,73 @@
         uploadedFiles.length > 0),
   );
 
-  // Short label for the currently selected model
-  const currentModelLabel = $derived(
-    currentModel
-      ? currentModel.split("/").pop() || currentModel
-      : modelDisplayOverride
-        ? modelDisplayOverride.split("/").pop() || modelDisplayOverride
-        : "",
-  );
+  // Extract available models from running instances
+  const availableModels = $derived(() => {
+    const models: Array<{ id: string; label: string; isImageModel: boolean }> =
+      [];
+    for (const [, instance] of Object.entries(instanceData)) {
+      const modelId = getInstanceModelId(instance);
+      if (
+        modelId &&
+        modelId !== "Unknown" &&
+        !models.some((m) => m.id === modelId)
+      ) {
+        models.push({
+          id: modelId,
+          label: modelId.split("/").pop() || modelId,
+          isImageModel: modelSupportsImageGeneration(modelId),
+        });
+      }
+    }
+    return models;
+  });
+
+  // Track previous model IDs to detect newly added models (plain variable to avoid reactive loop)
+  let previousModelIds: Set<string> = new Set();
+
+  // Auto-select the first available model if none is selected, if current selection is stale, or if a new model is added
+  $effect(() => {
+    const models = availableModels();
+    const currentModelIds = new Set(models.map((m) => m.id));
+
+    if (models.length > 0) {
+      // Find newly added models (in current but not in previous)
+      const newModels = models.filter((m) => !previousModelIds.has(m.id));
+
+      // If no model selected, select the first available
+      if (!currentModel) {
+        setSelectedChatModel(models[0].id);
+      }
+      // If current model is stale (no longer has a running instance), reset to first available
+      else if (!models.some((m) => m.id === currentModel)) {
+        setSelectedChatModel(models[0].id);
+      }
+      // If a new model was just added, select it
+      else if (newModels.length > 0 && previousModelIds.size > 0) {
+        setSelectedChatModel(newModels[0].id);
+      }
+    } else {
+      // No instances running - clear the selected model
+      if (currentModel) {
+        setSelectedChatModel("");
+      }
+    }
+
+    // Update previous model IDs for next comparison
+    previousModelIds = currentModelIds;
+  });
+
+  function getInstanceModelId(instanceWrapped: unknown): string {
+    if (!instanceWrapped || typeof instanceWrapped !== "object") return "";
+    const keys = Object.keys(instanceWrapped as Record<string, unknown>);
+    if (keys.length === 1) {
+      const instance = (instanceWrapped as Record<string, unknown>)[
+        keys[0]
+      ] as { shardAssignments?: { modelId?: string } };
+      return instance?.shardAssignments?.modelId || "";
+    }
+    return "";
+  }
 
   async function handleFiles(files: File[]) {
     if (files.length === 0) return;
@@ -216,15 +277,6 @@
     uploadedFiles = [];
     resetTextareaHeight();
 
-    // When onAutoSend is provided, the parent controls all send logic
-    // (including launching non-running models before sending)
-    if (onAutoSend) {
-      onAutoSend(content, files);
-      onSend?.();
-      setTimeout(() => textareaRef?.focus(), 10);
-      return;
-    }
-
     // Use image editing if in edit mode
     if (isEditMode && currentEditingImage && content) {
       editImage(content, currentEditingImage.imageDataUrl);
@@ -368,7 +420,7 @@
     {/if}
 
     <!-- Model selector (when enabled) -->
-    {#if showModelSelector}
+    {#if showModelSelector && availableModels().length > 0}
       <div
         class="flex items-center justify-between gap-2 px-3 py-2 border-b border-exo-medium-gray/30"
       >
@@ -377,22 +429,33 @@
             class="text-xs text-exo-light-gray uppercase tracking-wider flex-shrink-0"
             >MODEL:</span
           >
-          <!-- Model button — opens the full model picker -->
+          <!-- Custom dropdown -->
           <div class="relative flex-1 max-w-xs">
             <button
+              bind:this={dropdownButtonRef}
               type="button"
-              onclick={() => onOpenModelPicker?.()}
-              class="w-full bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-3 pr-8 py-1.5 text-xs font-mono text-left tracking-wide cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70"
+              onclick={() => (isModelDropdownOpen = !isModelDropdownOpen)}
+              class="w-full bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-3 pr-8 py-1.5 text-xs font-mono text-left tracking-wide cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70 {isModelDropdownOpen
+                ? 'border-exo-yellow/70'
+                : ''}"
             >
-              {#if currentModelLabel}
-                <span class="text-exo-yellow truncate">{currentModelLabel}</span
+              {#if availableModels().find((m) => m.id === currentModel)}
+                <span class="text-exo-yellow truncate"
+                  >{availableModels().find((m) => m.id === currentModel)
+                    ?.label}</span
+                >
+              {:else if availableModels().length > 0}
+                <span class="text-exo-yellow truncate"
+                  >{availableModels()[0].label}</span
                 >
               {:else}
                 <span class="text-exo-light-gray/50">— SELECT MODEL —</span>
               {/if}
             </button>
             <div
-              class="absolute right-2 top-1/2 -translate-y-1/2 pointer-events-none"
+              class="absolute right-2 top-1/2 -translate-y-1/2 pointer-events-none transition-transform duration-200 {isModelDropdownOpen
+                ? 'rotate-180'
+                : ''}"
             >
               <svg
                 class="w-3 h-3 text-exo-yellow/60"
@@ -409,6 +472,78 @@
               </svg>
             </div>
           </div>
+
+          {#if isModelDropdownOpen}
+            <!-- Backdrop to close dropdown -->
+            <button
+              type="button"
+              class="fixed inset-0 z-[9998] cursor-default"
+              onclick={() => (isModelDropdownOpen = false)}
+              aria-label="Close dropdown"
+            ></button>
+
+            <!-- Dropdown Panel - fixed positioning to escape overflow:hidden -->
+            <div
+              class="fixed bg-exo-dark-gray border border-exo-yellow/30 rounded shadow-lg shadow-black/50 z-[9999] max-h-48 overflow-y-auto"
+              style="bottom: calc(100vh - {dropdownPosition()
+                .top}px + 4px); left: {dropdownPosition()
+                .left}px; width: {dropdownPosition().width}px;"
+            >
+              <div class="py-1">
+                {#each availableModels() as model}
+                  <button
+                    type="button"
+                    onclick={() => {
+                      setSelectedChatModel(model.id);
+                      isModelDropdownOpen = false;
+                    }}
+                    class="w-full px-3 py-2 text-left text-xs font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {currentModel ===
+                    model.id
+                      ? 'bg-transparent text-exo-yellow'
+                      : 'text-exo-light-gray hover:text-exo-yellow'}"
+                  >
+                    {#if currentModel === model.id}
+                      <svg
+                        class="w-3 h-3 flex-shrink-0"
+                        fill="currentColor"
+                        viewBox="0 0 20 20"
+                      >
+                        <path
+                          fill-rule="evenodd"
+                          d="M16.707 5.293a1 1 0 010 1.414l-8 8a1 1 0 01-1.414 0l-4-4a1 1 0 011.414-1.414L8 12.586l7.293-7.293a1 1 0 011.414 0z"
+                          clip-rule="evenodd"
+                        />
+                      </svg>
+                    {:else}
+                      <span class="w-3"></span>
+                    {/if}
+                    {#if model.isImageModel}
+                      <svg
+                        class="w-3.5 h-3.5 flex-shrink-0 text-exo-yellow"
+                        fill="none"
+                        viewBox="0 0 24 24"
+                        stroke="currentColor"
+                        stroke-width="2"
+                        aria-label="Image generation model"
+                      >
+                        <rect
+                          x="3"
+                          y="3"
+                          width="18"
+                          height="18"
+                          rx="2"
+                          ry="2"
+                        />
+                        <circle cx="8.5" cy="8.5" r="1.5" />
+                        <polyline points="21 15 16 10 5 21" />
+                      </svg>
+                    {/if}
+                    <span class="truncate flex-1">{model.label}</span>
+                  </button>
+                {/each}
+              </div>
+            </div>
+          {/if}
         </div>
         <!-- Thinking toggle -->
         {#if modelSupportsThinking()}

dashboard/src/lib/components/ChatModelSelector.svelte

@@ -1,401 +0,0 @@
-<script lang="ts" module>
-  export interface ChatModelInfo {
-    id: string;
-    name: string;
-    base_model: string;
-    storage_size_megabytes: number;
-    capabilities: string[];
-    family: string;
-    quantization: string;
-  }
-
-  // Auto mode tier list (for when user just starts typing)
-  export const AUTO_TIERS: string[][] = [
-    // Tier 1 (frontier)
-    ["DeepSeek V3.1", "GLM-5", "Kimi K2.5", "Qwen3 Coder Next"],
-    // Tier 2 (excellent)
-    [
-      "Kimi K2",
-      "Qwen3 235B",
-      "MiniMax M2.5",
-      "Step 3.5 Flash",
-      "Qwen3 Next 80B",
-    ],
-    // Tier 3 (great)
-    [
-      "GLM 4.7",
-      "MiniMax M2.1",
-      "Qwen3 Coder 480B",
-      "GLM 4.5 Air",
-      "Llama 3.3 70B",
-    ],
-    // Tier 4 (good)
-    ["GPT-OSS 120B", "Qwen3 30B", "Llama 3.1 70B", "GLM 4.7 Flash"],
-    // Tier 5 (small/fast)
-    [
-      "Llama 3.1 8B",
-      "GPT-OSS 20B",
-      "Llama 3.2 3B",
-      "Qwen3 0.6B",
-      "Llama 3.2 1B",
-    ],
-  ];
-
-  /** Return the tier index (0 = best) for a base_model name. */
-  export function getAutoTierIndex(baseModel: string): number {
-    for (let i = 0; i < AUTO_TIERS.length; i++) {
-      if (AUTO_TIERS[i].includes(baseModel)) return i;
-    }
-    return AUTO_TIERS.length; // not in any tier → lowest priority
-  }
-
-  /** Auto mode: walk tiers top-down, pick biggest fitting variant from highest tier. */
-  export function pickAutoModel(
-    modelList: ChatModelInfo[],
-    memoryGB: number,
-  ): ChatModelInfo | null {
-    for (const tier of AUTO_TIERS) {
-      const candidates: ChatModelInfo[] = [];
-      for (const baseModel of tier) {
-        const variants = modelList
-          .filter(
-            (m) =>
-              m.base_model === baseModel &&
-              (m.storage_size_megabytes || 0) / 1024 <= memoryGB &&
-              (m.storage_size_megabytes || 0) > 0,
-          )
-          .sort(
-            (a, b) =>
-              (b.storage_size_megabytes || 0) - (a.storage_size_megabytes || 0),
-          );
-        if (variants[0]) candidates.push(variants[0]);
-      }
-      if (candidates.length > 0) {
-        candidates.sort(
-          (a, b) =>
-            (b.storage_size_megabytes || 0) - (a.storage_size_megabytes || 0),
-        );
-        return candidates[0];
-      }
-    }
-    return null;
-  }
-</script>
-
-<script lang="ts">
-  interface CategoryRecommendation {
-    category: "coding" | "writing" | "agentic" | "biggest";
-    label: string;
-    model: ChatModelInfo | null;
-    tooltip: string;
-  }
-
-  interface Props {
-    models: ChatModelInfo[];
-    clusterLabel: string;
-    totalMemoryGB: number;
-    onSelect: (modelId: string, category: string) => void;
-    onAddModel: () => void;
-    class?: string;
-  }
-
-  let {
-    models,
-    clusterLabel,
-    totalMemoryGB,
-    onSelect,
-    onAddModel,
-    class: className = "",
-  }: Props = $props();
-
-  // --- Hardcoded Rankings ---
-  const CODING_RANKING = [
-    "Qwen3 Coder Next",
-    "Qwen3 Coder 480B",
-    "Qwen3 30B",
-    "GPT-OSS 20B",
-    "Llama 3.1 8B",
-    "Llama 3.2 3B",
-    "Qwen3 0.6B",
-  ];
-
-  const WRITING_RANKING = [
-    "Kimi K2.5",
-    "Kimi K2",
-    "Qwen3 Next 80B",
-    "Llama 3.3 70B",
-    "MiniMax M2.5",
-    "GLM 4.5 Air",
-    "GLM 4.7 Flash",
-    "GPT-OSS 20B",
-    "Llama 3.1 8B",
-    "Llama 3.2 3B",
-    "Qwen3 0.6B",
-  ];
-
-  const AGENTIC_RANKING = [
-    "DeepSeek V3.1",
-    "GLM-5",
-    "Qwen3 235B",
-    "Step 3.5 Flash",
-    "GLM 4.7",
-    "MiniMax M2.1",
-    "GPT-OSS 120B",
-    "Llama 3.3 70B",
-    "Llama 3.1 70B",
-    "GLM 4.7 Flash",
-    "GPT-OSS 20B",
-    "Qwen3 30B",
-    "Llama 3.1 8B",
-    "Llama 3.2 3B",
-    "Qwen3 0.6B",
-  ];
-
-  function getModelSizeGB(m: ChatModelInfo): number {
-    return (m.storage_size_megabytes || 0) / 1024;
-  }
-
-  function fitsInMemory(m: ChatModelInfo): boolean {
-    return getModelSizeGB(m) <= totalMemoryGB && getModelSizeGB(m) > 0;
-  }
-
-  /** For a given base_model name, find the biggest quant variant that fits in memory. */
-  function pickBestVariant(baseModel: string): ChatModelInfo | null {
-    const variants = models
-      .filter((m) => m.base_model === baseModel && fitsInMemory(m))
-      .sort((a, b) => getModelSizeGB(b) - getModelSizeGB(a));
-    return variants[0] ?? null;
-  }
-
-  /** Walk a ranked list of base_model names, return the first that has a fitting variant. */
-  function pickFromRanking(ranking: string[]): ChatModelInfo | null {
-    for (const baseModel of ranking) {
-      const pick = pickBestVariant(baseModel);
-      if (pick) return pick;
-    }
-    return null;
-  }
-
-  /** Pick the single biggest model that fits. */
-  function pickBiggest(): ChatModelInfo | null {
-    const fitting = models
-      .filter((m) => fitsInMemory(m))
-      .sort((a, b) => getModelSizeGB(b) - getModelSizeGB(a));
-    return fitting[0] ?? null;
-  }
-
-  const recommendations = $derived.by((): CategoryRecommendation[] => {
-    return [
-      {
-        category: "coding",
-        label: "Best for Coding",
-        model: pickFromRanking(CODING_RANKING),
-        tooltip:
-          "Ranked by coding benchmark performance (LiveCodeBench, SWE-bench)",
-      },
-      {
-        category: "writing",
-        label: "Best for Writing",
-        model: pickFromRanking(WRITING_RANKING),
-        tooltip: "Ranked by creative writing quality and instruction following",
-      },
-      {
-        category: "agentic",
-        label: "Best Agentic",
-        model: pickFromRanking(AGENTIC_RANKING),
-        tooltip: "Ranked by reasoning, planning, and tool-use capability",
-      },
-      {
-        category: "biggest",
-        label: "Biggest",
-        model: pickBiggest(),
-        tooltip: "Largest model that fits in your available memory",
-      },
-    ];
-  });
-
-  function formatSize(mb: number): string {
-    const gb = mb / 1024;
-    if (gb >= 100) return `${Math.round(gb)} GB`;
-    return `${gb.toFixed(1)} GB`;
-  }
-
-  // Category icons (SVG paths)
-  const categoryIcons: Record<string, string> = {
-    coding:
-      "M8 9l3 3-3 3m5 0h3M5 20h14a2 2 0 002-2V6a2 2 0 00-2-2H5a2 2 0 00-2 2v12a2 2 0 002 2z",
-    writing:
-      "M11 5H6a2 2 0 00-2 2v11a2 2 0 002 2h11a2 2 0 002-2v-5m-1.414-9.414a2 2 0 112.828 2.828L11.828 15H9v-2.828l8.586-8.586z",
-    agentic:
-      "M9.663 17h4.673M12 3v1m6.364 1.636l-.707.707M21 12h-1M4 12H3m3.343-5.657l-.707-.707m2.828 9.9a5 5 0 117.072 0l-.548.547A3.374 3.374 0 0014 18.469V19a2 2 0 11-4 0v-.531c0-.895-.356-1.754-.988-2.386l-.548-.547z",
-    biggest:
-      "M19 11H5m14 0a2 2 0 012 2v6a2 2 0 01-2 2H5a2 2 0 01-2-2v-6a2 2 0 012-2m14 0V9a2 2 0 00-2-2M5 11V9a2 2 0 012-2m0 0V5a2 2 0 012-2h6a2 2 0 012 2v2M7 7h10",
-  };
-
-  let hoveredTooltip = $state<string | null>(null);
-  let tooltipAnchor = $state<{ x: number; y: number } | null>(null);
-
-  function showTooltip(category: string, e: MouseEvent | FocusEvent) {
-    hoveredTooltip = category;
-    const target = e.currentTarget as HTMLElement;
-    const rect = target.getBoundingClientRect();
-    tooltipAnchor = { x: rect.left + rect.width / 2, y: rect.top };
-  }
-
-  function hideTooltip() {
-    hoveredTooltip = null;
-    tooltipAnchor = null;
-  }
-</script>
-
-<div class="flex flex-col items-center justify-center gap-6 {className}">
-  <!-- Header -->
-  <div class="text-center">
-    <p class="text-xs text-exo-light-gray uppercase tracking-[0.2em] mb-1">
-      Recommended for your
-    </p>
-    <p class="text-sm text-white font-mono tracking-wide">{clusterLabel}</p>
-  </div>
-
-  <!-- Category Cards Grid -->
-  <div class="grid grid-cols-2 gap-3 w-full max-w-md">
-    {#each recommendations as rec}
-      {#if rec.model}
-        <button
-          type="button"
-          onclick={() => rec.model && onSelect(rec.model.id, rec.category)}
-          class="group relative flex flex-col items-start gap-2 p-4 rounded-lg border border-exo-medium-gray/50 bg-exo-dark-gray/50 hover:border-exo-yellow/40 hover:bg-exo-dark-gray transition-all duration-200 cursor-pointer text-left"
-        >
-          <!-- Category icon + label -->
-          <div class="flex items-center gap-2 w-full">
-            <svg
-              class="w-4 h-4 text-exo-yellow/70 group-hover:text-exo-yellow transition-colors flex-shrink-0"
-              fill="none"
-              viewBox="0 0 24 24"
-              stroke="currentColor"
-              stroke-width="1.5"
-            >
-              <path
-                stroke-linecap="round"
-                stroke-linejoin="round"
-                d={categoryIcons[rec.category]}
-              />
-            </svg>
-            <span
-              class="text-xs font-mono uppercase tracking-wider text-exo-light-gray group-hover:text-white transition-colors"
-            >
-              {rec.label}
-            </span>
-            <!-- Info tooltip -->
-            <div class="ml-auto flex-shrink-0">
-              <span
-                role="button"
-                tabindex="-1"
-                class="text-exo-light-gray/40 hover:text-exo-light-gray transition-colors cursor-help inline-flex"
-                onmouseenter={(e: MouseEvent) => showTooltip(rec.category, e)}
-                onmouseleave={() => hideTooltip()}
-                onclick={(e: MouseEvent) => {
-                  e.stopPropagation();
-                  if (hoveredTooltip === rec.category) {
-                    hideTooltip();
-                  } else {
-                    showTooltip(rec.category, e);
-                  }
-                }}
-              >
-                <svg
-                  class="w-3.5 h-3.5"
-                  fill="none"
-                  viewBox="0 0 24 24"
-                  stroke="currentColor"
-                  stroke-width="2"
-                >
-                  <circle cx="12" cy="12" r="10" />
-                  <path d="M12 16v-4m0-4h.01" />
-                </svg>
-              </span>
-            </div>
-          </div>
-
-          <!-- Model name + size -->
-          <div class="w-full">
-            <p class="text-sm text-white font-mono truncate">
-              {rec.model.base_model}
-            </p>
-            <p class="text-xs text-exo-light-gray/60 font-mono mt-0.5">
-              {formatSize(rec.model.storage_size_megabytes)}
-              {#if rec.model.quantization}
-                <span class="text-exo-light-gray/40"
-                  >&middot; {rec.model.quantization}</span
-                >
-              {/if}
-            </p>
-          </div>
-        </button>
-      {:else}
-        <!-- No model fits for this category -->
-        <div
-          class="flex flex-col items-start gap-2 p-4 rounded-lg border border-exo-medium-gray/30 bg-exo-dark-gray/30 opacity-50"
-        >
-          <div class="flex items-center gap-2">
-            <svg
-              class="w-4 h-4 text-exo-light-gray/40 flex-shrink-0"
-              fill="none"
-              viewBox="0 0 24 24"
-              stroke="currentColor"
-              stroke-width="1.5"
-            >
-              <path
-                stroke-linecap="round"
-                stroke-linejoin="round"
-                d={categoryIcons[rec.category]}
-              />
-            </svg>
-            <span
-              class="text-xs font-mono uppercase tracking-wider text-exo-light-gray/50"
-              >{rec.label}</span
-            >
-          </div>
-          <p class="text-xs text-exo-light-gray/40 font-mono">No model fits</p>
-        </div>
-      {/if}
-    {/each}
-  </div>
-
-  <!-- Add Model Button -->
-  <button
-    type="button"
-    onclick={onAddModel}
-    class="flex items-center gap-2 px-4 py-2 text-xs font-mono uppercase tracking-wider text-exo-light-gray hover:text-exo-yellow border border-exo-medium-gray/30 hover:border-exo-yellow/30 rounded-lg transition-all duration-200 cursor-pointer"
-  >
-    <svg
-      class="w-3.5 h-3.5"
-      fill="none"
-      viewBox="0 0 24 24"
-      stroke="currentColor"
-      stroke-width="2"
-    >
-      <path stroke-linecap="round" stroke-linejoin="round" d="M12 4v16m8-8H4" />
-    </svg>
-    Add Model
-  </button>
-
-  <!-- Auto hint -->
-  <p class="text-xs text-exo-light-gray/40 font-mono tracking-wide text-center">
-    Or just start typing &mdash; we'll pick the best model automatically
-  </p>
-</div>
-
-<!-- Fixed-position tooltip (escapes overflow-hidden ancestors) -->
-{#if hoveredTooltip && tooltipAnchor}
-  {@const rec = recommendations.find((r) => r.category === hoveredTooltip)}
-  {#if rec}
-    <div
-      class="fixed z-[9999] px-3 py-2 bg-exo-black border border-exo-medium-gray/50 rounded text-xs text-exo-light-gray whitespace-nowrap shadow-lg pointer-events-none"
-      style="left: {tooltipAnchor.x}px; top: {tooltipAnchor.y -
-        8}px; transform: translate(-50%, -100%);"
-    >
-      {rec.tooltip}
-    </div>
-  {/if}
-{/if}

dashboard/src/lib/components/ChatSidebar.svelte

@@ -2,6 +2,7 @@
   import {
     conversations,
     activeConversationId,
+    createConversation,
     loadConversation,
     deleteConversation,
     deleteAllConversations,
@@ -16,15 +17,9 @@
 
   interface Props {
     class?: string;
-    onNewChat?: () => void;
-    onSelectConversation?: () => void;
   }
 
-  let {
-    class: className = "",
-    onNewChat,
-    onSelectConversation,
-  }: Props = $props();
+  let { class: className = "" }: Props = $props();
 
   const conversationList = $derived(conversations());
   const activeId = $derived(activeConversationId());
@@ -47,11 +42,10 @@
   );
 
   function handleNewChat() {
-    onNewChat?.();
+    createConversation();
   }
 
   function handleSelectConversation(id: string) {
-    onSelectConversation?.();
     loadConversation(id);
   }
 

dashboard/src/lib/components/DeviceIcon.svelte

@@ -67,8 +67,8 @@
   const studioMemH = $derived((ramPercent / 100) * studioMemTotalH);
 
   // ── MacBook dimensions (same ratios as TopologyGraph) ──
-  const mbW = $derived((size * 1.6 * 0.85) / 1.15);
-  const mbH = $derived(size * 0.85);
+  const mbW = $derived(size * 1.6);
+  const mbH = $derived(size * 1.15);
   const mbX = $derived(cx - mbW / 2);
   const mbY = $derived(cy - mbH / 2);
 

dashboard/src/lib/components/HeaderNav.svelte

@@ -95,7 +95,7 @@
     {#if showHome}
       <button
         onclick={handleHome}
-        class="text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
+        class="text-sm text-exo-light-gray hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
         title="Back to topology view"
       >
         <svg
@@ -116,7 +116,7 @@
     {/if}
     <a
       href="/#/downloads"
-      class="text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
+      class="text-sm text-exo-light-gray hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
       title="View downloads overview"
     >
       {#if downloadProgress}

dashboard/src/lib/components/ModelFilterPopover.svelte

@@ -6,7 +6,6 @@
     capabilities: string[];
     sizeRange: { min: number; max: number } | null;
     downloadedOnly: boolean;
-    readyOnly: boolean;
   }
 
   type ModelFilterPopoverProps = {
@@ -191,58 +190,34 @@
       </div>
     </div>
 
-    <!-- Availability filters -->
+    <!-- Downloaded only -->
     <div>
       <h4 class="text-xs font-mono text-white/50 mb-2">Availability</h4>
-      <div class="flex flex-wrap gap-1.5">
-        <button
-          type="button"
-          class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.downloadedOnly
-            ? 'bg-green-500/20 text-green-400 border border-green-500/30'
-            : 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
-          onclick={() =>
-            onChange({ ...filters, downloadedOnly: !filters.downloadedOnly })}
+      <button
+        type="button"
+        class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.downloadedOnly
+          ? 'bg-green-500/20 text-green-400 border border-green-500/30'
+          : 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
+        onclick={() =>
+          onChange({ ...filters, downloadedOnly: !filters.downloadedOnly })}
+      >
+        <svg
+          class="w-3.5 h-3.5 inline-block"
+          viewBox="0 0 24 24"
+          fill="none"
+          stroke="currentColor"
+          stroke-width="2"
+          stroke-linecap="round"
+          stroke-linejoin="round"
         >
-          <svg
-            class="w-3.5 h-3.5 inline-block"
-            viewBox="0 0 24 24"
-            fill="none"
-            stroke="currentColor"
-            stroke-width="2"
-            stroke-linecap="round"
-            stroke-linejoin="round"
-          >
-            <path
-              class="text-white/40"
-              d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
-            />
-            <path class="text-green-400" d="m9 13 2 2 4-4" />
-          </svg>
-          <span class="ml-1">Downloaded</span>
-        </button>
-        <button
-          type="button"
-          class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.readyOnly
-            ? 'bg-green-500/20 text-green-400 border border-green-500/30'
-            : 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
-          onclick={() =>
-            onChange({ ...filters, readyOnly: !filters.readyOnly })}
-        >
-          <svg
-            class="w-3.5 h-3.5 inline-block"
-            viewBox="0 0 24 24"
-            fill="none"
-            stroke="currentColor"
-            stroke-width="2"
-            stroke-linecap="round"
-            stroke-linejoin="round"
-          >
-            <circle cx="12" cy="12" r="10" />
-            <path d="m9 12 2 2 4-4" />
-          </svg>
-          <span class="ml-1">Ready</span>
-        </button>
-      </div>
+          <path
+            class="text-white/40"
+            d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
+          />
+          <path class="text-green-400" d="m9 13 2 2 4-4" />
+        </svg>
+        <span class="ml-1">Downloaded</span>
+      </button>
     </div>
 
     <!-- Size range -->

dashboard/src/lib/components/ModelPickerGroup.svelte

@@ -41,7 +41,6 @@
     onShowInfo: (group: ModelGroup) => void;
     downloadStatusMap?: Map<string, DownloadAvailability>;
     launchedAt?: number;
-    instanceStatuses?: Record<string, { status: string; statusClass: string }>;
   };
 
   let {
@@ -57,7 +56,6 @@
     onShowInfo,
     downloadStatusMap,
     launchedAt,
-    instanceStatuses = {},
   }: ModelPickerGroupProps = $props();
 
   // Group-level download status: show if any variant is downloaded
@@ -94,12 +92,6 @@
   const anyVariantFits = $derived(
     group.variants.some((v) => canModelFit(v.id)),
   );
-  // Check if any variant has an active instance (ready, loading, downloading)
-  const anyVariantHasInstance = $derived(
-    instanceStatuses
-      ? group.variants.some((v) => instanceStatuses[v.id] != null)
-      : false,
-  );
   const groupFitStatus = $derived.by((): ModelFitStatus => {
     let hasClusterCapacityOnly = false;
     for (const variant of group.variants) {
@@ -130,35 +122,16 @@
     !group.hasMultipleVariants &&
       group.variants.some((v) => v.id === selectedModelId),
   );
-
-  // Group-level instance status: show the "best" status across all variants
-  const groupInstanceStatus = $derived.by(() => {
-    if (!instanceStatuses) return null;
-    const readyStatuses = ["READY", "LOADED", "RUNNING"];
-    const loadingStatuses = ["LOADING", "WARMING UP"];
-    let bestStatus: { status: string; statusClass: string } | null = null;
-    for (const variant of group.variants) {
-      const s = instanceStatuses[variant.id];
-      if (!s) continue;
-      if (readyStatuses.includes(s.status)) return s; // Ready is best
-      if (loadingStatuses.includes(s.status) || s.status === "DOWNLOADING") {
-        bestStatus = s;
-      }
-    }
-    return bestStatus;
-  });
 </script>
 
 <div
-  class="border-b border-white/5 last:border-b-0 {!anyVariantFits &&
-  !anyVariantHasInstance
+  class="border-b border-white/5 last:border-b-0 {!anyVariantFits
     ? 'opacity-50'
     : ''}"
 >
   <!-- Main row -->
   <div
-    class="flex items-center gap-2 px-3 py-2.5 transition-colors {anyVariantFits ||
-    anyVariantHasInstance
+    class="flex items-center gap-2 px-3 py-2.5 transition-colors {anyVariantFits
       ? 'hover:bg-white/5 cursor-pointer'
       : 'cursor-not-allowed'} {isMainSelected
       ? 'bg-exo-yellow/10 border-l-2 border-exo-yellow'
@@ -168,7 +141,7 @@
         onToggleExpand();
       } else {
         const modelId = group.variants[0]?.id;
-        if (modelId && (canModelFit(modelId) || instanceStatuses[modelId])) {
+        if (modelId && canModelFit(modelId)) {
           onSelectModel(modelId);
         }
       }
@@ -182,7 +155,7 @@
           onToggleExpand();
         } else {
           const modelId = group.variants[0]?.id;
-          if (modelId && (canModelFit(modelId) || instanceStatuses[modelId])) {
+          if (modelId && canModelFit(modelId)) {
             onSelectModel(modelId);
           }
         }
@@ -373,47 +346,6 @@
       </span>
     {/if}
 
-    <!-- Instance status badge -->
-    {#if groupInstanceStatus}
-      {#if groupInstanceStatus.status === "READY" || groupInstanceStatus.status === "LOADED" || groupInstanceStatus.status === "RUNNING"}
-        <span class="flex-shrink-0" title="Running">
-          <svg
-            class="w-3 h-3 text-green-400"
-            viewBox="0 0 12 12"
-            fill="currentColor"
-          >
-            <circle cx="6" cy="6" r="5" />
-          </svg>
-        </span>
-      {:else if groupInstanceStatus.status === "DOWNLOADING"}
-        <span class="flex-shrink-0 animate-pulse" title="Downloading">
-          <svg
-            class="w-3.5 h-3.5 text-blue-400"
-            viewBox="0 0 24 24"
-            fill="none"
-            stroke="currentColor"
-            stroke-width="2"
-            stroke-linecap="round"
-            stroke-linejoin="round"
-          >
-            <path d="M21 15v4a2 2 0 01-2 2H5a2 2 0 01-2-2v-4" />
-            <polyline points="7 10 12 15 17 10" />
-            <line x1="12" y1="15" x2="12" y2="3" />
-          </svg>
-        </span>
-      {:else if groupInstanceStatus.status === "LOADING" || groupInstanceStatus.status === "WARMING UP"}
-        <span class="flex-shrink-0 animate-pulse" title="Loading">
-          <svg
-            class="w-3 h-3 text-yellow-400"
-            viewBox="0 0 12 12"
-            fill="currentColor"
-          >
-            <circle cx="6" cy="6" r="5" />
-          </svg>
-        </span>
-      {/if}
-    {/if}
-
     <!-- Check mark if selected (single-variant) -->
     {#if isMainSelected}
       <svg
@@ -488,30 +420,20 @@
       {#each group.variants as variant}
         {@const fitStatus = getModelFitStatus(variant.id)}
         {@const modelCanFit = canModelFit(variant.id)}
-        {@const variantHasInstance = instanceStatuses[variant.id] != null}
         {@const isSelected = selectedModelId === variant.id}
-        <div
-          class="w-full flex items-center gap-3 px-3 py-2 pl-10 hover:bg-white/5 transition-colors text-left {!modelCanFit &&
-          !variantHasInstance
+        <button
+          type="button"
+          class="w-full flex items-center gap-3 px-3 py-2 pl-10 hover:bg-white/5 transition-colors text-left {!modelCanFit
             ? 'opacity-50 cursor-not-allowed'
             : 'cursor-pointer'} {isSelected
             ? 'bg-exo-yellow/10 border-l-2 border-exo-yellow'
             : 'border-l-2 border-transparent'}"
-          role="button"
-          tabindex="0"
+          disabled={!modelCanFit}
           onclick={() => {
-            if (modelCanFit || variantHasInstance) {
+            if (modelCanFit) {
               onSelectModel(variant.id);
             }
           }}
-          onkeydown={(e) => {
-            if (e.key === "Enter" || e.key === " ") {
-              e.preventDefault();
-              if (modelCanFit) {
-                onSelectModel(variant.id);
-              }
-            }
-          }}
         >
           <!-- Quantization badge -->
           <span
@@ -556,48 +478,6 @@
             {/if}
           {/if}
 
-          <!-- Instance status badge -->
-          {#if instanceStatuses[variant.id]}
-            {@const instStatus = instanceStatuses[variant.id]}
-            {#if instStatus.status === "READY" || instStatus.status === "LOADED" || instStatus.status === "RUNNING"}
-              <span class="flex-shrink-0" title="Running">
-                <svg
-                  class="w-3 h-3 text-green-400"
-                  viewBox="0 0 12 12"
-                  fill="currentColor"
-                >
-                  <circle cx="6" cy="6" r="5" />
-                </svg>
-              </span>
-            {:else if instStatus.status === "DOWNLOADING"}
-              <span class="flex-shrink-0 animate-pulse" title="Downloading">
-                <svg
-                  class="w-3.5 h-3.5 text-blue-400"
-                  viewBox="0 0 24 24"
-                  fill="none"
-                  stroke="currentColor"
-                  stroke-width="2"
-                  stroke-linecap="round"
-                  stroke-linejoin="round"
-                >
-                  <path d="M21 15v4a2 2 0 01-2 2H5a2 2 0 01-2-2v-4" />
-                  <polyline points="7 10 12 15 17 10" />
-                  <line x1="12" y1="15" x2="12" y2="3" />
-                </svg>
-              </span>
-            {:else if instStatus.status === "LOADING" || instStatus.status === "WARMING UP"}
-              <span class="flex-shrink-0 animate-pulse" title="Loading">
-                <svg
-                  class="w-3 h-3 text-yellow-400"
-                  viewBox="0 0 12 12"
-                  fill="currentColor"
-                >
-                  <circle cx="6" cy="6" r="5" />
-                </svg>
-              </span>
-            {/if}
-          {/if}
-
           <!-- Check mark if selected -->
           {#if isSelected}
             <svg
@@ -610,36 +490,7 @@
               />
             </svg>
           {/if}
-
-          <!-- Info button -->
-          <button
-            type="button"
-            class="p-1 rounded hover:bg-white/10 transition-colors flex-shrink-0"
-            onclick={(e) => {
-              e.stopPropagation();
-              onShowInfo({
-                id: variant.id,
-                name: variant.name || variant.id,
-                capabilities: group.capabilities,
-                family: group.family,
-                variants: [variant],
-                smallestVariant: variant,
-                hasMultipleVariants: false,
-              });
-            }}
-            title="View variant details"
-          >
-            <svg
-              class="w-4 h-4 text-white/30 hover:text-white/50"
-              viewBox="0 0 24 24"
-              fill="currentColor"
-            >
-              <path
-                d="M12 2C6.48 2 2 6.48 2 12s4.48 10 10 10 10-4.48 10-10S17.52 2 12 2zm1 15h-2v-6h2v6zm0-8h-2V7h2v2z"
-              />
-            </svg>
-          </button>
-        </div>
+        </button>
       {/each}
     </div>
   {/if}

dashboard/src/lib/components/ModelPickerModal.svelte

@@ -36,7 +36,6 @@
     capabilities: string[];
     sizeRange: { min: number; max: number } | null;
     downloadedOnly: boolean;
-    readyOnly: boolean;
   }
 
   interface HuggingFaceModel {
@@ -50,11 +49,6 @@
 
   type ModelFitStatus = "fits_now" | "fits_cluster_capacity" | "too_large";
 
-  export type InstanceStatus = {
-    status: string;
-    statusClass: string;
-  };
-
   type ModelPickerModalProps = {
     isOpen: boolean;
     models: ModelInfo[];
@@ -81,7 +75,6 @@
         macmon_info?: { memory?: { ram_total?: number } };
       }
     >;
-    instanceStatuses?: Record<string, InstanceStatus>;
   };
 
   let {
@@ -103,7 +96,6 @@
     usedMemoryGB,
     downloadsData,
     topologyNodes,
-    instanceStatuses = {},
   }: ModelPickerModalProps = $props();
 
   // Local state
@@ -115,7 +107,6 @@
     capabilities: [],
     sizeRange: null,
     downloadedOnly: false,
-    readyOnly: false,
   });
   let infoGroup = $state<ModelGroup | null>(null);
 
@@ -449,16 +440,6 @@
       );
     }
 
-    // Filter to ready/running models only
-    if (filters.readyOnly) {
-      result = result.filter((g) =>
-        g.variants.some((v) => {
-          const s = instanceStatuses[v.id];
-          return s && s.statusClass === "ready";
-        }),
-      );
-    }
-
     // Sort: fits-now first, then fits-cluster-capacity, then too-large
     result.sort((a, b) => {
       const getGroupFitRank = (group: ModelGroup): number => {
@@ -569,19 +550,13 @@
   }
 
   function clearFilters() {
-    filters = {
-      capabilities: [],
-      sizeRange: null,
-      downloadedOnly: false,
-      readyOnly: false,
-    };
+    filters = { capabilities: [], sizeRange: null, downloadedOnly: false };
   }
 
   const hasActiveFilters = $derived(
     filters.capabilities.length > 0 ||
       filters.sizeRange !== null ||
-      filters.downloadedOnly ||
-      filters.readyOnly,
+      filters.downloadedOnly,
   );
 </script>
 
@@ -850,7 +825,6 @@
                 onShowInfo={(g) => (infoGroup = g)}
                 downloadStatusMap={getVariantDownloadMap(group)}
                 launchedAt={recentTimestamps.get(group.variants[0]?.id ?? "")}
-                {instanceStatuses}
               />
             {/each}
           {/if}
@@ -918,7 +892,6 @@
               {onToggleFavorite}
               onShowInfo={(g) => (infoGroup = g)}
               downloadStatusMap={getVariantDownloadMap(group)}
-              {instanceStatuses}
             />
           {/each}
           <!-- Other models -->
@@ -945,7 +918,6 @@
               {onToggleFavorite}
               onShowInfo={(g) => (infoGroup = g)}
               downloadStatusMap={getVariantDownloadMap(group)}
-              {instanceStatuses}
             />
           {/each}
         {/if}
@@ -968,11 +940,6 @@
             >Downloaded</span
           >
         {/if}
-        {#if filters.readyOnly}
-          <span class="px-1.5 py-0.5 bg-green-500/20 text-green-400 rounded"
-            >Ready</span
-          >
-        {/if}
         {#if filters.sizeRange}
           <span class="px-1.5 py-0.5 bg-exo-yellow/20 text-exo-yellow rounded">
             {filters.sizeRange.min}GB - {filters.sizeRange.max}GB

dashboard/src/lib/components/index.ts

@@ -12,4 +12,3 @@ export { default as HuggingFaceResultItem } from "./HuggingFaceResultItem.svelte
 export { default as ModelFilterPopover } from "./ModelFilterPopover.svelte";
 export { default as ModelPickerGroup } from "./ModelPickerGroup.svelte";
 export { default as ModelPickerModal } from "./ModelPickerModal.svelte";
-export { default as ChatModelSelector } from "./ChatModelSelector.svelte";

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

@@ -847,10 +847,6 @@ class AppStore {
     this.thinkingEnabled = conversation.enableThinking ?? true;
     this.refreshConversationModelFromInstances();
 
-    // Sync global selection to the loaded conversation's model so reactive
-    // effects in +page.svelte can determine the correct chat launch state.
-    this.selectedChatModel = conversation.modelId || "";
-
     return true;
   }
 
@@ -3158,23 +3154,6 @@ 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)
    */
@@ -3318,5 +3297,3 @@ 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/downloads/+page.svelte

@@ -412,7 +412,7 @@
                   <div>{col.label}</div>
                   {#if col.diskAvailable != null}
                     <div
-                      class="text-[9px] text-white/70 normal-case tracking-normal mt-0.5"
+                      class="text-[9px] text-exo-light-gray/60 normal-case tracking-normal mt-0.5"
                     >
                       {formatBytes(col.diskAvailable)} free
                     </div>
@@ -436,7 +436,7 @@
                       </div>
                       {#if row.prettyName}
                         <div
-                          class="text-[10px] text-white/60"
+                          class="text-[10px] text-exo-light-gray/60"
                           title={row.modelId}
                         >
                           {row.modelId}
@@ -450,7 +450,7 @@
                       title="View model details"
                     >
                       <svg
-                        class="w-4 h-4 text-white/60 hover:text-white/80"
+                        class="w-4 h-4 text-white/30 hover:text-white/60"
                         viewBox="0 0 24 24"
                         fill="currentColor"
                       >
@@ -469,11 +469,11 @@
                   <td class="px-4 py-3 text-center align-middle">
                     {#if cell.kind === "completed"}
                       <div
-                        class="flex flex-col items-center gap-1"
+                        class="flex flex-col items-center gap-0.5"
                         title="Completed ({formatBytes(cell.totalBytes)})"
                       >
                         <svg
-                          class="w-7 h-7 text-green-400"
+                          class="w-5 h-5 text-green-400"
                           viewBox="0 0 20 20"
                           fill="currentColor"
                         >
@@ -483,18 +483,18 @@
                             clip-rule="evenodd"
                           ></path>
                         </svg>
-                        <span class="text-xs text-white/70"
+                        <span class="text-[10px] text-exo-light-gray/70"
                           >{formatBytes(cell.totalBytes)}</span
                         >
                         <button
                           type="button"
-                          class="text-white/50 hover:text-red-400 transition-colors mt-0.5 cursor-pointer"
+                          class="text-exo-light-gray/40 hover:text-red-400 transition-colors mt-0.5"
                           onclick={() =>
                             deleteDownload(col.nodeId, row.modelId)}
                           title="Delete from this node"
                         >
                           <svg
-                            class="w-5 h-5"
+                            class="w-3.5 h-3.5"
                             viewBox="0 0 20 20"
                             fill="none"
                             stroke="currentColor"
@@ -517,11 +517,11 @@
                           cell.speed,
                         )} - ETA {formatEta(cell.etaMs)}"
                       >
-                        <span class="text-exo-yellow text-sm font-medium"
+                        <span class="text-exo-yellow text-xs font-medium"
                           >{clampPercent(cell.percentage).toFixed(1)}%</span
                         >
                         <div
-                          class="w-16 h-2 bg-exo-black/60 rounded-sm overflow-hidden"
+                          class="w-14 h-1.5 bg-exo-black/60 rounded-sm overflow-hidden"
                         >
                           <div
                             class="h-full bg-gradient-to-r from-exo-yellow to-exo-yellow/70 transition-all duration-300"
@@ -530,25 +530,25 @@
                             ).toFixed(1)}%"
                           ></div>
                         </div>
-                        <span class="text-[10px] text-white/70"
+                        <span class="text-[9px] text-exo-light-gray/60"
                           >{formatSpeed(cell.speed)}</span
                         >
                       </div>
                     {:else if cell.kind === "pending"}
                       <div
-                        class="flex flex-col items-center gap-1"
+                        class="flex flex-col items-center gap-0.5"
                         title={cell.downloaded > 0
-                          ? `${formatBytes(cell.downloaded)} / ${formatBytes(cell.total)} downloaded (paused)`
+                          ? `${formatBytes(cell.downloaded)} / ${formatBytes(cell.total)} downloaded`
                           : "Download pending"}
                       >
                         {#if cell.downloaded > 0 && cell.total > 0}
-                          <span class="text-white/70 text-xs"
+                          <span class="text-exo-light-gray/70 text-[10px]"
                             >{formatBytes(cell.downloaded)} / {formatBytes(
                               cell.total,
                             )}</span
                           >
                           <div
-                            class="w-full h-1.5 bg-white/10 rounded-full overflow-hidden"
+                            class="w-full h-1 bg-white/10 rounded-full overflow-hidden"
                           >
                             <div
                               class="h-full bg-exo-light-gray/40 rounded-full"
@@ -558,65 +558,21 @@
                               ).toFixed(1)}%"
                             ></div>
                           </div>
-                          {#if row.shardMetadata}
-                            <button
-                              type="button"
-                              class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
-                              onclick={() =>
-                                startDownload(col.nodeId, row.shardMetadata!)}
-                              title="Resume download on this node"
-                            >
-                              <svg
-                                class="w-5 h-5"
-                                viewBox="0 0 20 20"
-                                fill="none"
-                                stroke="currentColor"
-                                stroke-width="2"
-                              >
-                                <path
-                                  d="M10 3v10m0 0l-3-3m3 3l3-3M3 17h14"
-                                  stroke-linecap="round"
-                                  stroke-linejoin="round"
-                                ></path>
-                              </svg>
-                            </button>
-                          {:else}
-                            <span class="text-white/50 text-[10px]">paused</span
-                            >
-                          {/if}
-                        {:else if row.shardMetadata}
-                          <button
-                            type="button"
-                            class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
-                            onclick={() =>
-                              startDownload(col.nodeId, row.shardMetadata!)}
-                            title="Start download on this node"
+                          <span class="text-exo-light-gray/40 text-[9px]"
+                            >paused</span
                           >
-                            <svg
-                              class="w-6 h-6"
-                              viewBox="0 0 20 20"
-                              fill="none"
-                              stroke="currentColor"
-                              stroke-width="2"
-                            >
-                              <path
-                                d="M10 3v10m0 0l-3-3m3 3l3-3M3 17h14"
-                                stroke-linecap="round"
-                                stroke-linejoin="round"
-                              ></path>
-                            </svg>
-                          </button>
                         {:else}
-                          <span class="text-white/40 text-sm">...</span>
+                          <span class="text-exo-light-gray/50 text-sm">...</span
+                          >
                         {/if}
                       </div>
                     {:else if cell.kind === "failed"}
                       <div
-                        class="flex flex-col items-center gap-1"
+                        class="flex flex-col items-center gap-0.5"
                         title="Download failed"
                       >
                         <svg
-                          class="w-7 h-7 text-red-400"
+                          class="w-5 h-5 text-red-400"
                           viewBox="0 0 20 20"
                           fill="currentColor"
                         >
@@ -629,13 +585,13 @@
                         {#if row.shardMetadata}
                           <button
                             type="button"
-                            class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
+                            class="text-exo-light-gray/40 hover:text-exo-yellow transition-colors"
                             onclick={() =>
                               startDownload(col.nodeId, row.shardMetadata!)}
                             title="Retry download on this node"
                           >
                             <svg
-                              class="w-5 h-5"
+                              class="w-3.5 h-3.5"
                               viewBox="0 0 20 20"
                               fill="none"
                               stroke="currentColor"
@@ -661,13 +617,13 @@
                         {#if row.shardMetadata}
                           <button
                             type="button"
-                            class="text-white/50 hover:text-exo-yellow transition-colors mt-0.5 opacity-0 group-hover:opacity-100 cursor-pointer"
+                            class="text-exo-light-gray/30 hover:text-exo-yellow transition-colors mt-0.5 opacity-0 group-hover:opacity-100"
                             onclick={() =>
                               startDownload(col.nodeId, row.shardMetadata!)}
                             title="Download to this node"
                           >
                             <svg
-                              class="w-5 h-5"
+                              class="w-3.5 h-3.5"
                               viewBox="0 0 20 20"
                               fill="none"
                               stroke="currentColor"

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

@@ -3,7 +3,6 @@
   import {
     listTraces,
     getTraceRawUrl,
-    deleteTraces,
     type TraceListItem,
   } from "$lib/stores/app.svelte";
   import HeaderNav from "$lib/components/HeaderNav.svelte";
@@ -11,51 +10,6 @@
   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";
@@ -155,16 +109,6 @@
         </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"
@@ -199,41 +143,14 @@
       </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
-            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);
-              }
-            }}
+            class="rounded border border-exo-medium-gray/30 bg-exo-black/30 p-4 flex items-center justify-between gap-4"
           >
             <div class="min-w-0 flex-1">
               <a
                 href="#/traces/{trace.taskId}"
-                class="text-sm font-mono transition-colors truncate block {isSelected
-                  ? 'text-exo-yellow'
-                  : 'text-white hover:text-exo-yellow'}"
-                onclick={(e) => e.stopPropagation()}
+                class="text-sm font-mono text-white hover:text-exo-yellow transition-colors truncate block"
               >
                 {trace.taskId}
               </a>
@@ -243,11 +160,7 @@
                 )}
               </div>
             </div>
-            <!-- svelte-ignore a11y_click_events_have_key_events -->
-            <div
-              class="flex items-center gap-2 shrink-0"
-              onclick={(e) => e.stopPropagation()}
-            >
+            <div class="flex items-center gap-2 shrink-0">
               <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,27 +183,6 @@ 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
@@ -287,7 +266,6 @@ 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,7 +108,6 @@
                 package = pkgsSwift.swiftPackages.swift-format;
               };
               shfmt.enable = true;
-              taplo.enable = true;
             };
           };
 

nix/mlx.nix

@@ -41,7 +41,7 @@ let
 
   mlx = stdenv.mkDerivation rec {
     pname = "mlx";
-    version = let v = "0.30.7.dev20260225+257d5692"; in
+    version = let v = "0.30.7.dev20260220+13998a05"; in
       assert v == uvLockMlxVersion || throw "MLX version mismatch: nix/mlx.nix has ${v} but uv.lock has ${uvLockMlxVersion}. Update both the version and hash in nix/mlx.nix.";
       v;
     pyproject = true;
@@ -49,8 +49,8 @@ let
     src = fetchFromGitHub {
       owner = "rltakashige";
       repo = "mlx-jaccl-fix-small-recv";
-      rev = "257d5692fc7af6bba3b8afaeb63c549b7d1e43d5";
-      hash = "sha256-GosFIWxIB48Egb1MqJrR3xhsUsQeWdRk5rV93USY6wQ=";
+      rev = "13998a054715edcdc93618fb1496c79c7c25ff7c";
+      hash = "sha256-fAqA3hFwNBx7FcoGnhQsIFpAIRbC2EerACm4Fvne0Cc=";
     };
 
     patches = [

pyproject.toml

@@ -1,35 +1,35 @@
 [project]
 name = "exo"
-version = "0.3.68"
+version = "0.3.0"
 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",
-  "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==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",
 ]
 
 [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,12 +57,15 @@ 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/ml-explore/mlx-lm", rev = "834fac934c4e04de9b3d723e2b9287a2c60cfd4a" }
+#mlx-lm = { git = "https://github.com/davidmcc73/mlx-lm", branch = "stable" }
 # 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 }
@@ -92,15 +95,7 @@ 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]]
@@ -114,18 +109,17 @@ 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"]
@@ -133,7 +127,13 @@ 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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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,24 +26,20 @@ 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,6 +2,7 @@
 # ruff: noqa: E501, F401
 
 import builtins
+import enum
 import typing
 
 @typing.final
@@ -10,6 +11,29 @@ 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"""
@@ -34,15 +58,21 @@ 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.
@@ -61,7 +91,18 @@ class NetworkingHandle:
         
         If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
         """
-    async def recv(self) -> PyFromSwarm: ...
+    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 NoPeersSubscribedToTopicError(builtins.Exception):
@@ -69,26 +110,11 @@ class NoPeersSubscribedToTopicError(builtins.Exception):
     def __repr__(self) -> builtins.str: ...
     def __str__(self) -> builtins.str: ...
 
-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: ...
-    
-    ...
+@typing.final
+class ConnectionUpdateType(enum.Enum):
+    r"""
+    Connection or disconnection event discriminant type.
+    """
+    Connected = ...
+    Disconnected = ...
 

rust/exo_pyo3_bindings/pyproject.toml

@@ -4,18 +4,21 @@ build-backend = "maturin"
 
 [project]
 name = "exo_pyo3_bindings"
-version = "0.2.0"
+version = "0.1.0"
 description = "Add your description here"
 readme = "README.md"
 authors = [
-  { name = "Andrei Cravtov", email = "the.andrei.cravtov@gmail.com" },
-  { name = "Evan Quiney", email = "evanev7@gmail.com" },
+    { name = "Andrei Cravtov", email = "the.andrei.cravtov@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,9 +155,6 @@ 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,24 +1,26 @@
-use std::pin::Pin;
-use std::sync::Arc;
+#![allow(
+    clippy::multiple_inherent_impl,
+    clippy::unnecessary_wraps,
+    clippy::unused_self,
+    clippy::needless_pass_by_value
+)]
 
 use crate::r#const::MPSC_CHANNEL_SIZE;
 use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
-use crate::ext::{ResultExt as _, TokioMpscSenderExt as _};
+use crate::ext::{ResultExt as _, TokioMpscReceiverExt as _, TokioMpscSenderExt as _};
 use crate::ident::PyKeypair;
-use crate::networking::exception::{
-    PyAllQueuesFullError, PyMessageTooLargeError, PyNoPeersSubscribedToTopicError,
-};
 use crate::pyclass;
-use futures_lite::{Stream, StreamExt as _};
-use libp2p::gossipsub::PublishError;
-use networking::swarm::{FromSwarm, ToSwarm, create_swarm};
-use pyo3::exceptions::PyRuntimeError;
+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 pyo3::prelude::{PyModule, PyModuleMethods as _};
 use pyo3::types::PyBytes;
-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 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 tokio::sync::{Mutex, mpsc, oneshot};
 
 mod exception {
@@ -96,78 +98,237 @@ mod exception {
             Self::MSG.to_string()
         }
     }
+}
 
-    #[gen_stub_pyclass]
-    #[pyclass(frozen, extends=PyException, name="MessageTooLargeError")]
-    pub struct PyMessageTooLargeError {}
+/// 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,
+}
 
-    impl PyMessageTooLargeError {
-        const MSG: &'static str = "Gossipsub message exceeds max_transmit_size. Reduce prompt length or increase the limit.";
+#[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,
 
-        pub(crate) fn new_err() -> PyErr {
-            PyErr::new::<Self, _>(())
+    /// 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 {
+                            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:?}");
+                    }
+                }
+            }
         }
     }
 
-    #[gen_stub_pymethods]
-    #[pymethods]
-    impl PyMessageTooLargeError {
-        #[new]
-        #[pyo3(signature = (*args))]
-        #[allow(unused_variables)]
-        pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self {
-            Self {}
-        }
-
-        fn __repr__(&self) -> String {
-            format!("MessageTooLargeError(\"{}\")", Self::MSG)
-        }
-
-        fn __str__(&self) -> String {
-            Self::MSG.to_string()
-        }
-    }
+    log::info!("RUST: networking task stopped");
 }
 
 #[gen_stub_pyclass]
 #[pyclass(name = "NetworkingHandle")]
+#[derive(Debug)]
 struct PyNetworkingHandle {
     // channels
-    pub to_swarm: mpsc::Sender<ToSwarm>,
-    pub swarm: Arc<Mutex<Pin<Box<dyn Stream<Item = FromSwarm> + Send>>>>,
+    to_task_tx: Option<mpsc::Sender<ToTask>>,
+    connection_update_rx: Mutex<mpsc::Receiver<PyConnectionUpdate>>,
+    gossipsub_message_rx: Mutex<mpsc::Receiver<(String, Vec<u8>)>>,
 }
 
-#[gen_stub_pyclass_complex_enum]
-#[pyclass]
-enum PyFromSwarm {
-    Connection {
-        peer_id: String,
-        connected: bool,
-    },
-    Message {
-        origin: String,
-        topic: String,
-        data: Py<PyBytes>,
-    },
+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
+    }
 }
-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(),
-            },
+
+#[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),
         }
     }
+
+    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]
@@ -181,36 +342,97 @@ impl PyNetworkingHandle {
 
     #[new]
     fn py_new(identity: Bound<'_, PyKeypair>) -> PyResult<Self> {
+        use pyo3_async_runtimes::tokio::get_runtime;
+
         // create communication channels
-        let (to_swarm, from_client) = mpsc::channel(MPSC_CHANNEL_SIZE);
+        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);
 
         // get identity
         let identity = identity.borrow().0.clone();
 
         // create networking swarm (within tokio context!! or it crashes)
-        let _guard = pyo3_async_runtimes::tokio::get_runtime().enter();
-        let swarm = create_swarm(identity, from_client).pyerr()?.into_stream();
+        let swarm = get_runtime()
+            .block_on(async { create_swarm(identity) })
+            .pyerr()?;
 
-        Ok(Self {
-            swarm: Arc::new(Mutex::new(swarm)),
-            to_swarm,
-        })
+        // 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,
+        ))
     }
 
     #[gen_stub(skip)]
-    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)
-        })
+    const fn __traverse__(&self, _visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
+        Ok(()) // This is needed purely so `__clear__` can work
     }
 
+    #[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.
@@ -220,10 +442,10 @@ impl PyNetworkingHandle {
         let (tx, rx) = oneshot::channel();
 
         // send off request to subscribe
-        self.to_swarm
-            .send_py(ToSwarm::Subscribe {
+        self.to_task_tx()
+            .send_py(ToTask::GossipsubSubscribe {
                 topic,
-                result_sender: tx,
+                result_tx: tx,
             })
             .allow_threads_py() // allow-threads-aware async call
             .await?;
@@ -232,7 +454,6 @@ impl PyNetworkingHandle {
         rx.allow_threads_py() // allow-threads-aware async call
             .await
             .map_err(|_| PyErr::receiver_channel_closed())?
-            .pyerr()
     }
 
     /// Unsubscribes from a `GossipSub` topic.
@@ -242,10 +463,10 @@ impl PyNetworkingHandle {
         let (tx, rx) = oneshot::channel();
 
         // send off request to unsubscribe
-        self.to_swarm
-            .send_py(ToSwarm::Unsubscribe {
+        self.to_task_tx()
+            .send_py(ToTask::GossipsubUnsubscribe {
                 topic,
-                result_sender: tx,
+                result_tx: tx,
             })
             .allow_threads_py() // allow-threads-aware async call
             .await?;
@@ -264,11 +485,11 @@ impl PyNetworkingHandle {
 
         // send off request to subscribe
         let data = Python::attach(|py| Vec::from(data.as_bytes(py)));
-        self.to_swarm
-            .send_py(ToSwarm::Publish {
+        self.to_task_tx()
+            .send_py(ToTask::GossipsubPublish {
                 topic,
                 data,
-                result_sender: tx,
+                result_tx: tx,
             })
             .allow_threads_py() // allow-threads-aware async call
             .await?;
@@ -277,35 +498,74 @@ impl PyNetworkingHandle {
         let _ = rx
             .allow_threads_py() // allow-threads-aware async call
             .await
-            .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()),
-            })?;
+            .map_err(|_| PyErr::receiver_channel_closed())??;
         Ok(())
     }
-}
 
-pyo3_stub_gen::inventory::submit! {
-    gen_methods_from_python! {
-        r#"
-            class PyNetworkingHandle:
-                async def recv() -> PyFromSwarm: ...
-        "#
+    // ---- 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()))
     }
+
+    /// 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<()> {
     m.add_class::<exception::PyNoPeersSubscribedToTopicError>()?;
     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,20 +1,16 @@
 import asyncio
 
 import pytest
-from exo_pyo3_bindings import (
-    Keypair,
-    NetworkingHandle,
-    NoPeersSubscribedToTopicError,
-    PyFromSwarm,
-)
+from exo_pyo3_bindings import Keypair, NetworkingHandle, NoPeersSubscribedToTopicError
 
 
 @pytest.mark.asyncio
 async def test_sleep_on_multiple_items() -> None:
     print("PYTHON: starting handle")
-    h = NetworkingHandle(Keypair.generate())
+    h = NetworkingHandle(Keypair.generate_ed25519())
 
-    rt = asyncio.create_task(_await_recv(h))
+    ct = asyncio.create_task(_await_cons(h))
+    mt = asyncio.create_task(_await_msg(h))
 
     # sleep for 4 ticks
     for i in range(4):
@@ -26,11 +22,13 @@ async def test_sleep_on_multiple_items() -> None:
             print("caught it", e)
 
 
-async def _await_recv(h: NetworkingHandle):
+async def _await_cons(h: NetworkingHandle):
     while True:
-        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}")
+        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}")

rust/networking/Cargo.toml

@@ -21,17 +21,13 @@ extend = { workspace = true }
 delegate = { workspace = true }
 
 # async
-async-stream = { workspace = true }
+tokio = { workspace = true, features = ["full"] }
 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
@@ -39,4 +35,3 @@ log = { workspace = true }
 
 # networking
 libp2p = { workspace = true, features = ["full"] }
-pin-project = "1.1.10"

rust/networking/examples/chatroom.rs

@@ -1,9 +1,7 @@
 use futures_lite::StreamExt;
-use libp2p::identity;
-use networking::swarm;
-use networking::swarm::{FromSwarm, ToSwarm};
-use tokio::sync::{mpsc, oneshot};
-use tokio::{io, io::AsyncBufReadExt as _};
+use libp2p::{gossipsub, identity, swarm::SwarmEvent};
+use networking::{discovery, swarm};
+use tokio::{io, io::AsyncBufReadExt as _, select};
 use tracing_subscriber::EnvFilter;
 use tracing_subscriber::filter::LevelFilter;
 
@@ -13,69 +11,64 @@ 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(), from_client)
-        .expect("Swarm creation failed")
-        .into_stream();
+    let mut swarm =
+        swarm::create_swarm(identity::Keypair::generate_ed25519()).expect("Swarm creation failed");
 
     // Create a Gossipsub topic & subscribe
-    let (tx, rx) = oneshot::channel();
-    _ = to_swarm
-        .send(ToSwarm::Subscribe {
-            topic: "test-net".to_string(),
-            result_sender: tx,
-        })
-        .await
-        .expect("should send");
+    let topic = gossipsub::IdentTopic::new("test-net");
+    swarm
+        .behaviour_mut()
+        .gossipsub
+        .subscribe(&topic)
+        .expect("Subscribing to topic failed");
 
     // 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 {
-        // on gossipsub outgoing
-        match swarm.next().await {
-            // on gossipsub incoming
-            Some(FromSwarm::Discovered { peer_id }) => {
-                println!("\n\nconnected to {peer_id}\n\n")
+        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:?}");
+                }
             }
-            Some(FromSwarm::Expired { peer_id }) => {
-                println!("\n\ndisconnected from {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::Message { from, topic, data }) => {
-                println!("{topic}/{from}:\n{}", String::from_utf8_lossy(&data))
-            }
-            None => {}
         }
     }
 }

rust/networking/src/swarm.rs

@@ -1,11 +1,9 @@
-use std::pin::Pin;
-
+use crate::alias;
 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};
+use libp2p::{SwarmBuilder, identity};
+
+pub type Swarm = libp2p::Swarm<Behaviour>;
 
 /// The current version of the network: this prevents devices running different versions of the
 /// software from interacting with each other.
@@ -17,136 +15,8 @@ use tokio::sync::{mpsc, oneshot};
 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,
-    from_client: mpsc::Receiver<ToSwarm>,
-) -> alias::AnyResult<Swarm> {
+pub fn create_swarm(keypair: identity::Keypair) -> alias::AnyResult<Swarm> {
     let mut swarm = SwarmBuilder::with_existing_identity(keypair)
         .with_tokio()
         .with_other_transport(tcp_transport)?
@@ -155,7 +25,7 @@ pub fn create_swarm(
 
     // Listen on all interfaces and whatever port the OS assigns
     swarm.listen_on("/ip4/0.0.0.0/tcp/0".parse()?)?;
-    Ok(Swarm { swarm, from_client })
+    Ok(swarm)
 }
 
 mod transport {

src/exo/download/coordinator.py

@@ -1,8 +1,11 @@
 import asyncio
+import socket
 from dataclasses import dataclass, field
+from random import random
 
 import anyio
 from anyio import current_time
+from anyio.abc import TaskGroup
 from loguru import logger
 
 from exo.download.download_utils import (
@@ -20,9 +23,13 @@ from exo.shared.types.commands import (
     ForwarderDownloadCommand,
     StartDownload,
 )
-from exo.shared.types.common import NodeId
+from exo.shared.types.common import NodeId, SessionId, SystemId
 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 (
@@ -33,28 +40,41 @@ from exo.shared.types.worker.downloads import (
     DownloadProgress,
 )
 from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
-from exo.utils.channels import Receiver, Sender
-from exo.utils.task_group import TaskGroup
+from exo.utils.channels import Receiver, Sender, channel
 
 
 @dataclass
 class DownloadCoordinator:
     node_id: NodeId
+    session_id: SessionId
     shard_downloader: ShardDownloader
     download_command_receiver: Receiver[ForwarderDownloadCommand]
-    event_sender: Sender[Event]
+    local_event_sender: Sender[LocalForwarderEvent]
+
+    # ack stuff
+    _global_event_receiver: Receiver[GlobalForwarderEvent]
+    _out_for_delivery: dict[EventId, LocalForwarderEvent] = field(default_factory=dict)
+
     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)
 
-    _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
+    # 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=anyio.create_task_group)
 
     # 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]()
+        if self.offline:
+            self.shard_downloader.set_internet_connection(False)
         self.shard_downloader.on_progress(self._download_progress_callback)
 
     def _model_dir(self, model_id: ModelId) -> str:
@@ -103,16 +123,65 @@ class DownloadCoordinator:
         logger.info(
             f"Starting DownloadCoordinator{' (offline mode)' if self.offline else ''}"
         )
+        if not self.offline:
+            self._test_internet_connection()
         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)
+                if not self.offline:
+                    tg.start_soon(self._check_internet_connection)
         finally:
             for task in self.active_downloads.values():
                 task.cancel()
 
+    def _test_internet_connection(self) -> None:
+        # Try multiple endpoints since some ISPs/networks block specific IPs
+        for host in ("1.1.1.1", "8.8.8.8", "1.0.0.1"):
+            try:
+                socket.create_connection((host, 443), timeout=3).close()
+                self.shard_downloader.set_internet_connection(True)
+                logger.debug(f"Internet connectivity: True (via {host})")
+                return
+            except OSError:
+                continue
+        self.shard_downloader.set_internet_connection(False)
+        logger.debug("Internet connectivity: False")
+
+    async def _check_internet_connection(self) -> None:
+        first_connection = True
+        while True:
+            await asyncio.sleep(10)
+
+            # Assume that internet connection is set to False on 443 errors.
+            if self.shard_downloader.internet_connection:
+                continue
+
+            self._test_internet_connection()
+
+            if first_connection and self.shard_downloader.internet_connection:
+                first_connection = False
+                self._tg.start_soon(self._emit_existing_download_progress)
+
     def shutdown(self) -> None:
-        self._tg.cancel_tasks()
+        self._tg.cancel_scope.cancel()
+
+    # 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:
@@ -286,6 +355,23 @@ 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

@@ -110,20 +110,54 @@ def map_repo_download_progress_to_download_progress_data(
     )
 
 
+def _resolve_hf_hub_model(search_dir: Path, normalized: str) -> Path | None:
+    """Try to find a model in HuggingFace Hub cache format.
+
+    HF Hub stores models as ``models--<org>--<name>/snapshots/<commit>/``
+    with symlinks to ``../../blobs/``.  The active commit is read from
+    ``refs/main``.
+    """
+    hf_model_dir = search_dir / f"models--{normalized}"
+    if not hf_model_dir.is_dir():
+        return None
+    # Resolve ref → snapshot
+    ref_file = hf_model_dir / "refs" / "main"
+    if ref_file.is_file():
+        commit_hash = ref_file.read_text().strip()
+        snapshot = hf_model_dir / "snapshots" / commit_hash
+        if snapshot.is_dir():
+            return snapshot
+    # Fallback: use latest snapshot by mtime
+    snapshots_dir = hf_model_dir / "snapshots"
+    if snapshots_dir.is_dir():
+        snapshots = sorted(
+            snapshots_dir.iterdir(), key=lambda p: p.stat().st_mtime, reverse=True
+        )
+        if snapshots:
+            return snapshots[0]
+    return None
+
+
 def resolve_model_in_path(model_id: ModelId) -> Path | None:
     """Search EXO_MODELS_PATH directories for a pre-existing model.
 
-    Checks each directory for the normalized name (org--model).  A candidate
-    is only returned if ``is_model_directory_complete`` confirms all weight
-    files are present.
+    Checks each directory for the normalized name (org--model) and the
+    HuggingFace Hub cache format (models--org--model/snapshots/<ref>/).
+    A candidate is only returned if ``is_model_directory_complete``
+    confirms all weight files are present.
     """
     if EXO_MODELS_PATH is None:
         return None
     normalized = model_id.normalize()
     for search_dir in EXO_MODELS_PATH:
+        # Try direct format: <dir>/<org--name>/
         candidate = search_dir / normalized
         if candidate.is_dir() and is_model_directory_complete(candidate):
             return candidate
+        # Try HF Hub cache format: <dir>/models--<org--name>/snapshots/<ref>/
+        hf_candidate = _resolve_hf_hub_model(search_dir, normalized)
+        if hf_candidate is not None and is_model_directory_complete(hf_candidate):
+            return hf_candidate
     return None
 
 
@@ -314,13 +348,9 @@ 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 cached file list for {model_id} - using local file list"
+                f"No internet and 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())
@@ -827,7 +857,6 @@ async def download_shard(
 
     for file in filtered_file_list:
         downloaded_bytes = await get_downloaded_size(target_dir / file.path)
-        final_file_exists = await aios.path.exists(target_dir / file.path)
         file_progress[file.path] = RepoFileDownloadProgress(
             repo_id=shard.model_card.model_id,
             repo_revision=revision,
@@ -837,9 +866,7 @@ async def download_shard(
             total=Memory.from_bytes(file.size or 0),
             speed=0,
             eta=timedelta(0),
-            status="complete"
-            if final_file_exists and downloaded_bytes == file.size
-            else "not_started",
+            status="complete" if downloaded_bytes == file.size else "not_started",
             start_time=time.time(),
         )
 

src/exo/download/impl_shard_downloader.py

@@ -15,13 +15,9 @@ from exo.shared.types.worker.shards import (
 )
 
 
-def exo_shard_downloader(
-    max_parallel_downloads: int = 8, offline: bool = False
-) -> ShardDownloader:
+def exo_shard_downloader(max_parallel_downloads: int = 8) -> ShardDownloader:
     return SingletonShardDownloader(
-        CachedShardDownloader(
-            ResumableShardDownloader(max_parallel_downloads, offline=offline)
-        )
+        CachedShardDownloader(ResumableShardDownloader(max_parallel_downloads))
     )
 
 
@@ -54,6 +50,10 @@ class SingletonShardDownloader(ShardDownloader):
         self.shard_downloader = shard_downloader
         self.active_downloads: dict[ShardMetadata, asyncio.Task[Path]] = {}
 
+    def set_internet_connection(self, value: bool) -> None:
+        self.internet_connection = value
+        self.shard_downloader.set_internet_connection(value)
+
     def on_progress(
         self,
         callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
@@ -90,6 +90,10 @@ class CachedShardDownloader(ShardDownloader):
         self.shard_downloader = shard_downloader
         self.cache: dict[tuple[str, ShardMetadata], Path] = {}
 
+    def set_internet_connection(self, value: bool) -> None:
+        self.internet_connection = value
+        self.shard_downloader.set_internet_connection(value)
+
     def on_progress(
         self,
         callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
@@ -119,9 +123,8 @@ class CachedShardDownloader(ShardDownloader):
 
 
 class ResumableShardDownloader(ShardDownloader):
-    def __init__(self, max_parallel_downloads: int = 8, offline: bool = False):
+    def __init__(self, max_parallel_downloads: int = 8):
         self.max_parallel_downloads = max_parallel_downloads
-        self.offline = offline
         self.on_progress_callbacks: list[
             Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]]
         ] = []
@@ -148,7 +151,8 @@ class ResumableShardDownloader(ShardDownloader):
             self.on_progress_wrapper,
             max_parallel_downloads=self.max_parallel_downloads,
             allow_patterns=allow_patterns,
-            skip_internet=self.offline,
+            skip_internet=not self.internet_connection,
+            on_connection_lost=lambda: self.set_internet_connection(False),
         )
         return target_dir
 
@@ -164,7 +168,8 @@ class ResumableShardDownloader(ShardDownloader):
                 shard,
                 self.on_progress_wrapper,
                 skip_download=True,
-                skip_internet=self.offline,
+                skip_internet=not self.internet_connection,
+                on_connection_lost=lambda: self.set_internet_connection(False),
             )
 
         semaphore = asyncio.Semaphore(self.max_parallel_downloads)
@@ -193,6 +198,7 @@ class ResumableShardDownloader(ShardDownloader):
             shard,
             self.on_progress_wrapper,
             skip_download=True,
-            skip_internet=self.offline,
+            skip_internet=not self.internet_connection,
+            on_connection_lost=lambda: self.set_internet_connection(False),
         )
         return progress

src/exo/download/shard_downloader.py

@@ -16,6 +16,11 @@ from exo.shared.types.worker.shards import (
 
 # TODO: the PipelineShardMetadata getting reinstantiated is a bit messy. Should this be a classmethod?
 class ShardDownloader(ABC):
+    internet_connection: bool = False
+
+    def set_internet_connection(self, value: bool) -> None:
+        self.internet_connection = value
+
     @abstractmethod
     async def ensure_shard(
         self, shard: ShardMetadata, config_only: bool = False

src/exo/download/tests/test_coordinator_ack.py

@@ -0,0 +1,98 @@
+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

@@ -7,6 +7,7 @@ from dataclasses import dataclass, field
 from typing import Self
 
 import anyio
+from anyio.abc import TaskGroup
 from loguru import logger
 from pydantic import PositiveInt
 
@@ -15,7 +16,6 @@ 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
@@ -23,14 +23,12 @@ from exo.shared.logging import logger_cleanup, logger_setup
 from exo.shared.types.common import NodeId, SessionId
 from exo.utils.channels import Receiver, channel
 from exo.utils.pydantic_ext import CamelCaseModel
-from exo.utils.task_group import TaskGroup
 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.
@@ -40,7 +38,7 @@ class Node:
 
     node_id: NodeId
     offline: bool
-    _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
+    _tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)
 
     @classmethod
     async def create(cls, args: "Args") -> Self:
@@ -54,12 +52,6 @@ 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}")
 
@@ -67,10 +59,13 @@ class Node:
         if not args.no_downloads:
             download_coordinator = DownloadCoordinator(
                 node_id,
-                exo_shard_downloader(offline=args.offline),
-                event_sender=event_router.sender(),
+                session_id,
+                exo_shard_downloader(),
                 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
@@ -78,8 +73,9 @@ class Node:
         if args.spawn_api:
             api = API(
                 node_id,
+                session_id,
                 port=args.api_port,
-                event_receiver=event_router.receiver(),
+                global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
                 command_sender=router.sender(topics.COMMANDS),
                 download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
                 election_receiver=router.receiver(topics.ELECTION_MESSAGES),
@@ -90,8 +86,9 @@ class Node:
         if not args.no_worker:
             worker = Worker(
                 node_id,
-                event_receiver=event_router.receiver(),
-                event_sender=event_router.sender(),
+                session_id,
+                global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
+                local_event_sender=router.sender(topics.LOCAL_EVENTS),
                 command_sender=router.sender(topics.COMMANDS),
                 download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
             )
@@ -102,7 +99,6 @@ 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),
@@ -125,7 +121,6 @@ class Node:
 
         return cls(
             router,
-            event_router,
             download_coordinator,
             worker,
             election,
@@ -141,7 +136,6 @@ 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)
@@ -155,11 +149,11 @@ class Node:
 
     def shutdown(self):
         # if this is our second call to shutdown, just sys.exit
-        if self._tg.cancel_called():
+        if self._tg.cancel_scope.cancel_called:
             import sys
 
             sys.exit(1)
-        self._tg.cancel_tasks()
+        self._tg.cancel_scope.cancel()
 
     async def _elect_loop(self):
         with self.election_result_receiver as results:
@@ -176,17 +170,6 @@ 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
@@ -200,7 +183,6 @@ 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),
@@ -223,16 +205,22 @@ 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,
-                            exo_shard_downloader(offline=self.offline),
-                            event_sender=self.event_router.sender(),
+                            result.session_id,
+                            exo_shard_downloader(),
                             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:
@@ -240,8 +228,11 @@ class Node:
                         # TODO: add profiling etc to resource monitor
                         self.worker = Worker(
                             self.node_id,
-                            event_receiver=self.event_router.receiver(),
-                            event_sender=self.event_router.sender(),
+                            result.session_id,
+                            global_event_receiver=self.router.receiver(
+                                topics.GLOBAL_EVENTS
+                            ),
+                            local_event_sender=self.router.sender(topics.LOCAL_EVENTS),
                             command_sender=self.router.sender(topics.COMMANDS),
                             download_command_sender=self.router.sender(
                                 topics.DOWNLOAD_COMMANDS
@@ -249,7 +240,7 @@ class Node:
                         )
                         self._tg.start_soon(self.worker.run)
                     if self.api:
-                        self.api.reset(result.won_clock, self.event_router.receiver())
+                        self.api.reset(result.session_id, result.won_clock)
                 else:
                     if self.api:
                         self.api.unpause(result.won_clock)
@@ -261,7 +252,7 @@ def main():
     target = min(max(soft, 65535), hard)
     resource.setrlimit(resource.RLIMIT_NOFILE, (target, hard))
 
-    mp.set_start_method("spawn", force=True)
+    mp.set_start_method("spawn")
     # TODO: Refactor the current verbosity system
     logger_setup(EXO_LOG, args.verbosity)
     logger.info("Starting EXO")
@@ -279,16 +270,9 @@ def main():
         logger.info("FAST_SYNCH forced OFF")
 
     node = anyio.run(Node.create, args)
-    try:
-        anyio.run(node.run)
-    except BaseException as exception:
-        logger.opt(exception=exception).critical(
-            "EXO terminated due to unhandled exception"
-        )
-        raise
-    finally:
-        logger.info("EXO Shutdown complete")
-        logger_cleanup()
+    anyio.run(node.run)
+    logger.info("EXO Shutdown complete")
+    logger_cleanup()
 
 
 class Args(CamelCaseModel):
@@ -299,7 +283,7 @@ class Args(CamelCaseModel):
     tb_only: bool = False
     no_worker: bool = False
     no_downloads: bool = False
-    offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true"
+    offline: bool = False
     fast_synch: bool | None = None  # None = auto, True = force on, False = force off
 
     @classmethod
@@ -350,7 +334,6 @@ class Args(CamelCaseModel):
         parser.add_argument(
             "--offline",
             action="store_true",
-            default=os.getenv("EXO_OFFLINE", "false").lower() == "true",
             help="Run in offline/air-gapped mode: skip internet checks, use only pre-staged local models",
         )
         fast_synch_group = parser.add_mutually_exclusive_group()

src/exo/master/api.py

@@ -11,7 +11,8 @@ from typing import Annotated, Literal, cast
 from uuid import uuid4
 
 import anyio
-from anyio import BrokenResourceError
+from anyio import BrokenResourceError, create_task_group
+from anyio.abc import TaskGroup
 from fastapi import FastAPI, File, Form, HTTPException, Query, Request, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
@@ -73,7 +74,6 @@ from exo.shared.types.api import (
     BenchChatCompletionResponse,
     BenchImageGenerationResponse,
     BenchImageGenerationTaskParams,
-    CancelCommandResponse,
     ChatCompletionChoice,
     ChatCompletionMessage,
     ChatCompletionRequest,
@@ -82,8 +82,6 @@ from exo.shared.types.api import (
     CreateInstanceResponse,
     DeleteDownloadResponse,
     DeleteInstanceResponse,
-    DeleteTracesRequest,
-    DeleteTracesResponse,
     ErrorInfo,
     ErrorResponse,
     FinishReason,
@@ -143,10 +141,11 @@ from exo.shared.types.commands import (
     TaskFinished,
     TextGeneration,
 )
-from exo.shared.types.common import CommandId, Id, NodeId, SystemId
+from exo.shared.types.common import CommandId, Id, NodeId, SessionId, SystemId
 from exo.shared.types.events import (
     ChunkGenerated,
     Event,
+    GlobalForwarderEvent,
     IndexedEvent,
     TracesMerged,
 )
@@ -174,7 +173,7 @@ 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.task_group import TaskGroup
+from exo.utils.event_buffer import OrderedBuffer
 
 _API_EVENT_LOG_DIR = EXO_EVENT_LOG_DIR / "api"
 ONBOARDING_COMPLETE_FILE = EXO_CACHE_HOME / "onboarding_complete"
@@ -197,9 +196,10 @@ class API:
     def __init__(
         self,
         node_id: NodeId,
+        session_id: SessionId,
         *,
         port: int,
-        event_receiver: Receiver[IndexedEvent],
+        global_event_receiver: Receiver[GlobalForwarderEvent],
         command_sender: Sender[ForwarderCommand],
         download_command_sender: Sender[ForwarderDownloadCommand],
         # This lets us pause the API if an election is running
@@ -210,9 +210,11 @@ class API:
         self._system_id = SystemId()
         self.command_sender = command_sender
         self.download_command_sender = download_command_sender
-        self.event_receiver = event_receiver
+        self.global_event_receiver = global_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
 
@@ -250,20 +252,19 @@ class API:
             CommandId, Sender[ImageChunk | ErrorChunk]
         ] = {}
         self._image_store = ImageStore(EXO_IMAGE_CACHE_DIR)
-        self._tg: TaskGroup = TaskGroup()
+        self._tg: TaskGroup | None = None
 
-    def reset(self, result_clock: int, event_receiver: Receiver[IndexedEvent]):
+    def reset(self, new_session_id: SessionId, result_clock: int):
         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")
@@ -323,7 +324,6 @@ 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,7 +344,6 @@ 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)
@@ -569,25 +568,6 @@ 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[
@@ -1611,7 +1591,8 @@ class API:
         shutdown_ev = anyio.Event()
 
         try:
-            async with self._tg as tg:
+            async with create_task_group() as tg:
+                self._tg = tg
                 logger.info("Starting API")
                 tg.start_soon(self._apply_state)
                 tg.start_soon(self._pause_on_new_election)
@@ -1626,7 +1607,7 @@ class API:
         finally:
             self._event_log.close()
             self.command_sender.close()
-            self.event_receiver.close()
+            self.global_event_receiver.close()
 
     async def run_api(self, ev: anyio.Event):
         cfg = Config()
@@ -1643,31 +1624,38 @@ class API:
             )
 
     async def _apply_state(self):
-        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
+        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))
 
-                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 = [
@@ -1731,12 +1719,8 @@ class API:
         await self._send_download(command)
         return DeleteDownloadResponse(command_id=command.command_id)
 
-    @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
+    def _get_trace_path(self, task_id: str) -> Path:
+        return EXO_TRACING_CACHE_DIR / f"trace_{task_id}.json"
 
     async def list_traces(self) -> TraceListResponse:
         traces: list[TraceListItem] = []
@@ -1835,18 +1819,6 @@ 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

@@ -1,6 +1,7 @@
 from datetime import datetime, timedelta, timezone
 
 import anyio
+from anyio.abc import TaskGroup
 from loguru import logger
 
 from exo.master.event_log import DiskEventLog
@@ -60,9 +61,8 @@ from exo.shared.types.tasks import (
     TextGeneration as TextGenerationTask,
 )
 from exo.shared.types.worker.instances import InstanceId
-from exo.utils.channels import Receiver, Sender
+from exo.utils.channels import Receiver, Sender, channel
 from exo.utils.event_buffer import MultiSourceBuffer
-from exo.utils.task_group import TaskGroup
 
 
 class Master:
@@ -72,21 +72,25 @@ 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 = anyio.create_task_group()
         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
-        self.event_sender = event_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._system_id = SystemId()
         self._multi_buffer = MultiSourceBuffer[SystemId, Event]()
         self._event_log = DiskEventLog(EXO_EVENT_LOG_DIR / "master")
@@ -100,16 +104,19 @@ 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")
-        self._tg.cancel_tasks()
+        self._tg.cancel_scope.cancel()
 
     async def _command_processor(self) -> None:
         with self.command_receiver as commands:
@@ -402,6 +409,22 @@ 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/placement.py

@@ -112,11 +112,7 @@ def place_instance(
         cycle for cycle in smallest_cycles if topology.is_rdma_cycle(cycle)
     ]
 
-    if command.instance_meta == InstanceMeta.MlxJaccl:
-        if not smallest_rdma_cycles:
-            raise ValueError(
-                "Requested RDMA (MlxJaccl) but no RDMA-connected cycles available"
-            )
+    if command.instance_meta == InstanceMeta.MlxJaccl and smallest_rdma_cycles != []:
         smallest_cycles = smallest_rdma_cycles
 
     cycles_with_leaf_nodes: list[Cycle] = [
@@ -133,11 +129,6 @@ def place_instance(
         ),
     )
 
-    # Single-node: force Pipeline/Ring (Tensor and Jaccl require multi-node)
-    if len(selected_cycle) == 1:
-        command.instance_meta = InstanceMeta.MlxRing
-        command.sharding = Sharding.Pipeline
-
     shard_assignments = get_shard_assignments(
         command.model_card, selected_cycle, command.sharding, node_memory
     )
@@ -147,6 +138,9 @@ def place_instance(
     instance_id = InstanceId()
     target_instances = dict(deepcopy(current_instances))
 
+    if len(selected_cycle) == 1:
+        command.instance_meta = InstanceMeta.MlxRing
+
     match command.instance_meta:
         case InstanceMeta.MlxJaccl:
             # TODO(evan): shard assignments should contain information about ranks, this is ugly

src/exo/master/tests/test_cancel_command.py

@@ -1,77 +0,0 @@
-# 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,7 +17,6 @@ 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,
@@ -51,22 +50,6 @@ 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] = []
 
@@ -84,7 +67,6 @@ 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,
@@ -93,7 +75,6 @@ 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")
@@ -216,5 +197,4 @@ async def test_master():
             input=[InputMessage(role="user", content="Hello, how are you?")],
         )
 
-        ev_send.close()
         await master.shutdown()

src/exo/master/tests/test_placement.py

@@ -14,12 +14,10 @@ from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
 from exo.shared.topology import Topology
 from exo.shared.types.commands import PlaceInstance
 from exo.shared.types.common import CommandId, NodeId
-from exo.shared.types.events import InstanceCreated, InstanceDeleted, TaskStatusUpdated
+from exo.shared.types.events import InstanceCreated, InstanceDeleted
 from exo.shared.types.memory import Memory
 from exo.shared.types.multiaddr import Multiaddr
 from exo.shared.types.profiling import NetworkInterfaceInfo, NodeNetworkInfo
-from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 from exo.shared.types.topology import Connection, SocketConnection
 from exo.shared.types.worker.instances import (
     Instance,
@@ -458,117 +456,3 @@ def test_tensor_rdma_backend_connectivity_matrix(
         else:
             ip_part = coordinator.split(":")[0]
             assert len(ip_part.split(".")) == 4
-
-
-def _make_task(
-    instance_id: InstanceId,
-    status: TaskStatus = TaskStatus.Running,
-) -> TextGeneration:
-    return TextGeneration(
-        task_id=TaskId(),
-        task_status=status,
-        instance_id=instance_id,
-        command_id=CommandId(),
-        task_params=TextGenerationTaskParams(
-            model=ModelId("test-model"),
-            input=[InputMessage(role="user", content="hello")],
-        ),
-    )
-
-
-def test_get_transition_events_delete_instance_cancels_running_tasks(
-    instance: Instance,
-):
-    # arrange
-    instance_id = InstanceId()
-    current_instances: dict[InstanceId, Instance] = {instance_id: instance}
-    target_instances: dict[InstanceId, Instance] = {}
-    task = _make_task(instance_id, TaskStatus.Running)
-    tasks = {task.task_id: task}
-
-    # act
-    events = get_transition_events(current_instances, target_instances, tasks)
-
-    # assert – cancellation event should come before the deletion event
-    assert len(events) == 2
-    assert isinstance(events[0], TaskStatusUpdated)
-    assert events[0].task_id == task.task_id
-    assert events[0].task_status == TaskStatus.Cancelled
-    assert isinstance(events[1], InstanceDeleted)
-    assert events[1].instance_id == instance_id
-
-
-def test_get_transition_events_delete_instance_cancels_pending_tasks(
-    instance: Instance,
-):
-    # arrange
-    instance_id = InstanceId()
-    current_instances: dict[InstanceId, Instance] = {instance_id: instance}
-    target_instances: dict[InstanceId, Instance] = {}
-    task = _make_task(instance_id, TaskStatus.Pending)
-    tasks = {task.task_id: task}
-
-    # act
-    events = get_transition_events(current_instances, target_instances, tasks)
-
-    # assert
-    assert len(events) == 2
-    assert isinstance(events[0], TaskStatusUpdated)
-    assert events[0].task_id == task.task_id
-    assert events[0].task_status == TaskStatus.Cancelled
-    assert isinstance(events[1], InstanceDeleted)
-
-
-def test_get_transition_events_delete_instance_ignores_completed_tasks(
-    instance: Instance,
-):
-    # arrange
-    instance_id = InstanceId()
-    current_instances: dict[InstanceId, Instance] = {instance_id: instance}
-    target_instances: dict[InstanceId, Instance] = {}
-    tasks = {
-        t.task_id: t
-        for t in [
-            _make_task(instance_id, TaskStatus.Complete),
-            _make_task(instance_id, TaskStatus.Failed),
-            _make_task(instance_id, TaskStatus.TimedOut),
-            _make_task(instance_id, TaskStatus.Cancelled),
-        ]
-    }
-
-    # act
-    events = get_transition_events(current_instances, target_instances, tasks)
-
-    # assert – only the InstanceDeleted event, no cancellations
-    assert len(events) == 1
-    assert isinstance(events[0], InstanceDeleted)
-
-
-def test_get_transition_events_delete_instance_cancels_only_matching_tasks(
-    instance: Instance,
-):
-    # arrange
-    instance_id_a = InstanceId()
-    instance_id_b = InstanceId()
-    current_instances: dict[InstanceId, Instance] = {
-        instance_id_a: instance,
-        instance_id_b: instance,
-    }
-    # only delete instance A, keep instance B
-    target_instances: dict[InstanceId, Instance] = {instance_id_b: instance}
-
-    task_a = _make_task(instance_id_a, TaskStatus.Running)
-    task_b = _make_task(instance_id_b, TaskStatus.Running)
-    tasks = {task_a.task_id: task_a, task_b.task_id: task_b}
-
-    # act
-    events = get_transition_events(current_instances, target_instances, tasks)
-
-    # assert – only task_a should be cancelled
-    cancel_events = [e for e in events if isinstance(e, TaskStatusUpdated)]
-    delete_events = [e for e in events if isinstance(e, InstanceDeleted)]
-    assert len(cancel_events) == 1
-    assert cancel_events[0].task_id == task_a.task_id
-    assert cancel_events[0].task_status == TaskStatus.Cancelled
-    assert len(delete_events) == 1
-    assert delete_events[0].instance_id == instance_id_a

src/exo/routing/connection_message.py

@@ -1,4 +1,6 @@
-from exo_pyo3_bindings import PyFromSwarm
+from enum import Enum
+
+from exo_pyo3_bindings import ConnectionUpdate, ConnectionUpdateType
 
 from exo.shared.types.common import NodeId
 from exo.utils.pydantic_ext import CamelCaseModel
@@ -6,10 +8,30 @@ 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
-    connected: bool
+    connection_type: ConnectionMessageType
+    remote_ipv4: str
+    remote_tcp_port: int
 
     @classmethod
-    def from_update(cls, update: PyFromSwarm.Connection) -> "ConnectionMessage":
-        return cls(node_id=NodeId(update.peer_id), connected=update.connected)
+    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,
+        )

src/exo/routing/event_router.py

@@ -1,161 +0,0 @@
-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

@@ -8,16 +8,16 @@ from typing import cast
 from anyio import (
     BrokenResourceError,
     ClosedResourceError,
+    create_task_group,
     move_on_after,
     sleep_forever,
 )
+from anyio.abc import TaskGroup
 from exo_pyo3_bindings import (
     AllQueuesFullError,
     Keypair,
-    MessageTooLargeError,
     NetworkingHandle,
     NoPeersSubscribedToTopicError,
-    PyFromSwarm,
 )
 from filelock import FileLock
 from loguru import logger
@@ -25,7 +25,6 @@ from loguru import logger
 from exo.shared.constants import EXO_NODE_ID_KEYPAIR
 from exo.utils.channels import Receiver, Sender, channel
 from exo.utils.pydantic_ext import CamelCaseModel
-from exo.utils.task_group import TaskGroup
 
 from .connection_message import ConnectionMessage
 from .topics import CONNECTION_MESSAGES, PublishPolicy, TypedTopic
@@ -112,9 +111,10 @@ class Router:
         self._net: NetworkingHandle = handle
         self._tmp_networking_sender: Sender[tuple[str, bytes]] | None = send
         self._id_count = count()
-        self._tg: TaskGroup = TaskGroup()
+        self._tg: TaskGroup | None = None
 
     async def register_topic[T: CamelCaseModel](self, topic: TypedTopic[T]):
+        assert self._tg is None, "Attempted to register topic after setup time"
         send = self._tmp_networking_sender
         if send:
             self._tmp_networking_sender = None
@@ -122,8 +122,7 @@ class Router:
             send = self.networking_receiver.clone_sender()
         router = TopicRouter[T](topic, send)
         self.topic_routers[topic.topic] = cast(TopicRouter[CamelCaseModel], router)
-        if self._tg.is_running():
-            await self._networking_subscribe(topic.topic)
+        await self._networking_subscribe(str(topic.topic))
 
     def sender[T: CamelCaseModel](self, topic: TypedTopic[T]) -> Sender[T]:
         router = self.topic_routers.get(topic.topic, None)
@@ -149,15 +148,14 @@ class Router:
     async def run(self):
         logger.debug("Starting Router")
         try:
-            async with self._tg as tg:
+            async with create_task_group() as tg:
+                self._tg = tg
                 for topic in self.topic_routers:
                     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:
@@ -167,7 +165,9 @@ class Router:
 
     async def shutdown(self):
         logger.debug("Shutting down Router")
-        self._tg.cancel_tasks()
+        if not self._tg:
+            return
+        self._tg.cancel_scope.cancel()
 
     async def _networking_subscribe(self, topic: str):
         await self._net.gossipsub_subscribe(topic)
@@ -178,41 +178,28 @@ class Router:
         logger.info(f"Unsubscribed from {topic}")
 
     async def _networking_recv(self):
-        try:
-            while True:
-                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"
+        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)
+
+    async def _networking_recv_connection_messages(self):
+        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}"
             )
-            raise
+            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)
 
     async def _networking_publish(self):
         with self.networking_receiver as networked_items:
@@ -224,10 +211,6 @@ class Router:
                     pass
                 except AllQueuesFullError:
                     logger.warning(f"All peer queues full, dropping message on {topic}")
-                except MessageTooLargeError:
-                    logger.warning(
-                        f"Message too large for gossipsub on {topic} ({len(data)} bytes), dropping"
-                    )
 
 
 def get_node_id_keypair(
@@ -258,6 +241,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()
+            keypair = Keypair.generate_ed25519()
             f.write(keypair.to_bytes())
             return keypair

src/exo/shared/constants.py

@@ -36,12 +36,26 @@ EXO_MODELS_DIR = (
 
 # Read-only search path for pre-downloaded models (colon-separated directories)
 _EXO_MODELS_PATH_ENV = os.environ.get("EXO_MODELS_PATH", None)
-EXO_MODELS_PATH: tuple[Path, ...] | None = (
-    tuple(Path(p).expanduser() for p in _EXO_MODELS_PATH_ENV.split(":") if p)
-    if _EXO_MODELS_PATH_ENV is not None
-    else None
+
+# Well-known model cache directories from other inference engines
+_WELL_KNOWN_MODEL_PATHS: tuple[Path, ...] = tuple(
+    p for p in (Path.home() / ".cache" / "huggingface" / "hub",) if p.is_dir()
 )
 
+
+def _build_models_path() -> tuple[Path, ...] | None:
+    if _EXO_MODELS_PATH_ENV is not None:
+        user_paths = tuple(
+            Path(p).expanduser() for p in _EXO_MODELS_PATH_ENV.split(":") if p
+        )
+        return user_paths + tuple(
+            p for p in _WELL_KNOWN_MODEL_PATHS if p not in user_paths
+        )
+    return _WELL_KNOWN_MODEL_PATHS if _WELL_KNOWN_MODEL_PATHS else None
+
+
+EXO_MODELS_PATH: tuple[Path, ...] | None = _build_models_path()
+
 _RESOURCES_DIR_ENV = os.environ.get("EXO_RESOURCES_DIR", None)
 RESOURCES_DIR = (
     find_resources() if _RESOURCES_DIR_ENV is None else Path.home() / _RESOURCES_DIR_ENV
@@ -78,6 +92,4 @@ EXO_ENABLE_IMAGE_MODELS = (
     os.getenv("EXO_ENABLE_IMAGE_MODELS", "false").lower() == "true"
 )
 
-EXO_OFFLINE = os.getenv("EXO_OFFLINE", "false").lower() == "true"
-
 EXO_TRACING_ENABLED = os.getenv("EXO_TRACING_ENABLED", "false").lower() == "true"

src/exo/shared/election.py

@@ -4,8 +4,10 @@ import anyio
 from anyio import (
     CancelScope,
     Event,
+    create_task_group,
     get_cancelled_exc_class,
 )
+from anyio.abc import TaskGroup
 from loguru import logger
 
 from exo.routing.connection_message import ConnectionMessage
@@ -13,7 +15,6 @@ from exo.shared.types.commands import ForwarderCommand
 from exo.shared.types.common import NodeId, SessionId
 from exo.utils.channels import Receiver, Sender
 from exo.utils.pydantic_ext import CamelCaseModel
-from exo.utils.task_group import TaskGroup
 
 DEFAULT_ELECTION_TIMEOUT = 3.0
 
@@ -81,12 +82,13 @@ class Election:
         self._candidates: list[ElectionMessage] = []
         self._campaign_cancel_scope: CancelScope | None = None
         self._campaign_done: Event | None = None
-        self._tg = TaskGroup()
+        self._tg: TaskGroup | None = None
 
     async def run(self):
         logger.info("Starting Election")
         try:
-            async with self._tg as tg:
+            async with create_task_group() as tg:
+                self._tg = tg
                 tg.start_soon(self._election_receiver)
                 tg.start_soon(self._connection_receiver)
                 tg.start_soon(self._command_counter)
@@ -122,7 +124,12 @@ class Election:
         )
 
     async def shutdown(self) -> None:
-        self._tg.cancel_tasks()
+        if not self._tg:
+            logger.warning(
+                "Attempted to shutdown election service that was not running"
+            )
+            return
+        self._tg.cancel_scope.cancel()
 
     async def _election_receiver(self) -> None:
         with self._em_receiver as election_messages:
@@ -136,6 +143,7 @@ class Election:
                 if message.clock > self.clock:
                     self.clock = message.clock
                     logger.debug(f"New clock: {self.clock}")
+                    assert self._tg is not None
                     logger.debug("Starting new campaign")
                     candidates: list[ElectionMessage] = [message]
                     logger.debug(f"Candidates: {candidates}")
@@ -170,6 +178,7 @@ class Election:
                 # These messages are strictly peer to peer
                 self.clock += 1
                 logger.debug(f"New clock: {self.clock}")
+                assert self._tg is not None
                 candidates: list[ElectionMessage] = []
                 self._candidates = candidates
                 logger.debug("Starting new campaign")

src/exo/shared/models/model_cards.py

@@ -90,7 +90,6 @@ class ModelCard(CamelCaseModel):
     base_model: str = ""
     capabilities: list[str] = []
     uses_cfg: bool = False
-    trust_remote_code: bool = True
 
     @field_validator("tasks", mode="before")
     @classmethod
@@ -138,7 +137,6 @@ class ModelCard(CamelCaseModel):
             hidden_size=config_data.hidden_size or 0,
             supports_tensor=config_data.supports_tensor,
             tasks=[ModelTask.TextGeneration],
-            trust_remote_code=False,
         )
         await mc.save_to_custom_dir()
         _card_cache[model_id] = mc
@@ -190,8 +188,6 @@ 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
+from exo.routing.connection_message import ConnectionMessage, ConnectionMessageType
 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,7 +327,14 @@ 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(), connected=True))
+            await cm_tx.send(
+                ConnectionMessage(
+                    node_id=NodeId(),
+                    connection_type=ConnectionMessageType.Connected,
+                    remote_ipv4="",
+                    remote_tcp_port=0,
+                )
+            )
 
             # Expect a broadcast for the new round at clock=1
             while True:

src/exo/shared/types/api.py

@@ -262,11 +262,6 @@ class DeleteInstanceResponse(BaseModel):
     instance_id: InstanceId
 
 
-class CancelCommandResponse(BaseModel):
-    message: str
-    command_id: CommandId
-
-
 ImageSize = Literal[
     "auto",
     "512x512",
@@ -442,12 +437,3 @@ 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,8 +2,6 @@
 
 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,
@@ -16,16 +14,3 @@ 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/utils/channels.py

@@ -5,7 +5,7 @@ from math import inf
 from multiprocessing.synchronize import Event
 from queue import Empty, Full
 from types import TracebackType
-from typing import Any, Self
+from typing import Self
 
 from anyio import (
     CapacityLimiter,
@@ -157,7 +157,7 @@ class MpSender[T]:
     ) -> None:
         self.close()
 
-    def __getstate__(self) -> dict[str, Any]:
+    def __getstate__(self):
         d = self.__dict__.copy()
         d.pop("__orig_class__", None)
         return d

src/exo/utils/info_gatherer/info_gatherer.py

@@ -8,7 +8,8 @@ from subprocess import CalledProcessError
 from typing import Self, cast
 
 import anyio
-from anyio import fail_after, open_process, to_thread
+from anyio import create_task_group, fail_after, open_process, to_thread
+from anyio.abc import TaskGroup
 from anyio.streams.buffered import BufferedByteReceiveStream
 from anyio.streams.text import TextReceiveStream
 from loguru import logger
@@ -29,7 +30,6 @@ from exo.shared.types.thunderbolt import (
 )
 from exo.utils.channels import Sender
 from exo.utils.pydantic_ext import TaggedModel
-from exo.utils.task_group import TaskGroup
 
 from .macmon import MacmonMetrics
 from .system_info import (
@@ -381,7 +381,7 @@ class InfoGatherer:
     static_info_poll_interval: float | None = 60
     rdma_ctl_poll_interval: float | None = 10 if IS_DARWIN else None
     disk_poll_interval: float | None = 30
-    _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
+    _tg: TaskGroup = field(init=False, default_factory=create_task_group)
 
     async def run(self):
         async with self._tg as tg:
@@ -408,7 +408,7 @@ class InfoGatherer:
                 await self.info_sender.send(nc)
 
     def shutdown(self):
-        self._tg.cancel_tasks()
+        self._tg.cancel_scope.cancel()
 
     async def _monitor_static_info(self):
         if self.static_info_poll_interval is None:

src/exo/utils/task_group.py

@@ -1,65 +0,0 @@
-from collections.abc import Awaitable, Callable
-from dataclasses import dataclass, field
-from types import TracebackType
-from typing import Any, Unpack
-
-from anyio import create_task_group
-from anyio.abc import TaskGroup as TaskGroupABC
-
-
-@dataclass
-class TaskGroup:
-    _tg: TaskGroupABC | None = field(default=None, init=False)
-    _queued: list[tuple[Any, Any, Any]] | None = field(default_factory=list, init=False)
-
-    def is_running(self) -> bool:
-        return self._tg is not None
-
-    def cancel_tasks(self):
-        assert self._tg
-        self._tg.cancel_scope.cancel()
-
-    def cancel_called(self) -> bool:
-        assert self._tg
-        return self._tg.cancel_scope.cancel_called
-
-    def start_soon[*T](
-        self,
-        func: Callable[[Unpack[T]], Awaitable[Any]],
-        *args: Unpack[T],
-        name: object = None,
-    ) -> None:
-        assert self._tg is not None
-        assert self._queued is None
-        self._tg.start_soon(func, *args, name=name)
-
-    def queue[*T](
-        self,
-        func: Callable[[Unpack[T]], Awaitable[Any]],
-        *args: Unpack[T],
-        name: object = None,
-    ) -> None:
-        assert self._tg is None
-        assert self._queued is not None
-        self._queued.append((func, args, name))
-
-    async def __aenter__(self) -> TaskGroupABC:
-        assert self._tg is None
-        assert self._queued is not None
-        self._tg = create_task_group()
-        r = await self._tg.__aenter__()
-        for func, args, name in self._queued:  # pyright: ignore[reportAny]
-            self._tg.start_soon(func, *args, name=name)  # pyright: ignore[reportAny]
-        self._queued = None
-        return r
-
-    async def __aexit__(
-        self,
-        exc_type: type[BaseException] | None,
-        exc_val: BaseException | None,
-        exc_tb: TracebackType | None,
-    ) -> bool:
-        """Exit the task group context waiting for all tasks to finish."""
-        assert self._tg is not None, "aenter sets self.lazy, so it exists when we aexit"
-        assert self._queued is None
-        return await self._tg.__aexit__(exc_type, exc_val, exc_tb)

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

@@ -0,0 +1,17 @@
+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,7 +16,6 @@ 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
@@ -32,22 +31,14 @@ 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_moe import Qwen3MoeSparseMoeBlock
 from mlx_lm.models.qwen3_next import Model as Qwen3NextModel
-from mlx_lm.models.qwen3_next import (
-    Qwen3NextDecoderLayer,
-    Qwen3NextGatedDeltaNet,
-    Qwen3NextSparseMoeBlock,
-)
+from mlx_lm.models.qwen3_next import Qwen3NextDecoderLayer, 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
+from transformers.models.qwen3.modeling_qwen3 import Qwen3DecoderLayer
 
 from exo.shared.logging import logger
 from exo.shared.types.worker.shards import PipelineShardMetadata
@@ -59,21 +50,6 @@ 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,
@@ -153,11 +129,11 @@ class PipelineFirstLayer(CustomMlxLayer):
 
     def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
         if self.r != 0:
-            # We want to avoid GPU timeout errors by evalling the distributed operation
-            # so that it stays on CPU, which does not have a timeout.
-            mx.eval(x)
             x = mx.distributed.recv_like(x, (self.r - 1), group=self.group)
-            mx.eval(x)
+            if self.is_prefill:
+                # We want to avoid GPU timeout errors by evalling the distributed operation
+                # so that it stays on CPU, which does not have a timeout.
+                mx.eval(x)
         return self.original_layer(x, *args, **kwargs)
 
 
@@ -175,7 +151,6 @@ 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(
@@ -184,34 +159,24 @@ class PipelineLastLayer(CustomMlxLayer):
 
         output: mx.array = self.original_layer(x, *args, **kwargs)
 
-        # Eval layer output to materialize it before send — this splits the graph
-        # so the send is isolated and the receiving rank's recv can complete.
-        mx.eval(output)
-
         if self.r != self.s - 1:
-            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
-                )
+            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
-                if hasattr(_cache, "keys"):  # pyright: ignore[reportAny]
-                    _cache.keys = mx.depends(_cache.keys, output)  # type: ignore
-            mx.eval(output)
-            if cache is not None and hasattr(_cache, "keys"):  # type: ignore
-                mx.eval(_cache.keys)  # type: ignore
+                _cache.keys = mx.depends(_cache.keys, output)  # type: ignore
+            if self.is_prefill:
+                mx.eval(output)
+                if cache is not None:
+                    mx.eval(_cache.keys)  # type: ignore
 
         if not self.is_prefill:
             output = mx.distributed.all_gather(output, group=self.group)[
                 -output.shape[0] :
             ]
-            mx.eval(output)
 
         return output
 
@@ -222,12 +187,6 @@ 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):
@@ -259,32 +218,6 @@ 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,
@@ -355,24 +288,6 @@ 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), (
@@ -402,8 +317,7 @@ 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
-            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
+            dep_cache.keys = mx.depends(dep_cache.keys, logits)  # type: ignore
 
         return logits
 
@@ -526,9 +440,7 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
         )
-    elif isinstance(
-        model, (Qwen3MoeModel, Qwen3NextModel, Qwen3_5TextModel, Qwen3_5MoeModel)
-    ):
+    elif isinstance(model, (Qwen3MoeModel, Qwen3NextModel)):
         tensor_parallel_sharding_strategy = QwenShardingStrategy(
             group,
             all_to_sharded_linear,
@@ -923,14 +835,12 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
         on_timeout: TimeoutCallback | None,
         on_layer_loaded: LayerLoadedCallback | None,
     ) -> nn.Module:
-        model = cast(
-            Qwen3MoeModel | Qwen3NextModel | Qwen3_5TextModel | Qwen3_5MoeModel, model
-        )
+        model = cast(Qwen3MoeModel | Qwen3NextModel, model)
         total = len(model.layers)
         for i, layer in enumerate(model.layers):
             eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
             # Shard the self attention
-            if isinstance(layer, Qwen3MoeDecoderLayer):
+            if isinstance(layer, Qwen3DecoderLayer):
                 layer.self_attn.q_proj = self.all_to_sharded_linear(
                     layer.self_attn.q_proj
                 )
@@ -943,42 +853,17 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                 layer.self_attn.o_proj = self.sharded_to_all_linear(
                     layer.self_attn.o_proj
                 )
-                layer.self_attn.n_heads //= self.N
-                layer.self_attn.n_kv_heads //= self.N
             else:
-                assert isinstance(layer, (Qwen3NextDecoderLayer, Qwen3_5DecoderLayer))
+                assert isinstance(layer, Qwen3NextDecoderLayer)
                 if hasattr(layer, "linear_attn"):
                     linear_attn = layer.linear_attn
 
-                    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.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
+                    )
                     linear_attn.out_proj = self.sharded_to_all_linear(
                         linear_attn.out_proj
                     )
@@ -1040,20 +925,11 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                     layer.self_attn.num_key_value_heads //= self.N
 
             # Shard the MoE.
-            if isinstance(
-                layer.mlp,
-                (
-                    Qwen3MoeSparseMoeBlock,
-                    Qwen3NextSparseMoeBlock,
-                    Qwen3_5SparseMoeBlock,
-                ),
-            ):
+            if isinstance(layer.mlp, (Qwen3MoeSparseMoeBlock, Qwen3NextSparseMoeBlock)):
                 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, Qwen3_5SparseMoeBlock)
-                ):
+                if isinstance(layer.mlp, Qwen3NextSparseMoeBlock):
                     self.all_to_sharded_linear_in_place(
                         layer.mlp.shared_expert.gate_proj
                     )

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

@@ -13,7 +13,8 @@ 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, Model
+from exo.shared.types.mlx import KVCacheType
+from exo.worker.engines.mlx import Model
 from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS
 from exo.worker.runner.bootstrap import logger
 
@@ -253,9 +254,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)
+                c.state = [None] * len(c.state)  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]
         else:
-            c.trim(num_tokens)
+            c.trim(num_tokens)  # pyright: ignore[reportUnknownMemberType]
 
 
 def encode_prompt(tokenizer: TokenizerWrapper, prompt: str) -> mx.array:

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

@@ -1,14 +1,9 @@
-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 (
-    maybe_quantize_kv_cache,
-    stream_generate,
-)
+from mlx_lm.generate import 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
@@ -23,19 +18,13 @@ 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, Model
+from exo.shared.types.mlx import KVCacheType
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 from exo.shared.types.worker.runner_response import (
     GenerationResponse,
 )
-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 import Model
+from exo.worker.engines.mlx.auto_parallel import set_pipeline_prefill
 from exo.worker.engines.mlx.cache import (
     CacheSnapshot,
     KVPrefixCache,
@@ -66,130 +55,6 @@ 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,
@@ -198,7 +63,6 @@ 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.
 
@@ -230,57 +94,31 @@ 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")
 
-    is_pipeline = _has_pipeline_communication_layer(model)
-
-    prefill_step_size = 4096
-
+    # 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
     try:
-        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
+        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
     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.
@@ -293,7 +131,7 @@ def prefill(
                 cache[i] = deepcopy(pre_gen.states[i])  # type: ignore
         else:
             assert not isinstance(c, (ArraysCache, RotatingKVCache))
-            c.trim(2)
+            c.trim(2)  # pyright: ignore[reportUnknownMemberType]
 
     elapsed = time.perf_counter() - start_time
     tokens_per_sec = num_tokens / elapsed if elapsed > 0 else 0.0
@@ -410,9 +248,6 @@ def extract_top_logprobs(
     for i in range(top_logprobs):
         token_id = int(top_indices[i].item())
         token_logprob = float(top_values[i].item())
-        if math.isnan(token_logprob):
-            continue
-
         # Decode token ID to string
         token_str = tokenizer.decode([token_id])
         # Get byte representation
@@ -436,8 +271,6 @@ 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()
@@ -499,7 +332,6 @@ def mlx_generate(
         caches,
         group,
         on_prefill_progress,
-        distributed_prompt_progress_callback,
     )
     cache_snapshots: list[CacheSnapshot] | None = ssm_snapshots_list or None
 
@@ -645,9 +477,6 @@ 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

@@ -2,7 +2,6 @@ import json
 import os
 import re
 import sys
-import tempfile
 import time
 from pathlib import Path
 from typing import Any, cast
@@ -24,7 +23,9 @@ from mlx_lm.models.deepseek_v3 import DeepseekV3Model
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
 from exo.shared.models.model_cards import ModelId
-from exo.worker.engines.mlx.constants import TRUST_REMOTE_CODE
+from exo.worker.engines.mlx.constants import (
+    TRUST_REMOTE_CODE,
+)
 
 try:
     from mlx_lm.tokenizer_utils import load_tokenizer
@@ -40,8 +41,6 @@ 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.tasks import TaskId, TextGeneration
 from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.instances import (
     BoundInstance,
@@ -54,6 +53,7 @@ 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,
@@ -100,13 +100,14 @@ def mlx_distributed_init(
     rank = bound_instance.bound_shard.device_rank
     logger.info(f"Starting initialization for rank {rank}")
 
-    with tempfile.TemporaryDirectory() as tmpdir:
-        coordination_file = str(
-            Path(tmpdir) / f"hosts_{bound_instance.instance.instance_id}_{rank}.json"
-        )
+    coordination_file = None
+    try:
         # TODO: singleton instances
         match bound_instance.instance:
             case MlxRingInstance(hosts_by_node=hosts_by_node, ephemeral_port=_):
+                coordination_file = (
+                    f"./hosts_{bound_instance.instance.instance_id}_{rank}.json"
+                )
                 hosts_for_node = hosts_by_node[bound_instance.bound_node_id]
                 hosts_json = HostList.from_hosts(hosts_for_node).model_dump_json()
 
@@ -129,6 +130,9 @@ def mlx_distributed_init(
                     jaccl_devices[i][i] is None for i in range(len(jaccl_devices))
                 )
                 # Use RDMA connectivity matrix
+                coordination_file = (
+                    f"./hosts_{bound_instance.instance.instance_id}_{rank}.json"
+                )
                 jaccl_devices_json = json.dumps(jaccl_devices)
 
                 with open(coordination_file, "w") as f:
@@ -148,6 +152,10 @@ def mlx_distributed_init(
         logger.info(f"Rank {rank} mlx distributed initialization complete")
 
         return group
+    finally:
+        with contextlib.suppress(FileNotFoundError):
+            if coordination_file:
+                os.remove(coordination_file)
 
 
 def initialize_mlx(
@@ -183,9 +191,10 @@ def load_mlx_items(
                 mx.eval(layer)  # type: ignore
                 if on_layer_loaded is not None:
                     on_layer_loaded(i, total)
-        except ValueError as e:
-            logger.opt(exception=e).debug(
+        except ValueError:
+            logger.debug(
                 "Model architecture doesn't support layer-by-layer progress tracking",
+                exc_info=True,
             )
         mx.eval(model)
         end_time = time.perf_counter()
@@ -208,8 +217,6 @@ def load_mlx_items(
 
     set_wired_limit_for_model(get_weights_size(bound_instance.bound_shard))
 
-    mx.clear_cache()
-
     return cast(Model, model), tokenizer
 
 
@@ -287,11 +294,7 @@ def shard_and_load(
 
 def get_tokenizer(model_path: Path, shard_metadata: ShardMetadata) -> TokenizerWrapper:
     """Load tokenizer for a model shard. Delegates to load_tokenizer_for_model_id."""
-    return load_tokenizer_for_model_id(
-        shard_metadata.model_card.model_id,
-        model_path,
-        trust_remote_code=shard_metadata.model_card.trust_remote_code,
-    )
+    return load_tokenizer_for_model_id(shard_metadata.model_card.model_id, model_path)
 
 
 def get_eos_token_ids_for_model(model_id: ModelId) -> list[int] | None:
@@ -319,14 +322,11 @@ 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
 
 
 def load_tokenizer_for_model_id(
-    model_id: ModelId, model_path: Path, *, trust_remote_code: bool = TRUST_REMOTE_CODE
+    model_id: ModelId, model_path: Path
 ) -> TokenizerWrapper:
     """
     Load tokenizer for a model given its ID and local path.
@@ -395,7 +395,7 @@ def load_tokenizer_for_model_id(
 
     tokenizer = load_tokenizer(
         model_path,
-        tokenizer_config_extra={"trust_remote_code": trust_remote_code},
+        tokenizer_config_extra={"trust_remote_code": TRUST_REMOTE_CODE},
         eos_token_ids=eos_token_ids,
     )
 
@@ -643,7 +643,7 @@ class NullKVCache(KVCache):
         raise NotImplementedError("We should not be setting a NullKVCache.")
 
 
-def mlx_force_oom(size: int = 200000) -> None:
+def mlx_force_oom(size: int = 40000) -> None:
     """
     Force an Out-Of-Memory (OOM) error in MLX by performing large tensor operations.
     """
@@ -670,7 +670,7 @@ def set_wired_limit_for_model(model_size: Memory):
         return
 
     max_rec_size = Memory.from_bytes(
-        int(mx.device_info()["max_recommended_working_set_size"])
+        int(mx.metal.device_info()["max_recommended_working_set_size"])
     )
     if model_size > 0.9 * max_rec_size:
         logger.warning(
@@ -749,55 +749,3 @@ def _parse_kimi_tool_calls(text: str):
         return [_parse_single_tool(match) for match in tool_matches]  # pyright: ignore[reportAny]
     else:
         return [_parse_single_tool(text)]
-
-
-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/main.py

@@ -1,8 +1,10 @@
 from collections import defaultdict
 from datetime import datetime, timezone
+from random import random
 
 import anyio
-from anyio import fail_after
+from anyio import CancelScope, create_task_group, fail_after
+from anyio.abc import TaskGroup
 from loguru import logger
 
 from exo.download.download_utils import resolve_model_in_path
@@ -12,13 +14,17 @@ 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, SystemId
+from exo.shared.types.common import CommandId, NodeId, SessionId, SystemId
 from exo.shared.types.events import (
     Event,
+    EventId,
+    GlobalForwarderEvent,
     IndexedEvent,
     InputChunkReceived,
+    LocalForwarderEvent,
     NodeDownloadProgress,
     NodeGatheredInfo,
     TaskCreated,
@@ -41,10 +47,10 @@ 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
-from exo.utils.task_group import TaskGroup
 from exo.worker.plan import plan
 from exo.worker.runner.runner_supervisor import RunnerSupervisor
 
@@ -53,26 +59,38 @@ class Worker:
     def __init__(
         self,
         node_id: NodeId,
+        session_id: SessionId,
         *,
-        event_receiver: Receiver[IndexedEvent],
-        event_sender: Sender[Event],
+        global_event_receiver: Receiver[GlobalForwarderEvent],
+        local_event_sender: Sender[LocalForwarderEvent],
         # 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.event_receiver = event_receiver
-        self.event_sender = event_sender
+        self.session_id: SessionId = session_id
+
+        self.global_event_receiver = global_event_receiver
+        self.local_event_sender = local_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._tg: TaskGroup = create_task_group()
+
+        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] = {}
@@ -90,12 +108,14 @@ 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.event_sender.close()
+            self.local_event_sender.close()
             self.command_sender.close()
             self.download_command_sender.close()
             for runner in self.runners.values():
@@ -113,22 +133,47 @@ class Worker:
                 )
 
     async def _event_applier(self):
-        with self.event_receiver as events:
-            async for event in events:
+        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]
+
                 # 2. for each event, apply it to the state
-                self.state = apply(self.state, event=event)
-                event = event.event
+                indexed_events = self.event_buffer.drain_indexed()
+                if indexed_events:
+                    self._nack_attempts = 0
 
-                # 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
+                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
                     )
+                    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:
@@ -272,7 +317,7 @@ class Worker:
                     await self._start_runner_task(task)
 
     def shutdown(self):
-        self._tg.cancel_tasks()
+        self._tg.cancel_scope.cancel()
 
     async def _start_runner_task(self, task: Task):
         if (instance := self.state.instances.get(task.instance_id)) is not None:
@@ -280,6 +325,43 @@ 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(
@@ -290,6 +372,21 @@ 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 or task.task_id in runner.in_progress:
+            if task.task_id in runner.completed:
                 continue
 
-            if isinstance(runner.status, (RunnerReady, RunnerRunning)) and all(
+            if isinstance(runner.status, RunnerReady) and all(
                 isinstance(all_runners[global_runner_id], (RunnerReady, RunnerRunning))
                 for global_runner_id in runner.bound_instance.instance.shard_assignments.runner_to_shard
             ):

src/exo/worker/runner/bootstrap.py

@@ -4,7 +4,7 @@ import loguru
 
 from exo.shared.types.events import Event, RunnerStatusUpdated
 from exo.shared.types.tasks import Task, TaskId
-from exo.shared.types.worker.instances import BoundInstance
+from exo.shared.types.worker.instances import BoundInstance, MlxJacclInstance
 from exo.shared.types.worker.runners import RunnerFailed
 from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender
 
@@ -18,33 +18,31 @@ def entrypoint(
     cancel_receiver: MpReceiver[TaskId],
     _logger: "loguru.Logger",
 ) -> None:
-    global logger
-    logger = _logger
-
     fast_synch_override = os.environ.get("EXO_FAST_SYNCH")
-    if fast_synch_override != "off":
+    if fast_synch_override == "on" or (
+        fast_synch_override != "off"
+        and (
+            isinstance(bound_instance.instance, MlxJacclInstance)
+            and len(bound_instance.instance.jaccl_devices) >= 2
+        )
+    ):
         os.environ["MLX_METAL_FAST_SYNCH"] = "1"
     else:
         os.environ["MLX_METAL_FAST_SYNCH"] = "0"
 
+    global logger
+    logger = _logger
+
     logger.info(f"Fast synch flag: {os.environ['MLX_METAL_FAST_SYNCH']}")
 
     # 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 Runner as ImageRunner
-
-            runner = ImageRunner(
-                bound_instance, event_sender, task_receiver, cancel_receiver
-            )
-            runner.main()
+            from exo.worker.runner.image_models.runner import main
         else:
-            from exo.worker.runner.llm_inference.runner import Runner
+            from exo.worker.runner.llm_inference.runner import main
 
-            runner = Runner(
-                bound_instance, event_sender, task_receiver, cancel_receiver
-            )
-            runner.main()
+        main(bound_instance, event_sender, task_receiver, cancel_receiver)
 
     except ClosedResourceError:
         logger.warning("Runner communication closed unexpectedly")

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

@@ -182,266 +182,272 @@ def _send_image_chunk(
         )
 
 
-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
+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))
 
-        soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
-        resource.setrlimit(resource.RLIMIT_NOFILE, (min(max(soft, 2048), hard), hard))
+    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)
 
-        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
+    setup_start_time = time.time()
+    cancelled_tasks = set[TaskId]()
 
-        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)
+    image_model: DistributedImageModel | None = None
+    group = None
 
-        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
+    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)
             )
-        )
-
-    def send_task_status(self, task: Task, status: TaskStatus):
-        self.event_sender.send(
-            TaskStatusUpdated(task_id=task.task_id, task_status=status)
-        )
-
-    def acknowledge_task(self, task: Task):
-        self.event_sender.send(TaskAcknowledged(task_id=task.task_id))
-
-    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
-                )
-                if not was_cancelled:
-                    self.send_task_status(task, TaskStatus.Complete)
-                self.update_status(self.current_status)
-
-                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),
-                                ),
-                            )
+            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
                         )
-                    raise
-                finally:
-                    _send_traces_if_enabled(
-                        self.event_sender, task.task_id, self.device_rank
+                    )
+                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
+                    group = initialize_mlx(bound_instance)
+
+                    logger.info("runner connected")
+                    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(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))
+
+                    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"
                     )
 
-                self.current_status = RunnerReady()
-                logger.info("runner ready")
+                    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),
-                                ),
-                            )
+                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
                         )
-                    raise
-                finally:
-                    _send_traces_if_enabled(
-                        self.event_sender, task.task_id, self.device_rank
                     )
+                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
 
-                self.current_status = RunnerReady()
-                logger.info("runner ready")
+                    try:
+                        image_index = 0
+                        for response in generate_image(
+                            model=image_model, task=task_params
+                        ):
+                            is_primary_output = _is_primary_output_node(shard_metadata)
 
-            case Shutdown():
-                logger.info("runner shutting down")
-                if not TYPE_CHECKING:
-                    del self.image_model, self.group
-                    mx.clear_cache()
-                    import gc
+                            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)
 
-                    gc.collect()
+                    current_status = RunnerReady()
+                    logger.info("runner ready")
 
-                self.update_status(RunnerShuttingDown())
-                self.acknowledge_task(task)
+                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))
 
-                self.current_status = RunnerShutdown()
-            case _:
-                raise ValueError(
-                    f"Received {task.__class__.__name__} outside of state machine in {self.current_status=}"
+                    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
+                    )
                 )
+            event_sender.send(
+                RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
+            )
+
+            if isinstance(current_status, RunnerShutdown):
+                break

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

@@ -1,293 +0,0 @@
-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
-
-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_all_gather_tasks,
-    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,
-        )
-        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

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

@@ -1,376 +0,0 @@
-from collections.abc import Generator
-from functools import cache
-
-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,
-) -> 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)
-
-    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
-) -> 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_tool_calls(
-                    "".join(tool_call_text_parts).strip()
-                )
-                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

@@ -58,16 +58,15 @@ def _flatten(p: dict[str, Any]) -> dict[str, str]:
     }
 
 
-def make_json_parser() -> ToolParser:
-    return ToolParser(
-        start_parsing="<tool_call>",
-        end_parsing="</tool_call>",
-        parse_tool_calls=_parse_json_calls,
-    )
+json_tool_parser = ToolParser(
+    start_parsing="<tool_call>",
+    end_parsing="</tool_call>",
+    parse_tool_calls=_parse_json_calls,
+)
 
 
 def infer_tool_parser(chat_template: str) -> ToolParser | None:
     """Attempt to auto-infer a tool parser from the chat template."""
     if "<tool_call>" in chat_template and "tool_call.name" in chat_template:
-        return make_json_parser()
+        return json_tool_parser
     return None

src/exo/worker/runner/runner_supervisor.py

@@ -1,7 +1,7 @@
 import contextlib
-import multiprocessing as mp
 import signal
 from dataclasses import dataclass, field
+from multiprocessing import Process
 from typing import Self
 
 import anyio
@@ -32,7 +32,6 @@ from exo.shared.types.worker.runners import (
 )
 from exo.shared.types.worker.shards import ShardMetadata
 from exo.utils.channels import MpReceiver, MpSender, Sender, mp_channel
-from exo.utils.task_group import TaskGroup
 from exo.worker.runner.bootstrap import entrypoint
 
 PREFILL_TIMEOUT_SECONDS = 60
@@ -43,21 +42,16 @@ DECODE_TIMEOUT_SECONDS = 5
 class RunnerSupervisor:
     shard_metadata: ShardMetadata
     bound_instance: BoundInstance
-    runner_process: mp.Process
+    runner_process: Process
     initialize_timeout: float
     _ev_recv: MpReceiver[Event]
     _task_sender: MpSender[Task]
     _event_sender: Sender[Event]
     _cancel_sender: MpSender[TaskId]
-    _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: set[TaskId] = field(default_factory=set, 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(
-        default_factory=anyio.CancelScope, init=False
-    )
 
     @classmethod
     def create(
@@ -71,7 +65,7 @@ class RunnerSupervisor:
         task_sender, task_recv = mp_channel[Task]()
         cancel_sender, cancel_recv = mp_channel[TaskId]()
 
-        runner_process = mp.Process(
+        runner_process = Process(
             target=entrypoint,
             args=(
                 bound_instance,
@@ -100,25 +94,15 @@ class RunnerSupervisor:
 
     async def run(self):
         self.runner_process.start()
-        async with self._tg as tg:
-            tg.start_soon(self._watch_runner)
-            tg.start_soon(self._forward_events)
+        await self._forward_events()
 
     def shutdown(self):
         logger.info("Runner supervisor shutting down")
-        self._tg.cancel_tasks()
-        if not self._cancel_watch_runner.cancel_called:
-            self._cancel_watch_runner.cancel()
-        with contextlib.suppress(ClosedResourceError):
-            self._ev_recv.close()
-        with contextlib.suppress(ClosedResourceError):
-            self._task_sender.close()
-        with contextlib.suppress(ClosedResourceError):
-            self._event_sender.close()
+        self._ev_recv.close()
+        self._task_sender.close()
         with contextlib.suppress(ClosedResourceError):
             self._cancel_sender.send(TaskId("CANCEL_CURRENT_TASK"))
-        with contextlib.suppress(ClosedResourceError):
-            self._cancel_sender.close()
+        self._cancel_sender.close()
         self.runner_process.join(5)
         if not self.runner_process.is_alive():
             logger.info("Runner process succesfully terminated")
@@ -158,7 +142,6 @@ 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:
@@ -168,14 +151,13 @@ class RunnerSupervisor:
             await self._check_runner(TimeoutError("cancel pipe blocked"))
 
     async def _forward_events(self):
-        try:
-            with self._ev_recv as events:
+        with self._ev_recv as events:
+            try:
                 async for event in events:
                     if isinstance(event, RunnerStatusUpdated):
                         self.status = event.runner_status
                     if isinstance(event, TaskAcknowledged):
                         self.pending.pop(event.task_id).set()
-                        self.in_progress.add(event.task_id)
                         continue
                     if (
                         isinstance(event, TaskStatusUpdated)
@@ -192,37 +174,27 @@ class RunnerSupervisor:
                                 RunnerShuttingDown,
                             ),
                         )
-                        self.in_progress.discard(event.task_id)
                         self.completed.add(event.task_id)
                     await self._event_sender.send(event)
-        except (ClosedResourceError, BrokenResourceError) as e:
-            await self._check_runner(e)
-        finally:
-            for tid in self.pending:
-                self.pending[tid].set()
+            except (ClosedResourceError, BrokenResourceError) as e:
+                await self._check_runner(e)
+                for tid in self.pending:
+                    self.pending[tid].set()
+        self._event_sender.close()
 
     def __del__(self) -> None:
         if self.runner_process.is_alive():
-            logger.critical("RunnerSupervisor was not stopped cleanly.")
+            logger.warning("RunnerSupervisor was not stopped cleanly.")
             with contextlib.suppress(ValueError):
                 self.runner_process.kill()
 
-    async def _watch_runner(self) -> None:
-        with self._cancel_watch_runner:
-            while True:
-                await anyio.sleep(5)
-                if not self.runner_process.is_alive():
-                    await self._check_runner(RuntimeError("Runner found to be dead"))
-
     async def _check_runner(self, e: Exception) -> None:
-        if not self._cancel_watch_runner.cancel_called:
-            self._cancel_watch_runner.cancel()
         logger.info("Checking runner's status")
         if self.runner_process.is_alive():
             logger.info("Runner was found to be alive, attempting to join process")
             await to_thread.run_sync(self.runner_process.join, 5)
         rc = self.runner_process.exitcode
-        logger.info(f"Runner exited with exit code {rc}")
+        logger.info(f"RunnerSupervisor exited with exit code {rc}")
         if rc == 0:
             return
 
@@ -235,19 +207,15 @@ class RunnerSupervisor:
         else:
             cause = f"exitcode={rc}"
 
-        logger.opt(exception=e).error(f"Runner terminated with {cause}")
+        logger.opt(exception=e).error(f"Runner terminated ({cause})")
 
         try:
-            self.status = RunnerFailed(error_message=f"Terminated ({cause})")
-            with anyio.CancelScope(shield=True):
-                await self._event_sender.send(
-                    RunnerStatusUpdated(
-                        runner_id=self.bound_instance.bound_runner_id,
-                        runner_status=RunnerFailed(
-                            error_message=f"Terminated ({cause})"
-                        ),
-                    )
+            await self._event_sender.send(
+                RunnerStatusUpdated(
+                    runner_id=self.bound_instance.bound_runner_id,
+                    runner_status=RunnerFailed(error_message=f"Terminated ({cause})"),
                 )
+            )
         except (ClosedResourceError, BrokenResourceError):
             logger.warning(
                 "Event sender already closed, unable to report runner failure"

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

@@ -20,8 +20,6 @@ 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,14 +143,7 @@ class TestKVPrefixCacheWithModel:
         cache = make_kv_cache(model)
 
         _, _, snapshots = prefill(
-            model,
-            tokenizer,
-            make_sampler(0.0),
-            tokens,
-            cache,
-            group=None,
-            on_prefill_progress=None,
-            distributed_prompt_progress_callback=None,
+            model, tokenizer, make_sampler(0.0), tokens, cache, group=None
         )
 
         # Cache should now hold the prompt tokens minus one
@@ -171,14 +164,7 @@ class TestKVPrefixCacheWithModel:
         cache = make_kv_cache(model)
 
         _, _, snapshots = prefill(
-            model,
-            tokenizer,
-            make_sampler(0.0),
-            tokens,
-            cache,
-            group=None,
-            on_prefill_progress=None,
-            distributed_prompt_progress_callback=None,
+            model, tokenizer, make_sampler(0.0), tokens, cache, group=None
         )
 
         kv_prefix_cache = KVPrefixCache(None)
@@ -214,14 +200,7 @@ class TestKVPrefixCacheWithModel:
         cache = make_kv_cache(model)
 
         _, _, snapshots = prefill(
-            model,
-            tokenizer,
-            make_sampler(0.0),
-            short_tokens,
-            cache,
-            group=None,
-            on_prefill_progress=None,
-            distributed_prompt_progress_callback=None,
+            model, tokenizer, make_sampler(0.0), short_tokens, cache, group=None
         )
 
         kv_prefix_cache = KVPrefixCache(None)
@@ -266,14 +245,7 @@ class TestKVPrefixCacheWithModel:
         cache = make_kv_cache(model)
 
         _, _, snapshots = prefill(
-            model,
-            tokenizer,
-            make_sampler(0.0),
-            tokens,
-            cache,
-            group=None,
-            on_prefill_progress=None,
-            distributed_prompt_progress_callback=None,
+            model, tokenizer, make_sampler(0.0), tokens, cache, group=None
         )
 
         kv_prefix_cache = KVPrefixCache(None)
@@ -313,14 +285,7 @@ class TestKVPrefixCacheWithModel:
         cache = make_kv_cache(model)
 
         _, _, snapshots = prefill(
-            model,
-            tokenizer,
-            make_sampler(0.0),
-            tokens,
-            cache,
-            group=None,
-            on_prefill_progress=None,
-            distributed_prompt_progress_callback=None,
+            model, tokenizer, make_sampler(0.0), tokens, cache, group=None
         )
 
         kv_prefix_cache = KVPrefixCache(None)
@@ -548,16 +513,7 @@ 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,
-                on_prefill_progress=None,
-                distributed_prompt_progress_callback=None,
-            )
+            prefill(model, tokenizer, make_sampler(0.0), tokens, cache, group=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)
@@ -582,16 +538,7 @@ 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,
-                on_prefill_progress=None,
-                distributed_prompt_progress_callback=None,
-            )
+            prefill(model, tokenizer, make_sampler(0.0), tokens, cache, group=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

@@ -1,512 +0,0 @@
-# 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 (