flush old addrs

.mlx_typings/mlx/core/__init__.pyi

@@ -2396,7 +2396,7 @@ def degrees(a: array, /, *, stream: Stream | Device | None = ...) -> array:
         array: The angles in degrees.
     """
 
-def depends[T](inputs: T, dependencies: array | Sequence[array]) -> T:
+def depends(inputs: array | Sequence[array], dependencies: array | Sequence[array]):
     """
     Insert dependencies between arrays in the graph. The outputs are
     identical to ``inputs`` but with dependencies on ``dependencies``.

.mlx_typings/mlx/nn/__init__.pyi

@@ -1,5 +1,9 @@
-from .layers import *
-from .utils import *
+"""
+This type stub file was generated by pyright.
+"""
+
+from layers import *
+from utils import *
 
 from . import init as init
 from . import losses as losses

.mlx_typings/mlx/nn/layers/__init__.pyi

@@ -1,16 +1,20 @@
-from .activations import *
-from .base import *
-from .containers import *
-from .convolution import *
-from .convolution_transpose import *
-from .distributed import *
-from .dropout import *
-from .embedding import *
-from .linear import *
-from .normalization import *
-from .pooling import *
-from .positional_encoding import *
-from .quantized import *
-from .recurrent import *
-from .transformer import *
-from .upsample import *
+"""
+This type stub file was generated by pyright.
+"""
+
+from activations import *
+from base import *
+from containers import *
+from convolution import *
+from convolution_transpose import *
+from distributed import *
+from dropout import *
+from embedding import *
+from linear import *
+from normalization import *
+from pooling import *
+from positional_encoding import *
+from quantized import *
+from recurrent import *
+from transformer import *
+from upsample import *

.mlx_typings/mlx/nn/layers/base.pyi

@@ -53,7 +53,7 @@ class Module(dict):
         mx.eval(model.parameters())
     """
 
-    def __call__(self, *args: Any, **kwargs: Any) -> mx.array: ...
+    __call__: Callable
     def __init__(self) -> None:
         """Should be called by the subclasses of ``Module``."""
 

.mlx_typings/mlx_lm/generate.pyi

@@ -30,7 +30,7 @@ def str2bool(string):  # -> bool:
 def setup_arg_parser():  # -> ArgumentParser:
     """Set up and return the argument parser."""
 
-generation_stream: mx.Stream
+generation_stream = ...
 
 @contextlib.contextmanager
 def wired_limit(
@@ -266,12 +266,12 @@ def _merge_caches(caches: Any) -> List[Any]: ...
 class Batch:
     uids: List[int]
     y: mx.array
-    logprobs: List[mx.array] | mx.array
+    logprobs: mx.array
     max_tokens: List[int]
     num_tokens: List[int]
     cache: List[Any]
-    samplers: List[Callable[[mx.array], mx.array] | None]
-    logits_processors: List[List[Callable[[mx.array, mx.array], mx.array]]]
+    samplers: List[Any]
+    logits_processors: List[Any]
     tokens: List[mx.array]
     def __len__(self) -> int: ...
     def filter(self, keep_idx: List[int]) -> None: ...
@@ -279,18 +279,13 @@ class Batch:
     def extract_cache(self, idx: int) -> List[Any]: ...
 
 class BatchGenerator:
-    model: nn.Module
-    sampler: Callable[[mx.array], mx.array]
-    stop_tokens: set[int]
+    model: Any
     max_kv_size: Optional[int]
     prefill_step_size: int
-    completion_batch_size: int
-    prefill_batch_size: int
     unprocessed_prompts: List[Any]
     active_batch: Optional[Batch]
     prompt_progress_callback: Callable[[List[Tuple[int, int, int]]], None]
     _stats: BatchStats
-    _next_count: int
 
     @dataclass
     class Response:

.mlx_typings/mlx_lm/models/cache.pyi

@@ -88,8 +88,8 @@ def create_attention_mask(
 ) -> array | Literal["causal"] | None: ...
 
 class _BaseCache(Cache):
-    keys: mx.array | None
-    values: mx.array | None
+    keys: mx.array
+    values: mx.array
     offset: int
     @property
     def state(self) -> tuple[mx.array | None, mx.array | None]: ...
@@ -268,14 +268,29 @@ class CacheList(_BaseCache):
         """
 
 class BatchKVCache(_BaseCache):
-    step: int
-    keys: array | None
-    values: array | None
-    offset: array
-    left_padding: array
-    _idx: int
-    def __init__(self, left_padding: List[int]) -> None: ...
-    def update_and_fetch(self, keys: array, values: array) -> tuple[array, array]: ...
+    step = ...
+    def __init__(self, left_padding: List[int]) -> None:
+        """
+        The BatchKV cache expects inputs to be left-padded.
+
+        E.g. the following prompts:
+
+            [1, 3, 5]
+            [7]
+            [2, 6, 8, 9]
+
+        Should be padded like so:
+
+            [0, 1, 3, 5]
+            [0, 0, 0, 7]
+            [2, 6, 8, 9]
+
+        And ``left_padding`` specifies the amount of padding for each.
+        In this case, ``left_padding = [1, 3, 0]``.
+        """
+
+    def update_and_fetch(self, keys, values):  # -> tuple[array | Any, array | Any]:
+        ...
     @property
     def state(
         self,
@@ -301,21 +316,12 @@ class BatchKVCache(_BaseCache):
         """
 
 class BatchRotatingKVCache(_BaseCache):
-    step: int
-    keys: array | None
-    values: array | None
-    offset: array
-    left_padding: array
-    max_size: int
-    _idx: int
-    _offset: int
-    rotated: bool
-    _lengths: array | None
-    def __init__(self, max_size: int, left_padding: List[int]) -> None: ...
-    def _trim(self, trim_size: int, v: array, append: array | None = ...) -> array: ...
-    def _update_in_place(self, keys: array, values: array) -> tuple[array, array]: ...
-    def _update_concat(self, keys: array, values: array) -> tuple[array, array]: ...
-    def update_and_fetch(self, keys: array, values: array) -> tuple[array, array]: ...
+    step = ...
+    def __init__(self, max_size, left_padding: List[int]) -> None: ...
+    def update_and_fetch(
+        self, keys, values
+    ):  # -> tuple[array | Any, array | Any] | tuple[array | Any, array | Any | None]:
+        ...
     @property
     def state(
         self,

.mlx_typings/mlx_lm/models/gated_delta.pyi

@@ -1,35 +0,0 @@
-from typing import Optional
-
-import mlx.core as mx
-
-def compute_g(A_log: mx.array, a: mx.array, dt_bias: mx.array) -> mx.array: ...
-def gated_delta_update(
-    q: mx.array,
-    k: mx.array,
-    v: mx.array,
-    a: mx.array,
-    b: mx.array,
-    A_log: mx.array,
-    dt_bias: mx.array,
-    state: Optional[mx.array] = ...,
-    mask: Optional[mx.array] = ...,
-    use_kernel: bool = ...,
-) -> tuple[mx.array, mx.array]: ...
-def gated_delta_ops(
-    q: mx.array,
-    k: mx.array,
-    v: mx.array,
-    g: mx.array,
-    beta: mx.array,
-    state: Optional[mx.array] = ...,
-    mask: Optional[mx.array] = ...,
-) -> tuple[mx.array, mx.array]: ...
-def gated_delta_kernel(
-    q: mx.array,
-    k: mx.array,
-    v: mx.array,
-    g: mx.array,
-    beta: mx.array,
-    state: mx.array,
-    mask: Optional[mx.array] = ...,
-) -> tuple[mx.array, mx.array]: ...

.mlx_typings/mlx_lm/models/rope_utils.pyi

@@ -1,51 +0,0 @@
-from typing import Any, Optional
-
-import mlx.nn as nn
-
-class YarnRoPE(nn.Module):
-    def __init__(
-        self,
-        dims: int,
-        traditional: bool = ...,
-        max_position_embeddings: int = ...,
-        base: float = ...,
-        scaling_factor: float = ...,
-        original_max_position_embeddings: int = ...,
-        beta_fast: float = ...,
-        beta_slow: float = ...,
-        mscale: float = ...,
-        mscale_all_dim: float = ...,
-    ) -> None: ...
-
-class Llama3RoPE(nn.Module):
-    def __init__(
-        self,
-        dims: int,
-        traditional: bool = ...,
-        max_position_embeddings: int = ...,
-        base: float = ...,
-        scaling_factor: float = ...,
-        original_max_position_embeddings: int = ...,
-        low_freq_factor: float = ...,
-        high_freq_factor: float = ...,
-    ) -> None: ...
-
-class SuScaledRoPE(nn.Module):
-    def __init__(
-        self,
-        dims: int,
-        traditional: bool = ...,
-        max_position_embeddings: int = ...,
-        base: float = ...,
-        short_factor: Any = ...,
-        long_factor: Any = ...,
-        original_max_position_embeddings: int = ...,
-    ) -> None: ...
-
-def initialize_rope(
-    dims: int,
-    base: float = ...,
-    traditional: bool = ...,
-    scaling_config: Optional[dict[str, Any]] = ...,
-    max_position_embeddings: Optional[int] = ...,
-) -> nn.Module: ...

CONTRIBUTING.md

@@ -11,18 +11,9 @@ To run EXO from source:
   ```bash
   brew install uv
   ```
-- [rust](https://github.com/rust-lang/rustup) (to build Rust bindings, nightly for now)
-  ```bash
-  curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
-  rustup toolchain install nightly
-  ```
 - [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
-  Use the pinned fork revision used by this repo instead of Homebrew `macmon`.
   ```bash
-  cargo install --git https://github.com/swiftraccoon/macmon \
-    --rev 9154d234f763fbeffdcb4135d0bbbaf80609699b \
-    macmon \
-    --force
+  brew install macmon
   ```
 
 ```bash

Cargo.toml

@@ -1,6 +1,11 @@
 [workspace]
 resolver = "3"
-members = ["rust/networking", "rust/exo_pyo3_bindings", "rust/util"]
+members = [
+  "rust/networking",
+  "rust/exo_pyo3_bindings",
+  "rust/util",
+  "rust/babblerd",
+]
 
 [workspace.package]
 version = "0.0.1"
@@ -43,7 +48,6 @@ log = "0.4"
 
 # networking
 libp2p = "0.56"
-libp2p-tcp = "0.44"
 
 [workspace.lints.rust]
 static_mut_refs = "warn"      # Or use "warn" instead of deny

README.md

@@ -95,10 +95,11 @@ Then restart the Nix daemon: `sudo launchctl kickstart -k system/org.nixos.nix-d
   /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"
   ```
 - [uv](https://github.com/astral-sh/uv) (for Python dependency management)
+- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
 - [node](https://github.com/nodejs/node) (for building the dashboard)
 
   ```bash
-  brew install uv node
+  brew install uv macmon node
   ```
 - [rust](https://github.com/rust-lang/rustup) (to build Rust bindings, nightly for now)
 
@@ -106,17 +107,6 @@ Then restart the Nix daemon: `sudo launchctl kickstart -k system/org.nixos.nix-d
   curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
   rustup toolchain install nightly
   ```
-- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
-
-  Install the pinned fork revision used by this repo instead of Homebrew `macmon`.
-  Homebrew `macmon 0.6.1` still crashes on Apple M5.
-
-  ```bash
-  cargo install --git https://github.com/swiftraccoon/macmon \
-    --rev 9154d234f763fbeffdcb4135d0bbbaf80609699b \
-    macmon \
-    --force
-  ```
 
 Clone the repo, build the dashboard, and run exo:
 

app/EXO/EXO/ExoProcessController.swift

@@ -4,7 +4,6 @@ import Foundation
 
 private let customNamespaceKey = "EXOCustomNamespace"
 private let hfTokenKey = "EXOHFToken"
-private let hfEndpointKey = "EXOHFEndpoint"
 private let enableImageModelsKey = "EXOEnableImageModels"
 private let offlineModeKey = "EXOOfflineMode"
 private let onboardingCompletedKey = "EXOOnboardingCompleted"
@@ -54,14 +53,6 @@ final class ExoProcessController: ObservableObject {
             UserDefaults.standard.set(hfToken, forKey: hfTokenKey)
         }
     }
-    @Published var hfEndpoint: String = {
-        return UserDefaults.standard.string(forKey: hfEndpointKey) ?? ""
-    }()
-    {
-        didSet {
-            UserDefaults.standard.set(hfEndpoint, forKey: hfEndpointKey)
-        }
-    }
     @Published var enableImageModels: Bool = {
         return UserDefaults.standard.bool(forKey: enableImageModelsKey)
     }()
@@ -282,9 +273,6 @@ final class ExoProcessController: ObservableObject {
         if !hfToken.isEmpty {
             environment["HF_TOKEN"] = hfToken
         }
-        if !hfEndpoint.isEmpty {
-            environment["HF_ENDPOINT"] = hfEndpoint
-        }
         if enableImageModels {
             environment["EXO_ENABLE_IMAGE_MODELS"] = "true"
         }

app/EXO/EXO/Views/SettingsView.swift

@@ -12,7 +12,6 @@ struct SettingsView: View {
 
     @State private var pendingNamespace: String = ""
     @State private var pendingHFToken: String = ""
-    @State private var pendingHFEndpoint: String = ""
     @State private var pendingEnableImageModels = false
     @State private var pendingOfflineMode = false
     @State private var needsRestart = false
@@ -43,7 +42,6 @@ struct SettingsView: View {
         .onAppear {
             pendingNamespace = controller.customNamespace
             pendingHFToken = controller.hfToken
-            pendingHFEndpoint = controller.hfEndpoint
             pendingEnableImageModels = controller.enableImageModels
             pendingOfflineMode = controller.offlineMode
             needsRestart = false
@@ -76,17 +74,6 @@ struct SettingsView: View {
                     .foregroundColor(.secondary)
             }
 
-            Section {
-                LabeledContent("HuggingFace Endpoint") {
-                    TextField("default", text: $pendingHFEndpoint)
-                        .textFieldStyle(.roundedBorder)
-                        .frame(width: 200)
-                }
-                Text("Defaults to huggingface.co. Use a mirror (e.g. hf-mirror.com) for China.")
-                    .font(.caption)
-                    .foregroundColor(.secondary)
-            }
-
             Section {
                 Toggle("Offline Mode", isOn: $pendingOfflineMode)
                 Text("Skip internet checks and use only locally available models.")
@@ -467,7 +454,6 @@ struct SettingsView: View {
 
     private var hasGeneralChanges: Bool {
         pendingNamespace != controller.customNamespace || pendingHFToken != controller.hfToken
-            || pendingHFEndpoint != controller.hfEndpoint
             || pendingOfflineMode != controller.offlineMode
     }
 
@@ -478,7 +464,6 @@ struct SettingsView: View {
     private func applyGeneralSettings() {
         controller.customNamespace = pendingNamespace
         controller.hfToken = pendingHFToken
-        controller.hfEndpoint = pendingHFEndpoint
         controller.offlineMode = pendingOfflineMode
         restartIfRunning()
     }

bench/exo_bench.py

@@ -501,21 +501,23 @@ def main() -> int:
                                     for x, _ in batch_results
                                     if x["stats"]["generation_tps"] > 0
                                 ]
-                                per_req_tps = (
+                                agg_gen_tps = (
                                     mean(valid_gen_tps) if valid_gen_tps else 0.0
                                 )
-                                agg_gen_tps = per_req_tps * concurrency
+                                gen_tps = agg_gen_tps / concurrency
                                 logger.info(
                                     f"[concurrent {concurrency}x]  "
                                     f"agg_gen_tps={agg_gen_tps:.2f}  "
-                                    f"per_req_tps={per_req_tps:.2f}  "
+                                    f"gen_tps={gen_tps:.2f}  "
                                     f"errors={batch_errors}"
                                 )
 
                     if runs:
                         prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
-                        per_req_tps = mean(x["stats"]["generation_tps"] for x in runs)
-                        gen_tps = per_req_tps * concurrency
+                        gen_tps = mean(
+                            x["stats"]["generation_tps"] / x["concurrency"]
+                            for x in runs
+                        )
                         ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
                         gtok = mean(x["stats"]["generation_tokens"] for x in runs)
                         peak = mean(

bench/test_mlx_bandwidth.py

@@ -0,0 +1,377 @@
+# type: ignore
+import argparse
+import json
+import os
+import statistics
+import sys
+import tempfile
+import time
+
+import mlx.core as mx
+
+DTYPE_MAP = {
+    "float32": (mx.float32, 4),
+    "float16": (mx.float16, 2),
+    "bfloat16": (mx.bfloat16, 2),
+}
+
+SIZES = [
+    1 * 1024,
+    4 * 1024,
+    16 * 1024,
+    64 * 1024,
+    256 * 1024,
+    1 * 1024 * 1024,
+    4 * 1024 * 1024,
+    16 * 1024 * 1024,
+    64 * 1024 * 1024,
+    256 * 1024 * 1024,
+    1 * 1024 * 1024 * 1024,
+    2 * 1024 * 1024 * 1024,
+    4 * 1024 * 1024 * 1024,
+    8 * 1024 * 1024 * 1024,
+]
+
+
+def format_bytes(n: int) -> str:
+    if n >= 1024 * 1024 * 1024:
+        return f"{n / (1024 * 1024 * 1024):.0f} GB"
+    if n >= 1024 * 1024:
+        return f"{n / (1024 * 1024):.0f} MB"
+    if n >= 1024:
+        return f"{n / 1024:.0f} KB"
+    return f"{n} B"
+
+
+def format_time(seconds: float) -> str:
+    if seconds >= 1.0:
+        return f"{seconds:.3f} s"
+    if seconds >= 0.001:
+        return f"{seconds * 1000:.2f} ms"
+    return f"{seconds * 1_000_000:.1f} us"
+
+
+def format_bandwidth(bytes_per_sec: float) -> str:
+    if bytes_per_sec >= 1024 * 1024 * 1024:
+        return f"{bytes_per_sec / (1024 * 1024 * 1024):.2f} GB/s"
+    if bytes_per_sec >= 1024 * 1024:
+        return f"{bytes_per_sec / (1024 * 1024):.1f} MB/s"
+    return f"{bytes_per_sec / 1024:.1f} KB/s"
+
+
+def barrier(group: mx.distributed.Group) -> None:
+    mx.eval(mx.distributed.all_sum(mx.array(1.0), group=group))
+
+
+def init_ring(
+    rank: int, self_ip: str, peer_ip: str, port: int, tmpdir: str
+) -> mx.distributed.Group:
+    if rank == 0:
+        hosts = [f"{self_ip}:{port}", f"{peer_ip}:{port}"]
+    else:
+        hosts = [f"{peer_ip}:{port}", f"{self_ip}:{port}"]
+
+    hostfile = os.path.join(tmpdir, "hosts.json")
+    with open(hostfile, "w") as f:
+        json.dump(hosts, f)
+
+    for var in ("MLX_HOSTFILE", "MLX_RANK", "MLX_IBV_DEVICES", "MLX_JACCL_COORDINATOR"):
+        os.environ.pop(var, None)
+
+    os.environ["MLX_HOSTFILE"] = hostfile
+    os.environ["MLX_RANK"] = str(rank)
+    return mx.distributed.init(backend="ring", strict=True)
+
+
+def init_jaccl(
+    rank: int, interface: str, coordinator: str, port: int, tmpdir: str
+) -> mx.distributed.Group:
+    devices = [[None, interface], [interface, None]]
+    devfile = os.path.join(tmpdir, "devices.json")
+    with open(devfile, "w") as f:
+        json.dump(devices, f)
+
+    for var in ("MLX_HOSTFILE", "MLX_RANK", "MLX_IBV_DEVICES", "MLX_JACCL_COORDINATOR"):
+        os.environ.pop(var, None)
+
+    os.environ["MLX_IBV_DEVICES"] = devfile
+    os.environ["MLX_RANK"] = str(rank)
+    if rank == 0:
+        os.environ["MLX_JACCL_COORDINATOR"] = f"0.0.0.0:{port}"
+    else:
+        os.environ["MLX_JACCL_COORDINATOR"] = coordinator
+
+    return mx.distributed.init(backend="jaccl", strict=True)
+
+
+def bench_unidirectional(
+    group: mx.distributed.Group,
+    rank: int,
+    size_bytes: int,
+    dtype: mx.Dtype,
+    element_size: int,
+    warmup: int,
+    iterations: int,
+) -> list[float]:
+    n_elements = size_bytes // element_size
+    tensor = mx.random.normal(shape=(n_elements,)).astype(dtype)
+    mx.eval(tensor)
+
+    for _ in range(warmup):
+        if rank == 0:
+            sent = mx.distributed.send(tensor, dst=1, group=group)
+            mx.eval(sent)
+        else:
+            received = mx.distributed.recv_like(tensor, src=0, group=group)
+            mx.eval(received)
+        barrier(group)
+
+    times: list[float] = []
+    for _ in range(iterations):
+        barrier(group)
+        t0 = time.perf_counter()
+        if rank == 0:
+            sent = mx.distributed.send(tensor, dst=1, group=group)
+            mx.eval(sent)
+        else:
+            received = mx.distributed.recv_like(tensor, src=0, group=group)
+            mx.eval(received)
+        barrier(group)
+        t1 = time.perf_counter()
+        times.append(t1 - t0)
+
+    return times
+
+
+def bench_rtt(
+    group: mx.distributed.Group,
+    rank: int,
+    size_bytes: int,
+    dtype: mx.Dtype,
+    element_size: int,
+    warmup: int,
+    iterations: int,
+) -> list[float]:
+    n_elements = size_bytes // element_size
+    tensor = mx.random.normal(shape=(n_elements,)).astype(dtype)
+    mx.eval(tensor)
+
+    for _ in range(warmup):
+        if rank == 0:
+            sent = mx.distributed.send(tensor, dst=1, group=group)
+            mx.eval(sent)
+            received = mx.distributed.recv_like(tensor, src=1, group=group)
+            mx.eval(received)
+        else:
+            received = mx.distributed.recv_like(tensor, src=0, group=group)
+            mx.eval(received)
+            sent = mx.distributed.send(received, dst=0, group=group)
+            mx.eval(sent)
+        barrier(group)
+
+    times: list[float] = []
+    for _ in range(iterations):
+        barrier(group)
+        t0 = time.perf_counter()
+        if rank == 0:
+            sent = mx.distributed.send(tensor, dst=1, group=group)
+            mx.eval(sent)
+            received = mx.distributed.recv_like(tensor, src=1, group=group)
+            mx.eval(received)
+        else:
+            received = mx.distributed.recv_like(tensor, src=0, group=group)
+            mx.eval(received)
+            sent = mx.distributed.send(received, dst=0, group=group)
+            mx.eval(sent)
+        barrier(group)
+        t1 = time.perf_counter()
+        times.append(t1 - t0)
+
+    return times
+
+
+def bench_all_gather(
+    group: mx.distributed.Group,
+    rank: int,
+    size_bytes: int,
+    dtype: mx.Dtype,
+    element_size: int,
+    warmup: int,
+    iterations: int,
+) -> list[float]:
+    n_elements = (size_bytes // 2) // element_size
+    tensor = mx.random.normal(shape=(n_elements,)).astype(dtype)
+    mx.eval(tensor)
+
+    for _ in range(warmup):
+        gathered = mx.distributed.all_gather(tensor, group=group)
+        mx.eval(gathered)
+        barrier(group)
+
+    times: list[float] = []
+    for _ in range(iterations):
+        barrier(group)
+        t0 = time.perf_counter()
+        gathered = mx.distributed.all_gather(tensor, group=group)
+        mx.eval(gathered)
+        t1 = time.perf_counter()
+        times.append(t1 - t0)
+
+    return times
+
+
+def print_table(title: str, rows: list[dict[str, str]]) -> None:
+    print(f"\n=== {title} ===")
+    headers = ["Size", "Median", "Min", "Max", "Bandwidth"]
+    widths = [
+        max(len(h), max((len(r[h]) for r in rows), default=0)) + 2 for h in headers
+    ]
+    header_line = "".join(h.ljust(w) for h, w in zip(headers, widths, strict=True))
+    print(header_line)
+    print("-" * len(header_line))
+    for row in rows:
+        print("".join(row[h].ljust(w) for h, w in zip(headers, widths, strict=True)))
+
+
+def run_bench(
+    name: str,
+    bench_fn,
+    group: mx.distributed.Group,
+    rank: int,
+    dtype: mx.Dtype,
+    element_size: int,
+    warmup: int,
+    iterations: int,
+    bw_multiplier: int = 1,
+) -> None:
+    rows: list[dict[str, str]] = []
+    for size in SIZES:
+        if rank == 0:
+            print(f"  {name}: {format_bytes(size)}...", end="", flush=True)
+        times = bench_fn(group, rank, size, dtype, element_size, warmup, iterations)
+        if rank == 0:
+            med = statistics.median(times)
+            mn = min(times)
+            mx_ = max(times)
+            bw = (size * bw_multiplier) / med
+            rows.append(
+                {
+                    "Size": format_bytes(size),
+                    "Median": format_time(med),
+                    "Min": format_time(mn),
+                    "Max": format_time(mx_),
+                    "Bandwidth": format_bandwidth(bw),
+                }
+            )
+            print(f" {format_bandwidth(bw)}")
+    if rank == 0:
+        print_table(name, rows)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="MLX Distributed Communication Benchmark"
+    )
+    subparsers = parser.add_subparsers(dest="backend", required=True)
+
+    ring_parser = subparsers.add_parser("ring")
+    ring_parser.add_argument("--rank", type=int, required=True, choices=[0, 1])
+    ring_parser.add_argument("--self-ip", required=True)
+    ring_parser.add_argument("--peer-ip", required=True)
+    ring_parser.add_argument("--port", type=int, default=5555)
+
+    jaccl_parser = subparsers.add_parser("jaccl")
+    jaccl_parser.add_argument("--rank", type=int, required=True, choices=[0, 1])
+    jaccl_parser.add_argument("--interface", required=True)
+    jaccl_parser.add_argument(
+        "--coordinator",
+        type=str,
+        default=None,
+        help="IP:PORT of rank 0 (required for rank 1)",
+    )
+    jaccl_parser.add_argument(
+        "--port", type=int, default=9999, help="Coordinator port (rank 0 only)"
+    )
+
+    for p in [ring_parser, jaccl_parser]:
+        p.add_argument("--warmup", type=int, default=3)
+        p.add_argument("--iterations", type=int, default=10)
+        p.add_argument("--dtype", choices=list(DTYPE_MAP.keys()), default="float32")
+
+    args = parser.parse_args()
+
+    if args.backend == "jaccl" and args.rank == 1 and args.coordinator is None:
+        jaccl_parser.error("--coordinator is required for rank 1")
+
+    return args
+
+
+def main() -> int:
+    args = parse_args()
+    dtype, element_size = DTYPE_MAP[args.dtype]
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        if args.backend == "ring":
+            print(f"Initializing ring backend (rank {args.rank})...")
+            group = init_ring(args.rank, args.self_ip, args.peer_ip, args.port, tmpdir)
+        else:
+            print(f"Initializing jaccl backend (rank {args.rank})...")
+            group = init_jaccl(
+                args.rank, args.interface, args.coordinator or "", args.port, tmpdir
+            )
+
+        print(f"Rank {group.rank()} of {group.size()} initialized")
+        barrier(group)
+
+        if args.rank == 0:
+            print("\nMLX Distributed Communication Benchmark")
+            print(
+                f"Backend: {args.backend} | Dtype: {args.dtype} | Warmup: {args.warmup} | Iterations: {args.iterations}"
+            )
+
+        run_bench(
+            "Unidirectional (rank 0 -> rank 1)",
+            bench_unidirectional,
+            group,
+            args.rank,
+            dtype,
+            element_size,
+            args.warmup,
+            args.iterations,
+        )
+        run_bench(
+            "Round-Trip (ping-pong)",
+            bench_rtt,
+            group,
+            args.rank,
+            dtype,
+            element_size,
+            args.warmup,
+            args.iterations,
+            bw_multiplier=2,
+        )
+        run_bench(
+            "All-Gather",
+            bench_all_gather,
+            group,
+            args.rank,
+            dtype,
+            element_size,
+            args.warmup,
+            args.iterations,
+        )
+
+        if args.rank == 0:
+            print("\nDone.")
+        else:
+            print("Rank 1 complete.")
+
+    return 0
+
+
+if __name__ == "__main__":
+    try:
+        sys.exit(main())
+    except KeyboardInterrupt:
+        print("\nInterrupted.")
+        sys.exit(1)

dashboard/src/lib/components/FamilyLogos.svelte

@@ -88,12 +88,6 @@
       d="M22.012 0h1.032v.927H24v.968h-.956V3.78h-1.032V1.896h-1.878v-.97h1.878V0zM2.6 12.371V1.87h.969v10.502h-.97zm10.423.66h10.95v.918h-6.208v9.579h-4.742V13.03zM5.629 3.333v12.356H0v4.51h10.386V8L20.859 8l-.003-4.668-15.227.001z"
     />
   </svg>
-{:else if family === "nemotron"}
-  <svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
-    <path
-      d="M8.948 8.798v-1.43a6.7 6.7 0 0 1 .424-.018c3.922-.124 6.493 3.374 6.493 3.374s-2.774 3.851-5.75 3.851c-.398 0-.787-.062-1.158-.185v-4.346c1.528.185 1.837.857 2.747 2.385l2.04-1.714s-1.492-1.952-4-1.952a6.016 6.016 0 0 0-.796.035m0-4.735v2.138l.424-.027c5.45-.185 9.01 4.47 9.01 4.47s-4.08 4.964-8.33 4.964c-.37 0-.733-.035-1.095-.097v1.325c.3.035.61.062.91.062 3.957 0 6.82-2.023 9.593-4.408.459.371 2.34 1.263 2.73 1.652-2.633 2.208-8.772 3.984-12.253 3.984-.335 0-.653-.018-.971-.053v1.864H24V4.063zm0 10.326v1.131c-3.657-.654-4.673-4.46-4.673-4.46s1.758-1.944 4.673-2.262v1.237H8.94c-1.528-.186-2.73 1.245-2.73 1.245s.68 2.412 2.739 3.11M2.456 10.9s2.164-3.197 6.5-3.533V6.201C4.153 6.59 0 10.653 0 10.653s2.35 6.802 8.948 7.42v-1.237c-4.84-.6-6.492-5.936-6.492-5.936z"
-    />
-  </svg>
 {:else}
   <svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
     <path

dashboard/src/lib/components/FamilySidebar.svelte

@@ -31,7 +31,6 @@
     kimi: "Kimi",
     flux: "FLUX",
     "qwen-image": "Qwen Img",
-    nemotron: "NVIDIA",
   };
 
   function getFamilyName(family: string): string {

dashboard/src/lib/components/ModelCard.svelte

@@ -16,9 +16,7 @@
       perNode?: Array<{
         nodeId: string;
         nodeName: string;
-        status: "completed" | "partial" | "pending" | "downloading";
-        percentage: number;
-        progress: DownloadProgress | null;
+        progress: DownloadProgress;
       }>;
     } | null;
     nodes?: Record<string, NodeInfo>;
@@ -147,7 +145,10 @@
     return `${s}s`;
   }
 
-  const perNode = $derived(downloadStatus?.perNode ?? []);
+  const isDownloading = $derived(downloadStatus?.isDownloading ?? false);
+  const progress = $derived(downloadStatus?.progress);
+  const percentage = $derived(progress?.percentage ?? 0);
+  let expandedNodes = $state<Set<string>>(new Set());
 
   function toggleNodeDetails(nodeId: string): void {
     const next = new Set(expandedNodes);
@@ -586,49 +587,23 @@
       </span>
     </div>
 
-    <!-- Download Status (per-node) -->
-    {#if perNode.length > 0}
+    <!-- Download Status -->
+    {#if isDownloading && progress}
       <div class="mb-2 space-y-1">
-        <div
-          class="text-[10px] font-mono text-white/20 tracking-widest uppercase"
-        >
-          Download progress
+        <div class="flex items-center justify-between text-xs font-mono">
+          <span class="text-blue-400 tracking-wider uppercase">Downloading</span
+          >
+          <span class="text-white/60"
+            >{percentage.toFixed(1)}% &middot; {formatSpeed(progress.speed)}
+            &middot; {formatEta(progress.etaMs)}</span
+          >
+        </div>
+        <div class="h-1 bg-exo-medium-gray/30 rounded overflow-hidden">
+          <div
+            class="h-full bg-blue-500/70 transition-all duration-300"
+            style="width: {percentage}%"
+          ></div>
         </div>
-        {#each perNode as node}
-          <div class="flex items-center gap-2 text-xs font-mono">
-            <span class="text-white/40 w-20 truncate" title={node.nodeId}
-              >{node.nodeName}</span
-            >
-            <div
-              class="flex-1 h-1 bg-exo-medium-gray/30 rounded overflow-hidden"
-            >
-              <div
-                class="h-full transition-all duration-300 {node.status ===
-                'downloading'
-                  ? 'bg-blue-500/70'
-                  : node.status === 'completed'
-                    ? 'bg-exo-yellow/40'
-                    : 'bg-white/20'}"
-                style="width: {node.percentage}%"
-              ></div>
-            </div>
-            <span
-              class="text-right {node.status === 'completed'
-                ? 'text-exo-yellow/60'
-                : node.status === 'downloading'
-                  ? 'text-blue-400/60'
-                  : 'text-white/30'}"
-            >
-              {#if node.status === "downloading" && node.progress}
-                {Math.round(node.percentage)}% {formatSpeed(
-                  node.progress.speed,
-                )}
-              {:else}
-                {node.percentage > 0 ? `${Math.round(node.percentage)}%` : "0%"}
-              {/if}
-            </span>
-          </div>
-        {/each}
       </div>
     {/if}
 
@@ -687,7 +662,15 @@
             {@const allConnections =
               isDebugMode && usedNodes.length > 1
                 ? (() => {
-                    const conns: Array = [];
+                    const conns: Array<{
+                      ip: string;
+                      iface: string | null;
+                      from: string;
+                      to: string;
+                      midX: number;
+                      midY: number;
+                      arrow: string;
+                    }> = [];
                     for (let i = 0; i < usedNodes.length; i++) {
                       for (let j = i + 1; j < usedNodes.length; j++) {
                         const n1 = usedNodes[i];
@@ -699,12 +682,7 @@
                           const toPos = nodePositions[c.to];
                           const arrow =
                             fromPos && toPos ? getArrow(fromPos, toPos) : "→";
-                          conns.push({
-                            ...c,
-                            midX,
-                            midY,
-                            arrow,
-                          });
+                          conns.push({ ...c, midX, midY, arrow });
                         }
                       }
                     }

dashboard/src/lib/components/ModelPickerModal.svelte

@@ -459,7 +459,6 @@
       "llama",
       "flux",
       "qwen-image",
-      "nemotron",
     ];
     return Array.from(families).sort((a, b) => {
       const aIdx = familyOrder.indexOf(a);

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

@@ -1793,14 +1793,6 @@ class AppStore {
             this.persistConversation(targetConversationId);
           }
         },
-        {
-          generation_stats: (data) => {
-            const stats = data as { generation_tps: number };
-            if (stats.generation_tps > 0) {
-              this.tps = stats.generation_tps;
-            }
-          },
-        },
       );
 
       // Final update
@@ -1998,14 +1990,6 @@ class AppStore {
             this.persistConversation(targetConversationId);
           }
         },
-        {
-          generation_stats: (data) => {
-            const stats = data as { generation_tps: number };
-            if (stats.generation_tps > 0) {
-              this.tps = stats.generation_tps;
-            }
-          },
-        },
       );
 
       // Final cleanup of the message (if conversation still exists)
@@ -2413,7 +2397,7 @@ class AppStore {
 
       let streamedContent = "";
       let streamedThinking = "";
-      let serverTpsReceived = false;
+
       interface ChatCompletionChunk {
         choices?: Array<{
           delta?: { content?: string; reasoning_content?: string };
@@ -2478,6 +2462,7 @@ class AppStore {
             tokenCount += 1;
             this.totalTokens = tokenCount;
 
+            // Update real-time TPS during streaming
             if (firstTokenTime !== null && tokenCount > 1) {
               const elapsed = performance.now() - firstTokenTime;
               this.tps = (tokenCount / elapsed) * 1000;
@@ -2528,24 +2513,16 @@ class AppStore {
               startedAt: this.prefillProgress?.startedAt ?? performance.now(),
             };
           },
-          generation_stats: (data) => {
-            const stats = data as { generation_tps: number };
-
-            if (stats.generation_tps > 0) {
-              this.tps = stats.generation_tps;
-              serverTpsReceived = true;
-            }
-          },
         },
       );
 
       // Clear prefill progress after stream ends
       this.prefillProgress = null;
 
-      // Use server-side TPS if available, otherwise fall back to client-side
-      if (!serverTpsReceived && firstTokenTime !== null && tokenCount > 1) {
+      // Calculate final TPS
+      if (firstTokenTime !== null && tokenCount > 1) {
         const totalGenerationTime = performance.now() - firstTokenTime;
-        this.tps = (tokenCount / totalGenerationTime) * 1000;
+        this.tps = (tokenCount / totalGenerationTime) * 1000; // tokens per second
       }
 
       // Final cleanup of the message (if conversation still exists)

dashboard/src/routes/+page.svelte

@@ -42,7 +42,6 @@
     setSelectedChatModel,
     selectedChatModel,
     sendMessage,
-    thinkingEnabled,
     generateImage,
     editImage,
     editingImage,
@@ -853,7 +852,7 @@
   ) {
     const model = selectedChatModel();
     if (!model) {
-      sendMessage(content, files, thinkingEnabled());
+      sendMessage(content, files, null);
       return;
     }
 
@@ -881,7 +880,7 @@
     }
 
     // Default: text chat
-    sendMessage(content, files, thinkingEnabled());
+    sendMessage(content, files, null);
   }
 
   let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline");
@@ -1536,44 +1535,34 @@
   }
 
   // Helper to get download status for a model (checks all downloads for matching model ID)
-  type NodeDownloadStatus = {
-    nodeId: string;
-    nodeName: string;
-    status: "completed" | "partial" | "pending" | "downloading";
-    percentage: number;
-    progress: DownloadProgress | null;
-  };
-
-  // Shared helper: collect per-node download status for a model across a set of nodes.
-  // Handles deduplication, entry parsing, and aggregation in one place.
-  function collectDownloadStatus(
-    modelId: string,
-    nodeIds?: string[],
-  ): {
+  function getModelDownloadStatus(modelId: string): {
     isDownloading: boolean;
     progress: DownloadProgress | null;
-    perNode: NodeDownloadStatus[];
-    failedError: string | null;
+    perNode: Array<{
+      nodeId: string;
+      nodeName: string;
+      progress: DownloadProgress;
+    }>;
   } {
-    const empty = {
-      isDownloading: false,
-      progress: null,
-      perNode: [] as NodeDownloadStatus[],
-      failedError: null,
-    };
-
     if (!downloadsData || Object.keys(downloadsData).length === 0) {
-      return empty;
+      return { isDownloading: false, progress: null, perNode: [] };
     }
 
-    // Deduplicate by nodeId — a node can have multiple entries for the same model
-    // (e.g. PipelineShardMetadata + TensorShardMetadata). Keep the last entry,
-    // which is the most recently applied event.
-    const perNodeMap = new Map<string, NodeDownloadStatus>();
+    let totalBytes = 0;
+    let downloadedBytes = 0;
+    let totalSpeed = 0;
+    let completedFiles = 0;
+    let totalFiles = 0;
+    let isDownloading = false;
+    const allFiles: DownloadProgress["files"] = [];
+    const perNode: Array<{
+      nodeId: string;
+      nodeName: string;
+      progress: DownloadProgress;
+    }> = [];
 
-    const nodeIdSet = nodeIds ? new Set(nodeIds) : null;
+    // Check all nodes for downloads matching this model
     for (const [nodeId, nodeDownloads] of Object.entries(downloadsData)) {
-      if (nodeIdSet && !nodeIdSet.has(nodeId)) continue;
       if (!Array.isArray(nodeDownloads)) continue;
 
       for (const downloadWrapped of nodeDownloads) {
@@ -1586,45 +1575,29 @@
         const downloadPayload = (downloadWrapped as Record<string, unknown>)[
           downloadKind
         ] as Record<string, unknown>;
-        if (!downloadPayload) continue;
-
-        const downloadModelId = extractModelIdFromDownload(downloadPayload);
-        if (!downloadModelId || downloadModelId !== modelId) continue;
-
-        // DownloadFailed — return with any data collected so far
-        if (downloadKind === "DownloadFailed") {
-          return {
-            isDownloading: false,
-            progress: null,
-            perNode: Array.from(perNodeMap.values()),
-            failedError:
-              (downloadPayload.errorMessage as string) ||
-              (downloadPayload.error_message as string) ||
-              "Download failed",
-          };
-        }
 
         if (
           downloadKind !== "DownloadOngoing" &&
-          downloadKind !== "DownloadPending" &&
-          downloadKind !== "DownloadCompleted"
+          downloadKind !== "DownloadPending"
         )
           continue;
+        if (!downloadPayload) continue;
 
-        const nodeName =
-          data?.nodes?.[nodeId]?.friendly_name ?? nodeId.slice(0, 8);
+        const downloadModelId = extractModelIdFromDownload(downloadPayload);
 
-        if (downloadKind === "DownloadCompleted") {
-          perNodeMap.set(nodeId, {
-            nodeId,
-            nodeName,
-            status: "completed",
-            percentage: 100,
-            progress: null,
-          });
-          continue;
+        // Match if the model ID contains or equals the requested model
+        // (handles cases like "mlx-community/Meta-Llama..." matching)
+        if (
+          !downloadModelId ||
+          !downloadModelId.includes(modelId.split("/").pop() || modelId)
+        ) {
+          // Try exact match or partial match
+          if (downloadModelId !== modelId) continue;
         }
 
+        // For DownloadPending with partial bytes (paused/resumed downloads),
+        // synthesize a progress object from the top-level downloaded/total fields
+        let progress: DownloadProgress | null;
         if (downloadKind === "DownloadPending") {
           const pendingDownloaded = getBytes(
             downloadPayload.downloaded ??
@@ -1637,67 +1610,44 @@
               downloadPayload.totalBytes,
           );
           if (pendingDownloaded <= 0 && pendingTotal <= 0) continue;
-          const pct =
-            pendingTotal > 0 ? (pendingDownloaded / pendingTotal) * 100 : 0;
-          perNodeMap.set(nodeId, {
-            nodeId,
-            nodeName,
-            status: pendingDownloaded > 0 ? "partial" : "pending",
-            percentage: pct,
-            progress: null,
-          });
-          continue;
+          isDownloading = true;
+          progress = {
+            totalBytes: pendingTotal,
+            downloadedBytes: pendingDownloaded,
+            speed: 0,
+            etaMs: 0,
+            percentage:
+              pendingTotal > 0 ? (pendingDownloaded / pendingTotal) * 100 : 0,
+            completedFiles: 0,
+            totalFiles: 0,
+            files: [],
+          };
+        } else {
+          isDownloading = true;
+          progress = parseDownloadProgress(downloadPayload);
         }
 
-        // DownloadOngoing
-        const progress = parseDownloadProgress(downloadPayload);
-        if (
-          !progress ||
-          (progress.downloadedBytes <= 0 && progress.totalBytes <= 0)
-        )
-          continue;
+        if (progress) {
+          // Sum all values across nodes - each node downloads independently
+          totalBytes += progress.totalBytes;
+          downloadedBytes += progress.downloadedBytes;
+          totalSpeed += progress.speed;
+          completedFiles += progress.completedFiles;
+          totalFiles += progress.totalFiles;
+          allFiles.push(...progress.files);
 
-        perNodeMap.set(nodeId, {
-          nodeId,
-          nodeName,
-          status: "downloading",
-          percentage: progress.percentage,
-          progress,
-        });
-      }
-    }
-
-    // Aggregate from deduplicated per-node entries
-    const perNode = Array.from(perNodeMap.values());
-    let totalBytes = 0;
-    let downloadedBytes = 0;
-    let totalSpeed = 0;
-    let completedFiles = 0;
-    let totalFiles = 0;
-    let isDownloading = false;
-    const allFiles: DownloadProgress["files"] = [];
-
-    for (const node of perNode) {
-      if (node.status === "downloading" && node.progress) {
-        isDownloading = true;
-        totalBytes += node.progress.totalBytes;
-        downloadedBytes += node.progress.downloadedBytes;
-        totalSpeed += node.progress.speed;
-        completedFiles += node.progress.completedFiles;
-        totalFiles += node.progress.totalFiles;
-        allFiles.push(...node.progress.files);
+          const nodeName =
+            data?.nodes?.[nodeId]?.friendly_name ?? nodeId.slice(0, 8);
+          perNode.push({ nodeId, nodeName, progress });
+        }
       }
     }
 
     if (!isDownloading) {
-      return {
-        isDownloading: false,
-        progress: null,
-        perNode,
-        failedError: null,
-      };
+      return { isDownloading: false, progress: null, perNode: [] };
     }
 
+    // ETA = total remaining bytes / total speed across all nodes
     const remainingBytes = totalBytes - downloadedBytes;
     const etaMs = totalSpeed > 0 ? (remainingBytes / totalSpeed) * 1000 : 0;
 
@@ -1714,21 +1664,9 @@
         files: allFiles,
       },
       perNode,
-      failedError: null,
     };
   }
 
-  function getModelDownloadStatus(
-    modelId: string,
-    nodeIds?: string[],
-  ): {
-    isDownloading: boolean;
-    progress: DownloadProgress | null;
-    perNode: NodeDownloadStatus[];
-  } {
-    return collectDownloadStatus(modelId, nodeIds);
-  }
-
   // Helper to get download status for an instance
   function getInstanceDownloadStatus(
     instanceId: string,
@@ -1739,9 +1677,26 @@
     errorMessage: string | null;
     progress: DownloadProgress | null;
     statusText: string;
-    perNode: NodeDownloadStatus[];
+    perNode: Array<{
+      nodeId: string;
+      nodeName: string;
+      progress: DownloadProgress;
+    }>;
   } {
-    // Unwrap the instance to get shard assignments
+    if (!downloadsData || Object.keys(downloadsData).length === 0) {
+      // No download data yet — defer to runner status instead of assuming RUNNING
+      const statusInfo = deriveInstanceStatus(instanceWrapped);
+      return {
+        isDownloading: false,
+        isFailed: false,
+        errorMessage: null,
+        progress: null,
+        statusText: statusInfo.statusText,
+        perNode: [],
+      };
+    }
+
+    // Unwrap the instance
     const [instanceTag, instance] = getTagged(instanceWrapped);
     if (!instance || typeof instance !== "object") {
       return {
@@ -1761,45 +1716,132 @@
         modelId?: string;
       };
     };
-    const instanceModelId = inst.shardAssignments?.modelId;
-
-    if (!instanceModelId) {
-      const statusInfo = deriveInstanceStatus(instanceWrapped);
-      return {
-        isDownloading: false,
-        isFailed: statusInfo.statusText === "FAILED",
-        errorMessage: null,
-        progress: null,
-        statusText: statusInfo.statusText,
-        perNode: [],
-      };
-    }
-
-    // Get node IDs assigned to this instance
     const nodeToRunner = inst.shardAssignments?.nodeToRunner || {};
     const runnerToShard = inst.shardAssignments?.runnerToShard || {};
+    const instanceModelId = inst.shardAssignments?.modelId;
+
+    // Build reverse mapping: runnerId -> nodeId
     const runnerToNode: Record<string, string> = {};
     for (const [nodeId, runnerId] of Object.entries(nodeToRunner)) {
       runnerToNode[runnerId] = nodeId;
     }
-    const instanceNodeIds = Object.keys(runnerToShard)
-      .map((runnerId) => runnerToNode[runnerId])
-      .filter(Boolean);
 
-    const result = collectDownloadStatus(instanceModelId, instanceNodeIds);
+    let totalBytes = 0;
+    let downloadedBytes = 0;
+    let totalSpeed = 0;
+    let completedFiles = 0;
+    let totalFiles = 0;
+    let isDownloading = false;
+    const allFiles: DownloadProgress["files"] = [];
+    const perNode: Array<{
+      nodeId: string;
+      nodeName: string;
+      progress: DownloadProgress;
+    }> = [];
 
-    if (result.failedError) {
-      return {
-        isDownloading: false,
-        isFailed: true,
-        errorMessage: result.failedError,
-        progress: null,
-        statusText: "FAILED",
-        perNode: [],
-      };
+    // Check downloads for nodes that are part of this instance
+    for (const runnerId of Object.keys(runnerToShard)) {
+      const nodeId = runnerToNode[runnerId];
+      if (!nodeId) continue;
+
+      const nodeDownloads = downloadsData[nodeId];
+      if (!Array.isArray(nodeDownloads)) continue;
+
+      for (const downloadWrapped of nodeDownloads) {
+        if (!downloadWrapped || typeof downloadWrapped !== "object") continue;
+
+        const keys = Object.keys(downloadWrapped as Record<string, unknown>);
+        if (keys.length !== 1) continue;
+
+        const downloadKind = keys[0];
+        const downloadPayload = (downloadWrapped as Record<string, unknown>)[
+          downloadKind
+        ] as Record<string, unknown>;
+
+        // Handle DownloadFailed - return immediately with error info
+        if (downloadKind === "DownloadFailed") {
+          const downloadModelId = extractModelIdFromDownload(downloadPayload);
+          if (
+            instanceModelId &&
+            downloadModelId &&
+            downloadModelId === instanceModelId
+          ) {
+            return {
+              isDownloading: false,
+              isFailed: true,
+              errorMessage:
+                (downloadPayload.errorMessage as string) || "Download failed",
+              progress: null,
+              statusText: "FAILED",
+              perNode: [],
+            };
+          }
+        }
+
+        if (
+          downloadKind !== "DownloadOngoing" &&
+          downloadKind !== "DownloadPending"
+        )
+          continue;
+        if (!downloadPayload) continue;
+
+        // Check if this download is for this instance's model
+        const downloadModelId = extractModelIdFromDownload(downloadPayload);
+        if (
+          instanceModelId &&
+          downloadModelId &&
+          downloadModelId === instanceModelId
+        ) {
+          // For DownloadPending with partial bytes, synthesize progress
+          let progress: DownloadProgress | null;
+          if (downloadKind === "DownloadPending") {
+            const pendingDownloaded = getBytes(
+              downloadPayload.downloaded ??
+                downloadPayload.downloaded_bytes ??
+                downloadPayload.downloadedBytes,
+            );
+            const pendingTotal = getBytes(
+              downloadPayload.total ??
+                downloadPayload.total_bytes ??
+                downloadPayload.totalBytes,
+            );
+            if (pendingDownloaded <= 0 && pendingTotal <= 0) continue;
+            isDownloading = true;
+            progress = {
+              totalBytes: pendingTotal,
+              downloadedBytes: pendingDownloaded,
+              speed: 0,
+              etaMs: 0,
+              percentage:
+                pendingTotal > 0 ? (pendingDownloaded / pendingTotal) * 100 : 0,
+              completedFiles: 0,
+              totalFiles: 0,
+              files: [],
+            };
+          } else {
+            isDownloading = true;
+            progress = parseDownloadProgress(downloadPayload);
+          }
+
+          if (progress) {
+            // Sum all values across nodes - each node downloads independently
+            totalBytes += progress.totalBytes;
+            downloadedBytes += progress.downloadedBytes;
+            totalSpeed += progress.speed;
+            completedFiles += progress.completedFiles;
+            totalFiles += progress.totalFiles;
+            allFiles.push(...progress.files);
+
+            const nodeName =
+              data?.nodes?.[nodeId]?.friendly_name ?? nodeId.slice(0, 8);
+            perNode.push({ nodeId, nodeName, progress });
+          }
+        }
+      }
     }
 
-    if (!result.isDownloading) {
+    if (!isDownloading) {
+      // Check runner status for other states
       const statusInfo = deriveInstanceStatus(instanceWrapped);
       return {
         isDownloading: false,
@@ -1807,17 +1849,30 @@
         errorMessage: null,
         progress: null,
         statusText: statusInfo.statusText,
-        perNode: result.perNode,
+        perNode: [],
       };
     }
 
+    // ETA = total remaining bytes / total speed across all nodes
+    const remainingBytes = totalBytes - downloadedBytes;
+    const etaMs = totalSpeed > 0 ? (remainingBytes / totalSpeed) * 1000 : 0;
+
     return {
       isDownloading: true,
       isFailed: false,
       errorMessage: null,
-      progress: result.progress,
+      progress: {
+        totalBytes,
+        downloadedBytes,
+        speed: totalSpeed,
+        etaMs,
+        percentage: totalBytes > 0 ? (downloadedBytes / totalBytes) * 100 : 0,
+        completedFiles,
+        totalFiles,
+        files: allFiles,
+      },
       statusText: "DOWNLOADING",
-      perNode: result.perNode,
+      perNode,
     };
   }
 
@@ -4575,7 +4630,7 @@
                 type="button"
                 onclick={() => {
                   completeOnboarding();
-                  sendMessage(chip, undefined, thinkingEnabled());
+                  sendMessage(chip);
                 }}
                 class="px-4 py-2 rounded-full border border-white/10 bg-white/5 text-sm text-white/60 hover:bg-white/10 hover:text-white/80 hover:border-white/20 transition-all duration-200 cursor-pointer"
               >
@@ -5314,10 +5369,10 @@
                             <div
                               class="mt-2 space-y-2 max-h-48 overflow-y-auto pr-1"
                             >
-                              {#each downloadInfo.perNode.filter((n) => n.status === "downloading" && n.progress) as nodeProg}
+                              {#each downloadInfo.perNode as nodeProg}
                                 {@const nodePercent = Math.min(
                                   100,
-                                  Math.max(0, nodeProg.percentage),
+                                  Math.max(0, nodeProg.progress.percentage),
                                 )}
                                 {@const isExpanded =
                                   instanceDownloadExpandedNodes.has(
@@ -5373,17 +5428,15 @@
                                     >
                                       <span
                                         >{formatBytes(
-                                          nodeProg.progress?.downloadedBytes ??
-                                            0,
+                                          nodeProg.progress.downloadedBytes,
                                         )} / {formatBytes(
-                                          nodeProg.progress?.totalBytes ?? 0,
+                                          nodeProg.progress.totalBytes,
                                         )}</span
                                       >
                                       <span
-                                        >{formatSpeed(
-                                          nodeProg.progress?.speed ?? 0,
-                                        )} • ETA {formatEta(
-                                          nodeProg.progress?.etaMs ?? 0,
+                                        >{formatSpeed(nodeProg.progress.speed)} •
+                                        ETA {formatEta(
+                                          nodeProg.progress.etaMs,
                                         )}</span
                                       >
                                     </div>
@@ -5391,14 +5444,14 @@
 
                                   {#if isExpanded}
                                     <div class="mt-2 space-y-1.5">
-                                      {#if nodeProg.progress?.files ?? [].length === 0}
+                                      {#if nodeProg.progress.files.length === 0}
                                         <div
                                           class="text-[11px] font-mono text-exo-light-gray/70"
                                         >
                                           No file details reported.
                                         </div>
                                       {:else}
-                                        {#each nodeProg.progress?.files ?? [] as f}
+                                        {#each nodeProg.progress.files as f}
                                           {@const filePercent = Math.min(
                                             100,
                                             Math.max(0, f.percentage ?? 0),
@@ -5874,15 +5927,12 @@
                 )}
                 {@const allPreviews = filteredPreviews()}
                 {#if selectedModel && allPreviews.length > 0}
+                  {@const downloadStatus = getModelDownloadStatus(
+                    selectedModel.id,
+                  )}
                   {@const tags = modelTags()[selectedModel.id] || []}
                   <div class="space-y-3">
                     {#each allPreviews as apiPreview, i}
-                      {@const downloadStatus = getModelDownloadStatus(
-                        selectedModel.id,
-                        apiPreview.memory_delta_by_node
-                          ? Object.keys(apiPreview.memory_delta_by_node)
-                          : undefined,
-                      )}
                       <div
                         role="group"
                         onmouseenter={() => {
@@ -6070,7 +6120,7 @@
                           onclick={() => {
                             chatLaunchState = "idle";
                             selectedChatCategory = null;
-                            sendMessage(prompt, undefined, thinkingEnabled());
+                            sendMessage(prompt);
                           }}
                           class="text-left px-3 py-2.5 text-xs text-exo-light-gray hover:text-white font-mono rounded-lg border border-exo-medium-gray/30 hover:border-exo-yellow/30 bg-exo-dark-gray/30 hover:bg-exo-dark-gray/60 transition-all duration-200 cursor-pointer"
                         >
@@ -6453,10 +6503,10 @@
                               <div
                                 class="mt-2 space-y-2 max-h-48 overflow-y-auto pr-1"
                               >
-                                {#each downloadInfo.perNode.filter((n) => n.status === "downloading" && n.progress) as nodeProg}
+                                {#each downloadInfo.perNode as nodeProg}
                                   {@const nodePercent = Math.min(
                                     100,
-                                    Math.max(0, nodeProg.percentage),
+                                    Math.max(0, nodeProg.progress.percentage),
                                   )}
                                   {@const isExpanded =
                                     instanceDownloadExpandedNodes.has(
@@ -6515,17 +6565,16 @@
                                       >
                                         <span
                                           >{formatBytes(
-                                            nodeProg.progress
-                                              ?.downloadedBytes ?? 0,
+                                            nodeProg.progress.downloadedBytes,
                                           )} / {formatBytes(
-                                            nodeProg.progress?.totalBytes ?? 0,
+                                            nodeProg.progress.totalBytes,
                                           )}</span
                                         >
                                         <span
                                           >{formatSpeed(
-                                            nodeProg.progress?.speed ?? 0,
+                                            nodeProg.progress.speed,
                                           )} • ETA {formatEta(
-                                            nodeProg.progress?.etaMs ?? 0,
+                                            nodeProg.progress.etaMs,
                                           )}</span
                                         >
                                       </div>
@@ -6533,14 +6582,14 @@
 
                                     {#if isExpanded}
                                       <div class="mt-2 space-y-1.5">
-                                        {#if nodeProg.progress?.files ?? [].length === 0}
+                                        {#if nodeProg.progress.files.length === 0}
                                           <div
                                             class="text-[11px] font-mono text-exo-light-gray/70"
                                           >
                                             No file details reported.
                                           </div>
                                         {:else}
-                                          {#each nodeProg.progress?.files ?? [] as f}
+                                          {#each nodeProg.progress.files as f}
                                             {@const filePercent = Math.min(
                                               100,
                                               Math.max(0, f.percentage ?? 0),

flake.nix

@@ -72,7 +72,7 @@
       ];
 
       perSystem =
-        { config, self', pkgs, lib, system, ... }:
+        { config, self', inputs', pkgs, lib, system, ... }:
         let
           # Use pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs)
           pkgsSwift = import inputs.nixpkgs-swift { inherit system; };
@@ -84,17 +84,6 @@
             config.allowUnfreePredicate = pkg: (pkg.pname or "") == "metal-toolchain";
             overlays = [
               (import ./nix/apple-sdk-overlay.nix)
-              (final: prev: {
-                macmon = prev.macmon.overrideAttrs (_: {
-                  version = "git";
-                  src = final.fetchFromGitHub {
-                    owner = "swiftraccoon";
-                    repo = "macmon";
-                    rev = "9154d234f763fbeffdcb4135d0bbbaf80609699b";
-                    hash = "sha256-CwhilKNbs5XL9/tF5DMwyPBlE/hpmjGNTuxQ36sM50M=";
-                  };
-                });
-              })
             ];
           };
           treefmt = {
@@ -123,7 +112,9 @@
             };
           };
 
-          packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin (
+          packages = {
+            babeld = pkgs.callPackage ./nix/babeld.nix { };
+          } // lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin (
             let
               uvLock = builtins.fromTOML (builtins.readFile ./uv.lock);
               mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx" && p.source ? git) uvLock.package);
@@ -168,9 +159,6 @@
                 just
                 jq
               ]
-              ++ lib.optionals stdenv.isLinux [
-                unixtools.ifconfig
-              ]
               ++ lib.optionals stdenv.isDarwin [
                 macmon
               ];

justfile

@@ -1,8 +1,5 @@
 export NIX_CONFIG := "extra-experimental-features = nix-command flakes"
 
-default: lint fmt
-all: lint fmt check
-
 fmt:
     treefmt || nix fmt
 
@@ -34,10 +31,6 @@ build-dashboard:
 package:
     uv run pyinstaller packaging/pyinstaller/exo.spec
 
-build-app: package
-    xcodebuild build -project app/EXO/EXO.xcodeproj -scheme EXO -configuration Debug -derivedDataPath app/EXO/build
-    @echo "\nBuild complete. Run with:\n  open {{justfile_directory()}}/app/EXO/build/Build/Products/Debug/EXO.app"
-
 clean:
     rm -rf **/__pycache__
     rm -rf target/

nix/babeld.nix

@@ -0,0 +1,25 @@
+{ stdenv
+, lib
+, fetchurl
+}:
+stdenv.mkDerivation rec {
+  pname = "babeld";
+  version = "1.13.1";
+
+  src = fetchurl {
+    url = "https://www.irif.fr/~jch/software/files/${pname}-${version}.tar.gz";
+    hash = "sha256-FfJNJtoMz8Bzq83vAwnygeRoTyqnESb4JlcsTIRejdk=";
+  };
+
+  outputs = [
+    "out"
+    "man"
+  ];
+
+  makeFlags = [
+    "PREFIX=${placeholder "out"}"
+    "ETCDIR=${placeholder "out"}/etc"
+  ]
+  ++ lib.optional stdenv.isDarwin "LDLIBS=''";
+}
+

packaging/pyinstaller/exo.spec

@@ -71,9 +71,7 @@ MACMON_PATH = shutil.which("macmon")
 if MACMON_PATH is None:
     raise SystemExit(
         "macmon binary not found in PATH. "
-        "Install the pinned fork used by exo via: "
-        "cargo install --git https://github.com/swiftraccoon/macmon "
-        "--rev 9154d234f763fbeffdcb4135d0bbbaf80609699b macmon --force"
+        "Install it via: brew install macmon"
     )
 
 BINARIES: list[tuple[str, str]] = [
@@ -122,3 +120,4 @@ coll = COLLECT(
     upx_exclude=[],
     name="exo",
 )
+

pyproject.toml

@@ -25,7 +25,8 @@ dependencies = [
   "openai-harmony>=0.0.8",
   "httpx>=0.28.1",
   "tomlkit>=0.14.0",
-  "mflux==0.17.2",
+  "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",
@@ -61,7 +62,7 @@ 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/rltakashige/mlx-lm", branch = "leo/fix-deepseek-v32-indexer" }
+mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/eval-left-padding-in-batched-rotation" }
 # 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 }

python/parts.nix

@@ -185,6 +185,7 @@
         exo-eval = mkBenchScript "exo-eval" (inputs.self + /bench/exo_eval.py);
         exo-eval-tool-calls = mkBenchScript "exo-eval-tool-calls" (inputs.self + /bench/eval_tool_calls.py);
         exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
+        exo-mlx-bandwidth-test = mkPythonScript "exo-mlx-bandwidth-test" (inputs.self + /bench/test_mlx_bandwidth.py);
       };
 
       checks = {

resources/inference_model_cards/mlx-community--DeepSeek-V3.2-4bit.toml

@@ -1,13 +0,0 @@
-model_id = "mlx-community/DeepSeek-V3.2-4bit"
-n_layers = 61
-hidden_size = 7168
-num_key_value_heads = 128
-supports_tensor = true
-tasks = ["TextGeneration"]
-family = "deepseek"
-quantization = "4bit"
-base_model = "DeepSeek V3.2"
-capabilities = ["text", "thinking", "thinking_toggle"]
-
-[storage_size]
-in_bytes = 378086226621

resources/inference_model_cards/mlx-community--DeepSeek-V3.2-8bit.toml

@@ -1,13 +0,0 @@
-model_id = "mlx-community/DeepSeek-V3.2-8bit"
-n_layers = 61
-hidden_size = 7168
-num_key_value_heads = 128
-supports_tensor = true
-tasks = ["TextGeneration"]
-family = "deepseek"
-quantization = "8bit"
-base_model = "DeepSeek V3.2"
-capabilities = ["text", "thinking", "thinking_toggle"]
-
-[storage_size]
-in_bytes = 755957120916

result

@@ -0,0 +1 @@
+/nix/store/41rd4g2p69a2106qscl3xvkzmh6j7nry-babblerd

rust/babblerd/Cargo.toml

@@ -0,0 +1,19 @@
+[package]
+name = "babblerd"
+version.workspace = true
+edition.workspace = true
+
+[dependencies]
+color-eyre = "0.6.5"
+futures-lite.workspace = true
+ipnet = "2.12.0"
+libc = "0.2.183"
+n0-watcher = "0.6.1"
+netwatch = "0.14.0"
+rand = "0.10.0"
+tokio = { workspace = true, features = ["full"] }
+tracing = "0.1.44"
+tracing-subscriber = { version = "0.3.22", features = ["env-filter"] }
+
+[lints]
+workspace = true

rust/babblerd/src/lib.rs

@@ -0,0 +1,548 @@
+pub use babel::{babel, handle_listener};
+pub use error::{BabbleError, Result};
+pub use if_watcher::{PREFIX, watch};
+pub mod error {
+    use std::io;
+
+    pub type Result<T> = core::result::Result<T, BabbleError>;
+    #[derive(Debug)]
+    pub enum BabbleError {
+        Io(io::Error),
+        Unspecified,
+        BabeldCrashed(Option<i32>),
+        FailedToSetIp,
+        Other(String),
+    }
+    impl std::fmt::Display for BabbleError {
+        // use the debug display for now
+        #[inline]
+        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+            write!(f, "{self:?}")
+        }
+    }
+    impl std::error::Error for BabbleError {}
+    impl From<io::Error> for BabbleError {
+        #[inline]
+        fn from(value: io::Error) -> Self {
+            Self::Io(value)
+        }
+    }
+}
+pub mod if_watcher {
+    #[cfg(target_os = "linux")]
+    use std::path::PathBuf;
+    use std::{
+        collections::HashSet,
+        net::{IpAddr, Ipv6Addr},
+    };
+
+    use futures_lite::StreamExt;
+    use ipnet::Ipv6Net;
+    use n0_watcher::Watcher;
+    use netwatch::interfaces::{Interface, IpNet};
+    use tokio::sync::mpsc;
+
+    use crate::{
+        BabbleError, Result,
+        babel::Babble,
+        ip_manager::{add_ip, remove_ip},
+    };
+
+    pub const PREFIX: Ipv6Net = Ipv6Net::new_assert(
+        Ipv6Addr::new(0xfde0, 0x20c6, 0x1fa7, 0xffff, 0, 0, 0, 0),
+        48,
+    );
+
+    trait IfaceExt {
+        fn has_link_local_v6(&self) -> bool;
+        fn is_real_interface(&self) -> bool;
+        fn will_babel(&self) -> bool;
+        fn get_v6_in(&self, prefix: Ipv6Net) -> Option<Ipv6Net>;
+    }
+    impl IfaceExt for Interface {
+        fn will_babel(&self) -> bool {
+            self.has_link_local_v6() && self.is_real_interface() && self.is_up()
+        }
+
+        fn has_link_local_v6(&self) -> bool {
+            let mut has = false;
+            for addr in self.addrs() {
+                let IpAddr::V6(a) = addr.addr() else {
+                    continue;
+                };
+                if a.is_unicast_link_local() {
+                    has = true;
+                    break;
+                }
+            }
+            has
+        }
+        fn is_real_interface(&self) -> bool {
+            // macos is weird. en0 & en1 are ethernet & wifi (varies which is which by device). en3+ is thunderbolt, but at some point becomes usb ethernet.
+            if self.name().strip_prefix("en").is_none()
+            //.and_then(|s| s.parse::<u8>().ok())
+            //.is_none_or(|_n| false)
+            {
+                return false;
+            }
+            #[cfg(target_os = "linux")]
+            {
+                if !PathBuf::from(format!("/sys/class/net/{}/device", self.name())).exists() {
+                    tracing::debug!(
+                        "skipping interface {} as it doesn't correspond to a physical link",
+                        self.name()
+                    );
+                    return false;
+                }
+                let dev_type_path = PathBuf::from(format!("/sys/class/net/{}/type", self.name()));
+                if !dev_type_path.exists() {
+                    tracing::debug!(
+                        "skipping interface {} with no type file at {:?}",
+                        self.name(),
+                        dev_type_path.to_str()
+                    );
+                    return false;
+                }
+                let Ok(dev_type) = std::fs::read_to_string(dev_type_path) else {
+                    return false;
+                };
+                if dev_type.trim() != "1" {
+                    tracing::debug!(
+                        "skipping interface {} with type {:?}",
+                        self.name(),
+                        dev_type
+                    );
+                    return false;
+                }
+            }
+            true
+        }
+        fn get_v6_in(&self, prefix: Ipv6Net) -> Option<Ipv6Net> {
+            for addr in self.addrs() {
+                if let IpNet::V6(v6) = addr
+                    && prefix.contains(&v6.addr())
+                {
+                    return Some(v6);
+                }
+            }
+            None
+        }
+    }
+
+    #[tracing::instrument(skip(send))]
+    pub async fn watch(my_range: Ipv6Net, send: mpsc::Sender<Babble>) -> Result<()> {
+        assert!(PREFIX.contains(&my_range));
+        let mut ready_ifaces = HashSet::new();
+        // 0 is reserved for the first loopback address
+        let mut iface_num: u16 = 1;
+
+        tracing::info!("starting interface monitor");
+        let mon = netwatch::netmon::Monitor::new()
+            .await
+            .map_err(|_| BabbleError::Unspecified)?;
+        let mut mon_stream = mon.interface_state().stream();
+
+        while let Some(s) = mon_stream.next().await {
+            for iface in s.interfaces.values() {
+                if !iface.is_real_interface() {
+                    continue;
+                }
+                for addr in iface.addrs() {
+                    if let IpNet::V6(v6) = addr
+                        && PREFIX.contains(&v6.addr())
+                        && !my_range.contains(&v6.addr())
+                    {
+                        tracing::info!("removing stale ip {v6} from {}", iface.name());
+                        // don't really care if this fails
+                        if let Err(e) = remove_ip(v6, iface).await {
+                            tracing::warn!(%e, "failed to remove ip");
+                        }
+                    }
+                }
+                if !iface.will_babel() {
+                    continue;
+                }
+                if iface.get_v6_in(my_range).is_none() {
+                    let addr = Ipv6Net::new_assert(
+                        Ipv6Addr::from_bits(my_range.addr().to_bits() | u128::from(iface_num)),
+                        128,
+                    );
+                    iface_num += 1;
+                    assert!(iface_num < u16::MAX, "Really? u16::MAX interfaces?");
+                    tracing::info!("adding new ip {addr} to {}", iface.name());
+                    add_ip(addr, iface).await?;
+                }
+
+                if ready_ifaces.insert(iface.name().to_owned()) {
+                    tracing::info!("telling babeld to watch {}", iface.name());
+                    let Ok(()) = send.send(Babble::AddIface(iface.name().to_owned())).await else {
+                        return Ok(());
+                    };
+                }
+            }
+        }
+        tracing::info!("stopping interface monitor");
+
+        Ok(())
+    }
+}
+pub mod babel {
+    #[tracing::instrument(skip_all)]
+    pub async fn handle_listener(sock: UnixStream, mut receiver: broadcast::Receiver<String>) {
+        tracing::info!("new socket conn");
+        let (reader, mut write) = sock.into_split();
+        let mut reader = BufReader::new(reader).lines();
+        loop {
+            tokio::select! {
+                read = reader.next_line() => {
+                    let Ok(Some(_)) = read else { break; };
+                }
+                recv = receiver.recv() => {
+                    let res = match recv {
+                        Err(broadcast::error::RecvError::Lagged(_)) => { tracing::warn!("receiver lagged"); continue; },
+                        Err(broadcast::error::RecvError::Closed) => { tracing::debug!("receiver closed, dropping connection"); break; },
+                        Ok(s) => write.write_all(format!("{s}\n").as_bytes()).await,
+                    };
+                    if let Err(e) = res { tracing::warn!(error=%e, "failed to write to socket"); break; }
+                }
+            };
+        }
+        tracing::info!("closing socket conn");
+        _ = write.shutdown().await;
+    }
+
+    #[cfg(target_os = "macos")]
+    const PRIVATE_SOCK_PATH: &str = "/var/run/babbler/private/babeld.sock";
+    #[cfg(target_os = "linux")]
+    const PRIVATE_SOCK_PATH: &str = "/run/babbler/private/babeld.sock";
+    #[cfg(target_os = "macos")]
+    const PRIVATE_DIR: &str = "/var/run/babbler/private";
+    #[cfg(target_os = "linux")]
+    const PRIVATE_DIR: &str = "/run/babbler/private";
+
+    use ipnet::Ipv6Net;
+    use std::fs::Permissions;
+    use std::io;
+    use std::os::unix::fs::PermissionsExt;
+    use tokio::time::Duration;
+
+    use futures_lite::FutureExt;
+    use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, Lines};
+    use tokio::net::UnixStream;
+    use tokio::net::unix::{OwnedReadHalf, OwnedWriteHalf};
+    use tokio::process::Command;
+    use tokio::sync::{broadcast, mpsc};
+
+    use crate::{BabbleError, Result};
+
+    #[derive(Debug)]
+    pub enum Babble {
+        AddIface(String),
+    }
+
+    pub struct BabeldProcess {
+        proc: tokio::process::Child,
+    }
+    impl Drop for BabeldProcess {
+        #[inline]
+        fn drop(&mut self) {
+            // emergency sigkill babeld process to prevent leakage
+            match self.proc.try_wait() {
+                Ok(None) => {}
+                Ok(Some(sc)) => {
+                    if !sc.success() {
+                        _ = self.proc.start_kill();
+                    }
+                }
+                _ => {
+                    _ = self.proc.start_kill();
+                }
+            }
+        }
+    }
+    impl BabeldProcess {
+        #[tracing::instrument]
+        async fn spawn(my_range: Ipv6Net, iface: String) -> Result<Self> {
+            tokio::fs::create_dir_all(PRIVATE_DIR).await?;
+            tokio::fs::set_permissions(PRIVATE_DIR, Permissions::from_mode(0o0700)).await?;
+            tracing::info!("spawning babeld socket in {PRIVATE_SOCK_PATH}");
+            let res = match Command::new("babeld")
+                .arg("-G")
+                .arg(PRIVATE_SOCK_PATH)
+                .arg("-I")
+                .arg(format!("{PRIVATE_DIR}/babeld.pid"))
+                .arg("-C")
+                .arg(format!("redistribute local ip {my_range}"))
+                .arg("-C")
+                .arg("redistribute local deny")
+                .arg(iface)
+                .spawn()
+            {
+                Ok(proc) => Ok(Self { proc }),
+                Err(e) => {
+                    tracing::warn!(error=%e, "failed to spawn babeld");
+                    Err(e.into())
+                }
+            };
+            tokio::time::sleep(Duration::from_millis(10)).await;
+            while !matches!(tokio::fs::try_exists(PRIVATE_SOCK_PATH).await, Ok(true)) {
+                tracing::info!("where is the sock");
+                tokio::time::sleep(Duration::from_secs(1)).await;
+            }
+            res
+        }
+
+        #[tracing::instrument(skip(read, write, send))]
+        async fn query(
+            read: &mut Lines<BufReader<OwnedReadHalf>>,
+            write: &mut OwnedWriteHalf,
+            send: &broadcast::Sender<String>,
+            cmd: &str,
+        ) -> io::Result<Option<bool>> {
+            write.write_all(cmd.as_bytes()).await?;
+            loop {
+                let Some(line) = read.next_line().await? else {
+                    tracing::warn!("babeld closed unexpectedly");
+                    return Ok(None);
+                };
+                tracing::info!("[babel] {:?}", line);
+                let ret = match line.as_str() {
+                    "ok" => Ok(Some(true)),
+                    "bad" => {
+                        tracing::warn!("malformed message sent to babeld");
+                        Ok(Some(false))
+                    }
+                    _ if line.starts_with("no") => {
+                        tracing::warn!("message rejected");
+                        Ok(Some(false))
+                    }
+                    _ => Ok(None),
+                };
+                let Ok(_) = send.send(line) else {
+                    return Ok(None);
+                };
+                if !matches!(ret, Ok(None)) {
+                    return ret;
+                }
+            }
+        }
+
+        #[tracing::instrument(skip_all)]
+        async fn supervise(
+            &self,
+            mut recv: mpsc::Receiver<Babble>,
+            send: broadcast::Sender<String>,
+        ) -> Result<()> {
+            std::fs::set_permissions(PRIVATE_SOCK_PATH, Permissions::from_mode(0o0600))?;
+            tracing::debug!("connecting to babeld.sock");
+            let (reader, mut writer) = UnixStream::connect(PRIVATE_SOCK_PATH).await?.into_split();
+            let mut babel_lines = BufReader::new(reader).lines();
+            while let Some(s) = babel_lines.next_line().await? {
+                tracing::debug!("[babeld] {}", s);
+                if s == "ok" {
+                    break;
+                }
+            }
+            tracing::info!("babeld ok");
+            /* TODO(evan): push rather than pull
+            if Self::query(&mut babel_lines, &mut writer, "monitor\n")
+                .await?
+                .is_none()
+            {
+                return Ok(());
+            };
+            */
+
+            let mut interval = tokio::time::interval(Duration::from_secs(5));
+            interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
+            loop {
+                tokio::select! {
+                    _ = interval.tick() => {
+                        if Self::query(&mut babel_lines, &mut writer, &send, "dump\n")
+                            .await?
+                            .is_none()
+                        {
+                            return Ok(());
+                        }
+                    },
+                    babble = recv.recv() => {
+                        tracing::debug!("[babble] {:?}", babble);
+                        let Some(babble) = babble else {
+                            break;
+                        };
+                        match babble {
+                            Babble::AddIface(iface) => {
+                                Self::query(&mut babel_lines, &mut writer, &send, format!("interface {iface}\n").as_ref()).await?;
+                            }
+                        }
+                    },
+                    line = babel_lines.next_line() => {
+                        let Ok(Some(line)) = line else {
+                            break;
+                        };
+                        tracing::debug!("[babeld] {}", line);
+                        let Ok(_) = send.send(line) else {
+                            break;
+                        };
+                    },
+                }
+            }
+            Ok(())
+        }
+
+        async fn shutdown(mut self) -> Result<()> {
+            let kill_res = if let Some(pid) = self.proc.id() {
+                let pid: libc::pid_t = pid.try_into().expect("pid overflow");
+                // SAFETY: pid >= 0, freshly checked.
+                let rc = unsafe { libc::kill(pid, libc::SIGINT) };
+                let rc_err = if rc != 0 && rc != libc::ESRCH {
+                    Err(io::Error::last_os_error().into())
+                } else {
+                    Ok(())
+                };
+                let exit_code = async { Some(self.proc.wait().await) }
+                    .or(async {
+                        tokio::time::sleep(Duration::from_secs(5)).await;
+                        None
+                    })
+                    .await;
+                match exit_code {
+                    Some(Ok(code)) => {
+                        if code.success() {
+                            rc_err
+                        } else {
+                            rc_err.and_then(|()| Err(BabbleError::BabeldCrashed(code.code())))
+                        }
+                    }
+                    Some(Err(e)) => Err(e.into()),
+                    None => {
+                        self.proc.kill().await?;
+                        rc_err.and(Err(BabbleError::BabeldCrashed(None)))
+                    }
+                }
+            } else {
+                Ok(())
+            };
+            let rem_res = match std::fs::remove_file(PRIVATE_SOCK_PATH) {
+                Ok(()) => Ok(()),
+                Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(()),
+                Err(e) => Err(e.into()),
+            };
+            kill_res.and(rem_res)
+        }
+    }
+
+    #[tracing::instrument(skip(send, recv))]
+    pub async fn babel(
+        my_range: Ipv6Net,
+        mut recv: mpsc::Receiver<Babble>,
+        send: broadcast::Sender<String>,
+    ) -> Result<()> {
+        let iface = loop {
+            match recv.recv().await {
+                Some(Babble::AddIface(iface)) => {
+                    break iface;
+                }
+                None => return Ok(()),
+            }
+        };
+
+        let babel = BabeldProcess::spawn(my_range, iface).await?;
+        let res1 = babel.supervise(recv, send).await;
+        let res2 = babel.shutdown().await;
+        res1.and(res2)
+    }
+}
+pub(crate) mod ip_manager {
+    pub use sys::add_ip;
+    pub use sys::remove_ip;
+
+    #[cfg(target_os = "linux")]
+    mod sys {
+        use ipnet::Ipv6Net;
+        use netwatch::interfaces::Interface;
+
+        use crate::{BabbleError, Result};
+        use tokio::process::Command;
+
+        #[tracing::instrument]
+        pub async fn add_ip(subnet: Ipv6Net, iface: &Interface) -> Result<()> {
+            let out = Command::new("ip")
+                .arg("addr")
+                .arg("add")
+                .arg(format!("{subnet}"))
+                .arg("dev")
+                .arg(iface.name())
+                .output()
+                .await?;
+            if out.status.success() {
+                Ok(())
+            } else {
+                Err(BabbleError::FailedToSetIp)
+            }
+        }
+
+        #[tracing::instrument]
+        pub async fn remove_ip(v6: Ipv6Net, iface: &Interface) -> Result<()> {
+            let out = Command::new("ip")
+                .arg("addr")
+                .arg("del")
+                .arg(format!("{v6}"))
+                .arg("dev")
+                .arg(iface.name())
+                .output()
+                .await?;
+            if out.status.success() {
+                Ok(())
+            } else {
+                let std_err = String::from_utf8_lossy(&out.stdout);
+                tracing::debug!(%std_err);
+                Err(BabbleError::FailedToSetIp)
+            }
+        }
+    }
+
+    #[cfg(target_os = "macos")]
+    mod sys {
+        use ipnet::Ipv6Net;
+        use netwatch::interfaces::Interface;
+
+        use crate::BabbleError;
+        use crate::Result;
+        use tokio::process::Command;
+
+        #[tracing::instrument]
+        pub async fn add_ip(subnet: Ipv6Net, iface: &Interface) -> Result<()> {
+            let out = Command::new("ifconfig")
+                .arg(iface.name())
+                .arg("inet6")
+                .arg(format!("{subnet}"))
+                .arg("add")
+                .output()
+                .await?;
+            if out.status.success() {
+                Ok(())
+            } else {
+                Err(BabbleError::FailedToSetIp)
+            }
+        }
+
+        #[tracing::instrument]
+        pub async fn remove_ip(v6: Ipv6Net, iface: &Interface) -> Result<()> {
+            let out = Command::new("ifconfig")
+                .arg(iface.name())
+                .arg("inet6")
+                .arg(format!("{v6}"))
+                .arg("delete")
+                .output()
+                .await?;
+            if out.status.success() {
+                Ok(())
+            } else {
+                let std_err = String::from_utf8_lossy(&out.stdout);
+                tracing::debug!(%std_err);
+                Err(BabbleError::FailedToSetIp)
+            }
+        }
+    }
+}

rust/babblerd/src/main.rs

@@ -0,0 +1,167 @@
+use std::{fs::Permissions, io, net::Ipv6Addr, os::unix::fs::PermissionsExt};
+
+use babblerd::babel::handle_listener;
+use color_eyre::eyre::{WrapErr, eyre};
+use ipnet::Ipv6Net;
+use tokio::{
+    net::UnixListener,
+    signal,
+    sync::{broadcast, mpsc},
+    task::{JoinHandle, JoinSet},
+};
+
+#[cfg(target_os = "macos")]
+const PUBLIC_DIR: &str = "/var/run/babbler";
+#[cfg(target_os = "linux")]
+const PUBLIC_DIR: &str = "/run/babbler";
+#[cfg(target_os = "macos")]
+const PUBLIC_SOCK_PATH: &str = "/var/run/babbler/babblerd.sock";
+#[cfg(target_os = "linux")]
+const PUBLIC_SOCK_PATH: &str = "/run/babbler/babblerd.sock";
+
+enum State {
+    Idle,
+    Active {
+        recv: broadcast::Receiver<String>,
+        babel: JoinHandle<babblerd::Result<()>>,
+        watcher: JoinHandle<babblerd::Result<()>>,
+        listeners: JoinSet<()>,
+    },
+}
+
+#[tokio::main]
+async fn main() -> color_eyre::Result<()> {
+    color_eyre::install()?;
+    tracing_subscriber::fmt()
+        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
+        .init();
+    // cleanup old data
+    match std::fs::remove_file(PUBLIC_SOCK_PATH) {
+        Err(e) if e.kind() != io::ErrorKind::NotFound => {
+            return Err(e.into());
+        }
+        Ok(()) => {
+            tracing::info!("cleaned up old file at {PUBLIC_SOCK_PATH}");
+        }
+        _ => {}
+    }
+    std::fs::create_dir_all(PUBLIC_DIR)?;
+    // make our directory world readable
+    if let Err(e) = std::fs::set_permissions(PUBLIC_DIR, Permissions::from_mode(0o0755)) {
+        if e.kind() == io::ErrorKind::PermissionDenied {
+            return Err(eyre!(
+                "Insufficient permissions to run daemon -- did you forget sudo?"
+            ));
+        }
+        return Err(e.into());
+    }
+
+    let res = inner_main().await;
+    _ = std::fs::remove_file(PUBLIC_SOCK_PATH);
+    res
+}
+
+async fn inner_main() -> color_eyre::Result<()> {
+    tracing::info!("creating socket at {PUBLIC_SOCK_PATH}");
+    let public_socket = UnixListener::bind(PUBLIC_SOCK_PATH)?;
+    // make our socket world accessible
+    std::fs::set_permissions(PUBLIC_SOCK_PATH, Permissions::from_mode(0o0666))?;
+    let mut babbler: State = State::Idle;
+
+    loop {
+        babbler = match babbler {
+            State::Idle => {
+                tracing::info!("waiting for connection");
+                tokio::select! {
+                    sig = signal::ctrl_c() => {
+                        sig?;
+                        break;
+                    }
+                    sock = public_socket.accept() => {
+                        let sock = sock?.0;
+                        tracing::info!("starting babeld");
+                        let (br_send, br_recv) = broadcast::channel(1024);
+                        let (mp_send, mp_recv) = mpsc::channel(32);
+                        // node id is a PREFIX/64, NODE_ID/48 and an INTERFACE_ID/16
+                        let ip_node_id = u128::from(rand::random::<u64>() & (u64::MAX>>16)) << 16;
+                        let my_range = Ipv6Net::new_assert(
+                            Ipv6Addr::from_bits(babblerd::PREFIX.addr().to_bits() | ip_node_id),
+                            112,
+                        );
+                        let babel = tokio::spawn(babblerd::babel(my_range, mp_recv, br_send));
+                        let watcher = tokio::spawn(babblerd::watch(my_range, mp_send));
+                        let mut listeners = JoinSet::new();
+                        listeners.spawn(handle_listener(sock, br_recv.resubscribe()));
+                        State::Active { recv: br_recv, babel, watcher, listeners }
+                    }
+                }
+            }
+            State::Active {
+                recv,
+                mut babel,
+                mut watcher,
+                mut listeners,
+            } => {
+                tokio::select! {
+                    sig = signal::ctrl_c() => {
+                        sig?;
+                        drop(recv);
+                        watcher.abort();
+                        if let Ok(e) = watcher.await {
+                            e.wrap_err("while ctrl-c")?;
+                        }
+                        babel.await?.wrap_err("while ctrl-c")?;
+                        while let Some(res) = listeners.join_next().await {
+                            res.wrap_err("while ctrl-c")?;
+                        }
+                        break;
+                    }
+                    next_join_result = listeners.join_next(), if !listeners.is_empty() => {
+                        next_join_result.expect("checked")?;
+                        tracing::info!("dropped a listener");
+                        if listeners.is_empty() {
+                            tracing::info!("stopping babeld");
+                            drop(recv);
+
+                            watcher.abort();
+                            if let Ok(e) = watcher.await {
+                                e.wrap_err("while closing listeners")?;
+                            }
+                            babel.await?.wrap_err("while closing listeners")?;
+
+                            State::Idle
+                        } else {
+                            State::Active { recv, babel, watcher, listeners }
+                        }
+                    }
+                    sock = public_socket.accept() => {
+                        listeners.spawn(handle_listener(sock?.0, recv.resubscribe()));
+                        State::Active { recv, babel, watcher, listeners }
+                    }
+                    res = &mut watcher => {
+                        res??;
+                        drop(recv);
+                        babel.await?.wrap_err("while closing watcher")?;
+                        while let Some(res2) = listeners.join_next().await {
+                            res2.wrap_err("while closing watcher")?;
+                        }
+                        State::Idle
+                    }
+                    res = &mut babel => {
+                        res??;
+                        drop(recv);
+                        watcher.abort();
+                        if let Ok(e) = watcher.await {
+                            e.wrap_err("while closing babeld")?;
+                        }
+                        while let Some(res2) = listeners.join_next().await {
+                            res2.wrap_err("while closing babeld")?;
+                        }
+                        State::Idle
+                    }
+                }
+            }
+        };
+    }
+    Ok(())
+}

rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi

@@ -42,7 +42,7 @@ class MessageTooLargeError(builtins.Exception):
 
 @typing.final
 class NetworkingHandle:
-    def __new__(cls, identity: Keypair, bootstrap_peers: typing.Sequence[builtins.str], listen_port: builtins.int) -> NetworkingHandle: ...
+    def __new__(cls, identity: Keypair) -> NetworkingHandle: ...
     async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
         r"""
         Subscribe to a `GossipSub` topic.

rust/exo_pyo3_bindings/src/networking.rs

@@ -180,12 +180,7 @@ impl PyNetworkingHandle {
     // ---- Lifecycle management methods ----
 
     #[new]
-    #[pyo3(signature = (identity, bootstrap_peers, listen_port))]
-    fn py_new(
-        identity: Bound<'_, PyKeypair>,
-        bootstrap_peers: Vec<String>,
-        listen_port: u16,
-    ) -> PyResult<Self> {
+    fn py_new(identity: Bound<'_, PyKeypair>) -> PyResult<Self> {
         // create communication channels
         let (to_swarm, from_client) = mpsc::channel(MPSC_CHANNEL_SIZE);
 
@@ -194,9 +189,7 @@ impl PyNetworkingHandle {
 
         // 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, bootstrap_peers, listen_port)
-            .pyerr()?
-            .into_stream();
+        let swarm = create_swarm(identity, from_client).pyerr()?.into_stream();
 
         Ok(Self {
             swarm: Arc::new(Mutex::new(swarm)),

rust/networking/examples/chatroom.rs

@@ -16,14 +16,9 @@ async fn main() {
     let (to_swarm, from_client) = mpsc::channel(20);
 
     // Configure swarm
-    let mut swarm = swarm::create_swarm(
-        identity::Keypair::generate_ed25519(),
-        from_client,
-        vec![],
-        0,
-    )
-    .expect("Swarm creation failed")
-    .into_stream();
+    let mut swarm = swarm::create_swarm(identity::Keypair::generate_ed25519(), from_client)
+        .expect("Swarm creation failed")
+        .into_stream();
 
     // Create a Gossipsub topic & subscribe
     let (tx, rx) = oneshot::channel();

rust/networking/src/discovery.rs

@@ -104,7 +104,6 @@ pub struct Behaviour {
     // state-tracking for managed behaviors & mDNS-discovered peers
     managed: managed::Behaviour,
     mdns_discovered: HashMap<PeerId, BTreeSet<Multiaddr>>,
-    bootstrap_peers: Vec<Multiaddr>,
 
     retry_delay: Delay, // retry interval
 
@@ -113,11 +112,10 @@ pub struct Behaviour {
 }
 
 impl Behaviour {
-    pub fn new(keypair: &identity::Keypair, bootstrap_peers: Vec<Multiaddr>) -> io::Result<Self> {
+    pub fn new(keypair: &identity::Keypair) -> io::Result<Self> {
         Ok(Self {
             managed: managed::Behaviour::new(keypair)?,
             mdns_discovered: HashMap::new(),
-            bootstrap_peers,
             retry_delay: Delay::new(RETRY_CONNECT_INTERVAL),
             pending_events: WakerDeque::new(),
         })
@@ -370,12 +368,6 @@ impl NetworkBehaviour for Behaviour {
                     self.dial(p, ma)
                 }
             }
-            // dial bootstrap peers (for environments where mDNS is unavailable)
-            for addr in &self.bootstrap_peers {
-                self.pending_events.push_back(ToSwarm::Dial {
-                    opts: DialOpts::unknown_peer_id().address(addr.clone()).build(),
-                })
-            }
             self.retry_delay.reset(RETRY_CONNECT_INTERVAL) // reset timeout
         }
 

rust/networking/src/swarm.rs

@@ -142,29 +142,19 @@ fn filter_swarm_event(event: SwarmEvent<BehaviourEvent>) -> Option<FromSwarm> {
     }
 }
 
-/// Create and configure a swarm.
-///
-/// - `listen_port`: TCP port to listen on. `0` lets the OS assign one.
-/// - `bootstrap_peers`: multiaddrs to dial for environments without mDNS.
+/// Create and configure a swarm which listens to all ports on OS
 pub fn create_swarm(
     keypair: identity::Keypair,
     from_client: mpsc::Receiver<ToSwarm>,
-    bootstrap_peers: Vec<String>,
-    listen_port: u16,
 ) -> alias::AnyResult<Swarm> {
-    let parsed_bootstrap_peers: Vec<libp2p::Multiaddr> = bootstrap_peers
-        .iter()
-        .filter(|s| !s.is_empty())
-        .filter_map(|s| s.parse().ok())
-        .collect();
-
     let mut swarm = SwarmBuilder::with_existing_identity(keypair)
         .with_tokio()
         .with_other_transport(tcp_transport)?
-        .with_behaviour(|keypair| Behaviour::new(keypair, parsed_bootstrap_peers))?
+        .with_behaviour(Behaviour::new)?
         .build();
 
-    swarm.listen_on(format!("/ip4/0.0.0.0/tcp/{listen_port}").parse()?)?;
+    // 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 })
 }
 
@@ -256,12 +246,9 @@ mod behaviour {
     }
 
     impl Behaviour {
-        pub fn new(
-            keypair: &identity::Keypair,
-            bootstrap_peers: Vec<libp2p::Multiaddr>,
-        ) -> alias::AnyResult<Self> {
+        pub fn new(keypair: &identity::Keypair) -> alias::AnyResult<Self> {
             Ok(Self {
-                discovery: discovery::Behaviour::new(keypair, bootstrap_peers)?,
+                discovery: discovery::Behaviour::new(keypair)?,
                 gossipsub: gossipsub_behaviour(keypair),
             })
         }

rust/networking/tests/bootstrap_peers.rs

@@ -1,107 +0,0 @@
-use futures_lite::StreamExt;
-use networking::swarm::{FromSwarm, create_swarm};
-use std::time::Duration;
-use tokio::sync::mpsc;
-use tokio::time::timeout;
-
-/// Helper: find a free TCP port.
-fn free_port() -> u16 {
-    let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
-    listener.local_addr().unwrap().port()
-}
-
-/// Two nodes connect via bootstrap peers — no mDNS needed.
-///
-/// Node A listens on a fixed port. Node B bootstraps to A's address.
-/// We verify that B emits `FromSwarm::Discovered` for A's peer ID.
-#[tokio::test]
-async fn two_nodes_connect_via_bootstrap_peers() {
-    let port_a = free_port();
-
-    // Node A: listens on a known port, no bootstrap peers
-    let keypair_a = libp2p::identity::Keypair::generate_ed25519();
-    let peer_id_a = keypair_a.public().to_peer_id();
-    let (_tx_a, rx_a) = mpsc::channel(16);
-    let swarm_a = create_swarm(keypair_a, rx_a, vec![], port_a).expect("create swarm A");
-    let mut stream_a = swarm_a.into_stream();
-
-    // Node B: bootstraps to A's address
-    let keypair_b = libp2p::identity::Keypair::generate_ed25519();
-    let (_tx_b, rx_b) = mpsc::channel(16);
-    let swarm_b = create_swarm(
-        keypair_b,
-        rx_b,
-        vec![format!("/ip4/127.0.0.1/tcp/{port_a}")],
-        0,
-    )
-    .expect("create swarm B");
-    let mut stream_b = swarm_b.into_stream();
-
-    // Wait for B to discover A (connection established)
-    let connected = timeout(Duration::from_secs(10), async {
-        loop {
-            tokio::select! {
-                Some(event) = stream_a.next() => {
-                    // A will also see B connect, but we check from B's perspective
-                    let _ = event;
-                }
-                Some(event) = stream_b.next() => {
-                    if let FromSwarm::Discovered { peer_id } = event {
-                        if peer_id == peer_id_a {
-                            return true;
-                        }
-                    }
-                }
-            }
-        }
-    })
-    .await;
-
-    assert!(
-        connected.is_ok() && connected.unwrap(),
-        "Node B should discover Node A via bootstrap peer"
-    );
-}
-
-/// Empty bootstrap peers should work (backward compatible).
-#[tokio::test]
-async fn create_swarm_with_empty_bootstrap_peers() {
-    let keypair = libp2p::identity::Keypair::generate_ed25519();
-    let (_tx, rx) = mpsc::channel(16);
-    let swarm = create_swarm(keypair, rx, vec![], 0);
-    assert!(
-        swarm.is_ok(),
-        "create_swarm with no bootstrap peers should succeed"
-    );
-}
-
-/// Invalid multiaddr strings are silently filtered out.
-#[tokio::test]
-async fn create_swarm_ignores_invalid_bootstrap_addrs() {
-    let keypair = libp2p::identity::Keypair::generate_ed25519();
-    let (_tx, rx) = mpsc::channel(16);
-    let swarm = create_swarm(
-        keypair,
-        rx,
-        vec![
-            "not-a-valid-multiaddr".to_string(),
-            "".to_string(),
-            "/ip4/10.0.0.1/tcp/30000".to_string(), // valid
-        ],
-        0,
-    );
-    assert!(
-        swarm.is_ok(),
-        "create_swarm should succeed even with invalid bootstrap addrs"
-    );
-}
-
-/// Fixed listen port works correctly.
-#[tokio::test]
-async fn create_swarm_with_fixed_port() {
-    let port = free_port();
-    let keypair = libp2p::identity::Keypair::generate_ed25519();
-    let (_tx, rx) = mpsc::channel(16);
-    let swarm = create_swarm(keypair, rx, vec![], port);
-    assert!(swarm.is_ok(), "create_swarm with fixed port should succeed");
-}

rust/parts.nix

@@ -1,7 +1,7 @@
 { inputs, ... }:
 {
   perSystem =
-    { inputs', pkgs, lib, ... }:
+    { inputs', self', pkgs, lib, ... }:
     let
       # Fenix nightly toolchain with all components
       rustToolchain = inputs'.fenix.packages.stable.withComponents [
@@ -79,7 +79,7 @@
       };
 
       config = {
-        packages = {
+        packages = rec {
           # Python bindings wheel via maturin
           exo_pyo3_bindings = craneLib.buildPackage (
             commonArgs
@@ -110,6 +110,19 @@
               '';
             }
           );
+          babblerd-unwrapped = craneLib.buildPackage (
+            commonArgs // {
+              inherit cargoArtifacts;
+              pname = "babblerd-unwrapped";
+            }
+          );
+          babblerd = pkgs.writeShellApplication {
+            name = "babblerd";
+            runtimeInputs = [ self'.packages.babeld ];
+            text = ''
+              exec ${babblerd-unwrapped}/bin/babblerd "$@"
+            '';
+          };
         };
 
         checks = {

src/exo/api/types/__init__.py

@@ -1,57 +0,0 @@
-from .api import AddCustomModelParams as AddCustomModelParams
-from .api import AdvancedImageParams as AdvancedImageParams
-from .api import BenchChatCompletionRequest as BenchChatCompletionRequest
-from .api import BenchChatCompletionResponse as BenchChatCompletionResponse
-from .api import BenchImageGenerationResponse as BenchImageGenerationResponse
-from .api import BenchImageGenerationTaskParams as BenchImageGenerationTaskParams
-from .api import CancelCommandResponse as CancelCommandResponse
-from .api import ChatCompletionChoice as ChatCompletionChoice
-from .api import ChatCompletionMessage as ChatCompletionMessage
-from .api import ChatCompletionMessageText as ChatCompletionMessageText
-from .api import ChatCompletionRequest as ChatCompletionRequest
-from .api import ChatCompletionResponse as ChatCompletionResponse
-from .api import CompletionTokensDetails as CompletionTokensDetails
-from .api import CreateInstanceParams as CreateInstanceParams
-from .api import CreateInstanceResponse as CreateInstanceResponse
-from .api import DeleteDownloadResponse as DeleteDownloadResponse
-from .api import DeleteInstanceResponse as DeleteInstanceResponse
-from .api import DeleteTracesRequest as DeleteTracesRequest
-from .api import DeleteTracesResponse as DeleteTracesResponse
-from .api import ErrorInfo as ErrorInfo
-from .api import ErrorResponse as ErrorResponse
-from .api import FinishReason as FinishReason
-from .api import GenerationStats as GenerationStats
-from .api import HuggingFaceSearchResult as HuggingFaceSearchResult
-from .api import ImageData as ImageData
-from .api import ImageEditsTaskParams as ImageEditsTaskParams
-from .api import ImageGenerationResponse as ImageGenerationResponse
-from .api import ImageGenerationStats as ImageGenerationStats
-from .api import ImageGenerationTaskParams as ImageGenerationTaskParams
-from .api import ImageListItem as ImageListItem
-from .api import ImageListResponse as ImageListResponse
-from .api import ImageSize as ImageSize
-from .api import Logprobs as Logprobs
-from .api import LogprobsContentItem as LogprobsContentItem
-from .api import ModelList as ModelList
-from .api import ModelListModel as ModelListModel
-from .api import NodePowerStats as NodePowerStats
-from .api import PlaceInstanceParams as PlaceInstanceParams
-from .api import PlacementPreview as PlacementPreview
-from .api import PlacementPreviewResponse as PlacementPreviewResponse
-from .api import PowerUsage as PowerUsage
-from .api import PromptTokensDetails as PromptTokensDetails
-from .api import StartDownloadParams as StartDownloadParams
-from .api import StartDownloadResponse as StartDownloadResponse
-from .api import StreamingChoiceResponse as StreamingChoiceResponse
-from .api import ToolCall as ToolCall
-from .api import ToolCallItem as ToolCallItem
-from .api import TopLogprobItem as TopLogprobItem
-from .api import TraceCategoryStats as TraceCategoryStats
-from .api import TraceEventResponse as TraceEventResponse
-from .api import TraceListItem as TraceListItem
-from .api import TraceListResponse as TraceListResponse
-from .api import TraceRankStats as TraceRankStats
-from .api import TraceResponse as TraceResponse
-from .api import TraceStatsResponse as TraceStatsResponse
-from .api import Usage as Usage
-from .api import normalize_image_size as normalize_image_size

src/exo/download/download_utils.py

@@ -744,30 +744,13 @@ async def download_shard(
         logger.debug(f"Downloading {shard.model_card.model_id=} with {allow_patterns=}")
 
     all_start_time = time.time()
-    try:
-        file_list = await fetch_file_list_with_cache(
-            shard.model_card.model_id,
-            revision,
-            recursive=True,
-            skip_internet=skip_internet,
-            on_connection_lost=on_connection_lost,
-        )
-    except FileNotFoundError:
-        not_started_progress = RepoDownloadProgress(
-            repo_id=str(shard.model_card.model_id),
-            repo_revision=revision,
-            shard=shard,
-            completed_files=0,
-            total_files=0,
-            downloaded=Memory.from_bytes(0),
-            downloaded_this_session=Memory.from_bytes(0),
-            total=Memory.from_bytes(0),
-            overall_speed=0.0,
-            overall_eta=timedelta(0),
-            status="not_started",
-            file_progress={},
-        )
-        return target_dir, not_started_progress
+    file_list = await fetch_file_list_with_cache(
+        shard.model_card.model_id,
+        revision,
+        recursive=True,
+        skip_internet=skip_internet,
+        on_connection_lost=on_connection_lost,
+    )
     filtered_file_list = list(
         filter_repo_objects(
             file_list, allow_patterns=allow_patterns, key=lambda x: x.path

src/exo/main.py

@@ -11,9 +11,9 @@ from loguru import logger
 from pydantic import PositiveInt
 
 import exo.routing.topics as topics
-from exo.api.main import API
 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
@@ -47,11 +47,7 @@ class Node:
         keypair = get_node_id_keypair()
         node_id = NodeId(keypair.to_node_id())
         session_id = SessionId(master_node_id=node_id, election_clock=0)
-        router = Router.create(
-            keypair,
-            bootstrap_peers=args.bootstrap_peers,
-            listen_port=args.libp2p_port,
-        )
+        router = Router.create(keypair)
         await router.register_topic(topics.GLOBAL_EVENTS)
         await router.register_topic(topics.LOCAL_EVENTS)
         await router.register_topic(topics.COMMANDS)
@@ -268,17 +264,12 @@ def main():
     mp.set_start_method("spawn", force=True)
     # TODO: Refactor the current verbosity system
     logger_setup(EXO_LOG, args.verbosity)
-    logger.info(f"{'=' * 40}")
-    logger.info(f"Starting EXO | pid={os.getpid()}")
-    logger.info(f"{'=' * 40}")
+    logger.info("Starting EXO")
     logger.info(f"EXO_LIBP2P_NAMESPACE: {os.getenv('EXO_LIBP2P_NAMESPACE')}")
 
     if args.offline:
         logger.info("Running in OFFLINE mode — no internet checks, local models only")
 
-    if args.bootstrap_peers:
-        logger.info(f"Bootstrap peers: {args.bootstrap_peers}")
-
     if args.no_batch:
         os.environ["EXO_NO_BATCH"] = "1"
         logger.info("Continuous batching disabled (--no-batch)")
@@ -315,8 +306,6 @@ class Args(CamelCaseModel):
     offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true"
     no_batch: bool = False
     fast_synch: bool | None = None  # None = auto, True = force on, False = force off
-    bootstrap_peers: list[str] = []
-    libp2p_port: int
 
     @classmethod
     def parse(cls) -> Self:
@@ -374,22 +363,6 @@ class Args(CamelCaseModel):
             action="store_true",
             help="Disable continuous batching, use sequential generation",
         )
-        parser.add_argument(
-            "--bootstrap-peers",
-            type=lambda s: [p for p in s.split(",") if p],
-            default=os.getenv("EXO_BOOTSTRAP_PEERS", "").split(",")
-            if os.getenv("EXO_BOOTSTRAP_PEERS")
-            else [],
-            dest="bootstrap_peers",
-            help="Comma-separated libp2p multiaddrs to dial on startup (env: EXO_BOOTSTRAP_PEERS)",
-        )
-        parser.add_argument(
-            "--libp2p-port",
-            type=int,
-            default=0,
-            dest="libp2p_port",
-            help="Fixed TCP port for libp2p to listen on (0 = OS-assigned).",
-        )
         fast_synch_group = parser.add_mutually_exclusive_group()
         fast_synch_group.add_argument(
             "--fast-synch",

src/exo/master/adapters/chat_completions.py

@@ -4,7 +4,7 @@ import time
 from collections.abc import AsyncGenerator
 from typing import Any
 
-from exo.api.types import (
+from exo.shared.types.api import (
     ChatCompletionChoice,
     ChatCompletionMessage,
     ChatCompletionMessageText,
@@ -202,8 +202,6 @@ async def generate_chat_stream(
                     usage=last_usage,
                 )
                 yield f"data: {tool_response.model_dump_json()}\n\n"
-                if chunk.stats is not None:
-                    yield f": generation_stats {chunk.stats.model_dump_json()}\n\n"
                 yield "data: [DONE]\n\n"
                 return
 
@@ -218,10 +216,7 @@ async def generate_chat_stream(
                 yield f"data: {chunk_response.model_dump_json()}\n\n"
 
                 if chunk.finish_reason is not None:
-                    if chunk.stats is not None:
-                        yield f": generation_stats {chunk.stats.model_dump_json()}\n\n"
                     yield "data: [DONE]\n\n"
-                    return
 
 
 async def collect_chat_response(

src/exo/master/adapters/claude.py

@@ -5,8 +5,14 @@ import re
 from collections.abc import AsyncGenerator
 from typing import Any
 
-from exo.api.types import FinishReason, Usage
-from exo.api.types.claude_api import (
+from exo.shared.types.api import FinishReason, Usage
+from exo.shared.types.chunks import (
+    ErrorChunk,
+    PrefillProgressChunk,
+    TokenChunk,
+    ToolCallChunk,
+)
+from exo.shared.types.claude_api import (
     ClaudeContentBlock,
     ClaudeContentBlockDeltaEvent,
     ClaudeContentBlockStartEvent,
@@ -29,12 +35,6 @@ from exo.api.types.claude_api import (
     ClaudeToolUseBlock,
     ClaudeUsage,
 )
-from exo.shared.types.chunks import (
-    ErrorChunk,
-    PrefillProgressChunk,
-    TokenChunk,
-    ToolCallChunk,
-)
 from exo.shared.types.common import CommandId
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 

src/exo/master/adapters/ollama.py

@@ -4,7 +4,14 @@ import json
 from collections.abc import AsyncGenerator
 from typing import Any
 
-from exo.api.types.ollama_api import (
+from exo.shared.types.chunks import (
+    ErrorChunk,
+    PrefillProgressChunk,
+    TokenChunk,
+    ToolCallChunk,
+)
+from exo.shared.types.common import CommandId
+from exo.shared.types.ollama_api import (
     OllamaChatRequest,
     OllamaChatResponse,
     OllamaDoneReason,
@@ -14,13 +21,6 @@ from exo.api.types.ollama_api import (
     OllamaToolCall,
     OllamaToolFunction,
 )
-from exo.shared.types.chunks import (
-    ErrorChunk,
-    PrefillProgressChunk,
-    TokenChunk,
-    ToolCallChunk,
-)
-from exo.shared.types.common import CommandId
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 
 

src/exo/master/adapters/responses.py

@@ -4,8 +4,15 @@ from collections.abc import AsyncGenerator
 from itertools import count
 from typing import Any
 
-from exo.api.types import Usage
-from exo.api.types.openai_responses import (
+from exo.shared.types.api import Usage
+from exo.shared.types.chunks import (
+    ErrorChunk,
+    PrefillProgressChunk,
+    TokenChunk,
+    ToolCallChunk,
+)
+from exo.shared.types.common import CommandId
+from exo.shared.types.openai_responses import (
     FunctionCallInputItem,
     ResponseCompletedEvent,
     ResponseContentPart,
@@ -35,13 +42,6 @@ from exo.api.types.openai_responses import (
     ResponseTextDoneEvent,
     ResponseUsage,
 )
-from exo.shared.types.chunks import (
-    ErrorChunk,
-    PrefillProgressChunk,
-    TokenChunk,
-    ToolCallChunk,
-)
-from exo.shared.types.common import CommandId
 from exo.shared.types.text_generation import (
     InputMessage,
     TextGenerationTaskParams,

src/exo/master/api.py

@@ -21,17 +21,17 @@ from hypercorn.config import Config
 from hypercorn.typing import ASGIFramework
 from loguru import logger
 
-from exo.api.adapters.chat_completions import (
+from exo.master.adapters.chat_completions import (
     chat_request_to_text_generation,
     collect_chat_response,
     generate_chat_stream,
 )
-from exo.api.adapters.claude import (
+from exo.master.adapters.claude import (
     claude_request_to_text_generation,
     collect_claude_response,
     generate_claude_stream,
 )
-from exo.api.adapters.ollama import (
+from exo.master.adapters.ollama import (
     collect_ollama_chat_response,
     collect_ollama_generate_response,
     generate_ollama_chat_stream,
@@ -39,12 +39,34 @@ from exo.api.adapters.ollama import (
     ollama_generate_request_to_text_generation,
     ollama_request_to_text_generation,
 )
-from exo.api.adapters.responses import (
+from exo.master.adapters.responses import (
     collect_responses_response,
     generate_responses_stream,
     responses_request_to_text_generation,
 )
-from exo.api.types import (
+from exo.master.event_log import DiskEventLog
+from exo.master.image_store import ImageStore
+from exo.master.placement import place_instance as get_instance_placements
+from exo.shared.apply import apply
+from exo.shared.constants import (
+    DASHBOARD_DIR,
+    EXO_CACHE_HOME,
+    EXO_EVENT_LOG_DIR,
+    EXO_IMAGE_CACHE_DIR,
+    EXO_MAX_CHUNK_SIZE,
+    EXO_TRACING_CACHE_DIR,
+)
+from exo.shared.election import ElectionMessage
+from exo.shared.logging import InterceptLogger
+from exo.shared.models.model_cards import (
+    ModelCard,
+    ModelId,
+    delete_custom_card,
+    get_model_cards,
+    is_custom_card,
+)
+from exo.shared.tracing import TraceEvent, compute_stats, export_trace, load_trace_file
+from exo.shared.types.api import (
     AddCustomModelParams,
     AdvancedImageParams,
     BenchChatCompletionRequest,
@@ -92,47 +114,6 @@ from exo.api.types import (
     TraceStatsResponse,
     normalize_image_size,
 )
-from exo.api.types.claude_api import (
-    ClaudeMessagesRequest,
-    ClaudeMessagesResponse,
-)
-from exo.api.types.ollama_api import (
-    OllamaChatRequest,
-    OllamaChatResponse,
-    OllamaGenerateRequest,
-    OllamaGenerateResponse,
-    OllamaModelDetails,
-    OllamaModelTag,
-    OllamaPsModel,
-    OllamaPsResponse,
-    OllamaShowRequest,
-    OllamaShowResponse,
-    OllamaTagsResponse,
-)
-from exo.api.types.openai_responses import (
-    ResponsesRequest,
-    ResponsesResponse,
-)
-from exo.master.image_store import ImageStore
-from exo.master.placement import place_instance as get_instance_placements
-from exo.shared.apply import apply
-from exo.shared.constants import (
-    DASHBOARD_DIR,
-    EXO_CACHE_HOME,
-    EXO_EVENT_LOG_DIR,
-    EXO_IMAGE_CACHE_DIR,
-    EXO_MAX_CHUNK_SIZE,
-    EXO_TRACING_CACHE_DIR,
-)
-from exo.shared.election import ElectionMessage
-from exo.shared.logging import InterceptLogger
-from exo.shared.models.model_cards import (
-    ModelCard,
-    ModelId,
-    get_card,
-    get_model_cards,
-)
-from exo.shared.tracing import TraceEvent, compute_stats, export_trace, load_trace_file
 from exo.shared.types.chunks import (
     ErrorChunk,
     ImageChunk,
@@ -141,11 +122,13 @@ from exo.shared.types.chunks import (
     TokenChunk,
     ToolCallChunk,
 )
+from exo.shared.types.claude_api import (
+    ClaudeMessagesRequest,
+    ClaudeMessagesResponse,
+)
 from exo.shared.types.commands import (
-    AddCustomModelCard,
     Command,
     CreateInstance,
-    DeleteCustomModelCard,
     DeleteDownload,
     DeleteInstance,
     DownloadCommand,
@@ -168,13 +151,29 @@ from exo.shared.types.events import (
     TracesMerged,
 )
 from exo.shared.types.memory import Memory
+from exo.shared.types.ollama_api import (
+    OllamaChatRequest,
+    OllamaChatResponse,
+    OllamaGenerateRequest,
+    OllamaGenerateResponse,
+    OllamaModelDetails,
+    OllamaModelTag,
+    OllamaPsModel,
+    OllamaPsResponse,
+    OllamaShowRequest,
+    OllamaShowResponse,
+    OllamaTagsResponse,
+)
+from exo.shared.types.openai_responses import (
+    ResponsesRequest,
+    ResponsesResponse,
+)
 from exo.shared.types.state import State
 from exo.shared.types.worker.downloads import DownloadCompleted
 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.disk_event_log import DiskEventLog
 from exo.utils.power_sampler import PowerSampler
 from exo.utils.task_group import TaskGroup
 
@@ -1559,7 +1558,7 @@ class API:
                     storage_size_megabytes=card.storage_size.in_mb,
                     supports_tensor=card.supports_tensor,
                     tasks=[task.value for task in card.tasks],
-                    is_custom=card.is_custom,
+                    is_custom=is_custom_card(card.model_id),
                     family=card.family,
                     quantization=card.quantization,
                     base_model=card.base_model,
@@ -1570,7 +1569,7 @@ class API:
         )
 
     async def add_custom_model(self, payload: AddCustomModelParams) -> ModelListModel:
-        """Fetch a model from HuggingFace and save as a custom model card, then sync across the cluster."""
+        """Fetch a model from HuggingFace and save as a custom model card."""
         try:
             card = await ModelCard.fetch_from_hf(payload.model_id)
         except Exception as exc:
@@ -1578,13 +1577,6 @@ class API:
                 status_code=400, detail=f"Failed to fetch model: {exc}"
             ) from exc
 
-        await self.command_sender.send(
-            ForwarderCommand(
-                origin=self._system_id,
-                command=AddCustomModelCard(model_card=card),
-            )
-        )
-
         return ModelListModel(
             id=card.model_id,
             hugging_face_id=card.model_id,
@@ -1598,18 +1590,10 @@ class API:
         )
 
     async def delete_custom_model(self, model_id: ModelId) -> JSONResponse:
-        """Delete a user-added custom model card and sync deletion across the cluster."""
-        card = get_card(model_id)
-        if card is None or not card.is_custom:
+        """Delete a user-added custom model card."""
+        deleted = await delete_custom_card(model_id)
+        if not deleted:
             raise HTTPException(status_code=404, detail="Custom model card not found")
-
-        await self.command_sender.send(
-            ForwarderCommand(
-                origin=self._system_id,
-                command=DeleteCustomModelCard(model_id=model_id),
-            )
-        )
-
         return JSONResponse(
             {"message": "Model card deleted", "model_id": str(model_id)}
         )

src/exo/master/main.py

@@ -3,6 +3,7 @@ from datetime import datetime, timedelta, timezone
 import anyio
 from loguru import logger
 
+from exo.master.event_log import DiskEventLog
 from exo.master.placement import (
     add_instance_to_placements,
     cancel_unnecessary_downloads,
@@ -13,9 +14,7 @@ from exo.master.placement import (
 from exo.shared.apply import apply
 from exo.shared.constants import EXO_EVENT_LOG_DIR, EXO_TRACING_ENABLED
 from exo.shared.types.commands import (
-    AddCustomModelCard,
     CreateInstance,
-    DeleteCustomModelCard,
     DeleteInstance,
     ForwarderCommand,
     ForwarderDownloadCommand,
@@ -31,8 +30,6 @@ from exo.shared.types.commands import (
 )
 from exo.shared.types.common import CommandId, NodeId, SessionId, SystemId
 from exo.shared.types.events import (
-    CustomModelCardAdded,
-    CustomModelCardDeleted,
     Event,
     GlobalForwarderEvent,
     IndexedEvent,
@@ -64,7 +61,6 @@ from exo.shared.types.tasks import (
 )
 from exo.shared.types.worker.instances import InstanceId
 from exo.utils.channels import Receiver, Sender
-from exo.utils.disk_event_log import DiskEventLog
 from exo.utils.event_buffer import MultiSourceBuffer
 from exo.utils.task_group import TaskGroup
 
@@ -298,7 +294,6 @@ class Master:
                                 self.state.instances,
                                 self.state.node_memory,
                                 self.state.node_network,
-                                download_status=self.state.downloads,
                             )
                             transition_events = get_transition_events(
                                 self.state.instances, placement, self.state.tasks
@@ -349,14 +344,6 @@ class Master:
                                     f"Finished command {command.finished_command_id} finished"
                                 )
 
-                        case AddCustomModelCard():
-                            generated_events.append(
-                                CustomModelCardAdded(model_card=command.model_card)
-                            )
-                        case DeleteCustomModelCard():
-                            generated_events.append(
-                                CustomModelCardDeleted(model_id=command.model_id)
-                            )
                         case RequestEventLog():
                             # We should just be able to send everything, since other buffers will ignore old messages
                             # rate limit to 1000 at a time

src/exo/master/placement.py

@@ -32,10 +32,7 @@ from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import (
-    DownloadCompleted,
-    DownloadFailed,
     DownloadOngoing,
-    DownloadPending,
     DownloadProgress,
 )
 from exo.shared.types.worker.instances import (
@@ -63,45 +60,6 @@ def add_instance_to_placements(
     return {**current_instances, command.instance.instance_id: command.instance}
 
 
-def _get_node_download_fraction(
-    node_id: NodeId,
-    model_id: ModelId,
-    download_status: Mapping[NodeId, Sequence[DownloadProgress]],
-) -> float:
-    """Return the download fraction (0.0–1.0) for a model on a given node."""
-    for progress in download_status.get(node_id, []):
-        if progress.shard_metadata.model_card.model_id != model_id:
-            continue
-        match progress:
-            case DownloadCompleted():
-                return 1.0
-            case DownloadOngoing():
-                total = progress.download_progress.total.in_bytes
-                return (
-                    progress.download_progress.downloaded.in_bytes / total
-                    if total > 0
-                    else 0.0
-                )
-            case DownloadPending():
-                total = progress.total.in_bytes
-                return progress.downloaded.in_bytes / total if total > 0 else 0.0
-            case DownloadFailed():
-                return 0.0
-    return 0.0
-
-
-def _cycle_download_score(
-    cycle: Cycle,
-    model_id: ModelId,
-    download_status: Mapping[NodeId, Sequence[DownloadProgress]],
-) -> float:
-    """Sum of download fractions across all nodes in a cycle."""
-    return sum(
-        _get_node_download_fraction(node_id, model_id, download_status)
-        for node_id in cycle
-    )
-
-
 def place_instance(
     command: PlaceInstance,
     topology: Topology,
@@ -109,7 +67,6 @@ def place_instance(
     node_memory: Mapping[NodeId, MemoryUsage],
     node_network: Mapping[NodeId, NodeNetworkInfo],
     required_nodes: set[NodeId] | None = None,
-    download_status: Mapping[NodeId, Sequence[DownloadProgress]] | None = None,
 ) -> dict[InstanceId, Instance]:
     cycles = topology.get_cycles()
     candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles))
@@ -173,21 +130,11 @@ def place_instance(
         if any(topology.node_is_leaf(node_id) for node_id in cycle)
     ]
 
-    resolved_download_status = download_status or {}
-    candidate_cycles = (
-        cycles_with_leaf_nodes if cycles_with_leaf_nodes != [] else smallest_cycles
-    )
-
     selected_cycle = max(
-        candidate_cycles,
-        key=lambda cycle: (
-            _cycle_download_score(
-                cycle, command.model_card.model_id, resolved_download_status
-            ),
-            sum(
-                (node_memory[node_id].ram_available for node_id in cycle),
-                start=Memory(),
-            ),
+        cycles_with_leaf_nodes if cycles_with_leaf_nodes != [] else smallest_cycles,
+        key=lambda cycle: sum(
+            (node_memory[node_id].ram_available for node_id in cycle),
+            start=Memory(),
         ),
     )
 

src/exo/master/tests/test_api_error_handling.py

@@ -4,11 +4,10 @@ from typing import Any
 from fastapi import FastAPI, HTTPException
 from fastapi.testclient import TestClient
 
-from exo.api.main import API
-
 
 def test_http_exception_handler_formats_openai_style() -> None:
     """Test that HTTPException is converted to OpenAI-style error format."""
+    from exo.master.api import API
 
     app = FastAPI()
 

src/exo/master/tests/test_cancel_command.py

@@ -5,12 +5,12 @@ from unittest.mock import AsyncMock, MagicMock
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
 
-from exo.api.main import API
 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)

src/exo/master/tests/test_claude_api.py

@@ -3,11 +3,11 @@
 import pydantic
 import pytest
 
-from exo.api.adapters.claude import (
+from exo.master.adapters.claude import (
     claude_request_to_text_generation,
     finish_reason_to_claude_stop_reason,
 )
-from exo.api.types.claude_api import (
+from exo.shared.types.claude_api import (
     ClaudeMessage,
     ClaudeMessagesRequest,
     ClaudeTextBlock,

src/exo/master/tests/test_claude_tool_use.py

@@ -4,12 +4,12 @@ import json
 from collections.abc import AsyncGenerator
 from typing import Any, cast
 
-from exo.api.adapters.claude import (
+from exo.master.adapters.claude import (
     ClaudeMessagesResponse,
     collect_claude_response,
     generate_claude_stream,
 )
-from exo.api.types import ToolCallItem
+from exo.shared.types.api import ToolCallItem
 from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
 from exo.shared.types.common import CommandId, ModelId
 

src/exo/master/tests/test_event_log.py

@@ -2,8 +2,8 @@ from pathlib import Path
 
 import pytest
 
+from exo.master.event_log import DiskEventLog
 from exo.shared.types.events import TestEvent
-from exo.utils.disk_event_log import DiskEventLog
 
 
 @pytest.fixture

src/exo/master/tests/test_openai_responses_api.py

@@ -7,11 +7,11 @@ The responses adapter converts it to TextGenerationTaskParams for the pipeline.
 import pydantic
 import pytest
 
-from exo.api.types.openai_responses import (
+from exo.shared.types.common import ModelId
+from exo.shared.types.openai_responses import (
     ResponseInputMessage,
     ResponsesRequest,
 )
-from exo.shared.types.common import ModelId
 
 
 class TestResponsesRequestValidation:

src/exo/master/tests/test_placement.py

@@ -25,12 +25,6 @@ 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.downloads import (
-    DownloadCompleted,
-    DownloadFailed,
-    DownloadOngoing,
-    DownloadProgressData,
-)
 from exo.shared.types.worker.instances import (
     Instance,
     InstanceId,
@@ -39,7 +33,7 @@ from exo.shared.types.worker.instances import (
     MlxRingInstance,
 )
 from exo.shared.types.worker.runners import ShardAssignments
-from exo.shared.types.worker.shards import PipelineShardMetadata, Sharding
+from exo.shared.types.worker.shards import Sharding
 
 
 @pytest.fixture
@@ -582,183 +576,3 @@ def test_get_transition_events_delete_instance_cancels_only_matching_tasks(
     assert cancel_events[0].task_status == TaskStatus.Cancelled
     assert len(delete_events) == 1
     assert delete_events[0].instance_id == instance_id_a
-
-
-def _make_shard_metadata(model_card: ModelCard) -> PipelineShardMetadata:
-    return PipelineShardMetadata(
-        model_card=model_card,
-        device_rank=0,
-        world_size=1,
-        start_layer=0,
-        end_layer=model_card.n_layers,
-        n_layers=model_card.n_layers,
-    )
-
-
-def test_placement_prefers_cycle_with_downloaded_model(
-    model_card: ModelCard,
-) -> None:
-    """When two cycles are otherwise equal, prefer the one with the model already downloaded."""
-    topology = Topology()
-
-    model_card.storage_size = Memory.from_bytes(500)
-
-    node_a = NodeId()
-    node_b = NodeId()
-
-    node_memory = {
-        node_a: create_node_memory(1000),
-        node_b: create_node_memory(1000),
-    }
-    node_network = {
-        node_a: create_node_network(),
-        node_b: create_node_network(),
-    }
-
-    topology.add_node(node_a)
-    topology.add_node(node_b)
-    # No connections between them — two single-node cycles
-
-    shard_meta = _make_shard_metadata(model_card)
-
-    # node_b has the model fully downloaded, node_a does not
-    download_status = {
-        node_b: [
-            DownloadCompleted(
-                node_id=node_b,
-                shard_metadata=shard_meta,
-                total=model_card.storage_size,
-            ),
-        ],
-    }
-
-    cic = place_instance_command(model_card)
-    placements = place_instance(
-        cic, topology, {}, node_memory, node_network, download_status=download_status
-    )
-
-    assert len(placements) == 1
-    instance = list(placements.values())[0]
-    assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
-    assert assigned_nodes == {node_b}
-
-
-def test_placement_prefers_cycle_with_higher_download_progress(
-    model_card: ModelCard,
-) -> None:
-    """When two cycles are otherwise equal, prefer the one with more download progress."""
-    topology = Topology()
-
-    model_card.storage_size = Memory.from_bytes(1000)
-
-    node_a = NodeId()
-    node_b = NodeId()
-
-    node_memory = {
-        node_a: create_node_memory(1000),
-        node_b: create_node_memory(1000),
-    }
-    node_network = {
-        node_a: create_node_network(),
-        node_b: create_node_network(),
-    }
-
-    topology.add_node(node_a)
-    topology.add_node(node_b)
-
-    shard_meta = _make_shard_metadata(model_card)
-
-    # node_a: 30% downloaded, node_b: 80% downloaded
-    download_status = {
-        node_a: [
-            DownloadOngoing(
-                node_id=node_a,
-                shard_metadata=shard_meta,
-                download_progress=DownloadProgressData(
-                    total=Memory.from_bytes(1000),
-                    downloaded=Memory.from_bytes(300),
-                    downloaded_this_session=Memory.from_bytes(300),
-                    completed_files=0,
-                    total_files=1,
-                    speed=0.0,
-                    eta_ms=0,
-                    files={},
-                ),
-            ),
-        ],
-        node_b: [
-            DownloadOngoing(
-                node_id=node_b,
-                shard_metadata=shard_meta,
-                download_progress=DownloadProgressData(
-                    total=Memory.from_bytes(1000),
-                    downloaded=Memory.from_bytes(800),
-                    downloaded_this_session=Memory.from_bytes(800),
-                    completed_files=0,
-                    total_files=1,
-                    speed=0.0,
-                    eta_ms=0,
-                    files={},
-                ),
-            ),
-        ],
-    }
-
-    cic = place_instance_command(model_card)
-    placements = place_instance(
-        cic, topology, {}, node_memory, node_network, download_status=download_status
-    )
-
-    assert len(placements) == 1
-    instance = list(placements.values())[0]
-    assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
-    assert assigned_nodes == {node_b}
-
-
-def test_placement_does_not_prefer_cycle_with_failed_download(
-    model_card: ModelCard,
-) -> None:
-    """A failed download should count as 0% — not preferred over a node with no download history."""
-    topology = Topology()
-
-    model_card.storage_size = Memory.from_bytes(500)
-
-    node_a = NodeId()
-    node_b = NodeId()
-
-    # node_a has slightly more RAM so it would win on the RAM tiebreaker
-    node_memory = {
-        node_a: create_node_memory(1001),
-        node_b: create_node_memory(1000),
-    }
-    node_network = {
-        node_a: create_node_network(),
-        node_b: create_node_network(),
-    }
-
-    topology.add_node(node_a)
-    topology.add_node(node_b)
-
-    shard_meta = _make_shard_metadata(model_card)
-
-    # node_b has a failed download — should not be preferred
-    download_status = {
-        node_b: [
-            DownloadFailed(
-                node_id=node_b,
-                shard_metadata=shard_meta,
-                error_message="connection reset",
-            ),
-        ],
-    }
-
-    cic = place_instance_command(model_card)
-    placements = place_instance(
-        cic, topology, {}, node_memory, node_network, download_status=download_status
-    )
-
-    assert len(placements) == 1
-    instance = list(placements.values())[0]
-    assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
-    # node_a should win on RAM tiebreaker since failed download scores 0.0
-    assert assigned_nodes == {node_a}

src/exo/routing/router.py

@@ -1,4 +1,3 @@
-from collections.abc import Sequence
 from copy import copy
 from itertools import count
 from math import inf
@@ -103,15 +102,8 @@ class TopicRouter[T: CamelCaseModel]:
 
 class Router:
     @classmethod
-    def create(
-        cls,
-        identity: Keypair,
-        bootstrap_peers: Sequence[str] = (),
-        listen_port: int = 0,
-    ) -> "Router":
-        return cls(
-            handle=NetworkingHandle(identity, list(bootstrap_peers), listen_port)
-        )
+    def create(cls, identity: Keypair) -> "Router":
+        return cls(handle=NetworkingHandle(identity))
 
     def __init__(self, handle: NetworkingHandle):
         self.topic_routers: dict[str, TopicRouter[CamelCaseModel]] = {}

src/exo/shared/apply.py

@@ -7,8 +7,6 @@ from loguru import logger
 from exo.shared.types.common import NodeId
 from exo.shared.types.events import (
     ChunkGenerated,
-    CustomModelCardAdded,
-    CustomModelCardDeleted,
     Event,
     IndexedEvent,
     InputChunkReceived,
@@ -67,8 +65,6 @@ def event_apply(event: Event, state: State) -> State:
             | InputChunkReceived()
             | TracesCollected()
             | TracesMerged()
-            | CustomModelCardAdded()
-            | CustomModelCardDeleted()
         ):  # Pass-through events that don't modify state
             return state
         case InstanceCreated():
@@ -119,13 +115,7 @@ def apply_node_download_progress(event: NodeDownloadProgress, state: State) -> S
 
     replaced = False
     for i, existing_dp in enumerate(current):
-        # TODO(ciaran): deduplicate by model_id for now. Will need to use
-        # shard_metadata again when pipeline and tensor downloads differ.
-        # For now this is fine
-        if (
-            existing_dp.shard_metadata.model_card.model_id
-            == dp.shard_metadata.model_card.model_id
-        ):
+        if existing_dp.shard_metadata == dp.shard_metadata:
             current[i] = dp
             replaced = True
             break

src/exo/shared/logging.py

@@ -66,7 +66,7 @@ def logger_setup(log_file: Path | None, verbosity: int = 0):
     else:
         logger.add(
             sys.__stderr__,  # type: ignore
-            format="[ {time:YYYY-MM-DD HH:mm:ss.SSS} | <level>{level: <8}</level> | {name}:{function}:{line} ] <level>{message}</level>",
+            format="[ {time:HH:mm:ss.SSS} | <level>{level: <8}</level> | {name}:{function}:{line} ] <level>{message}</level>",
             level="DEBUG",
             colorize=True,
             enqueue=True,
@@ -76,7 +76,7 @@ def logger_setup(log_file: Path | None, verbosity: int = 0):
         logger.add(
             log_file,
             format="[ {time:YYYY-MM-DD HH:mm:ss.SSS} | {level: <8} | {name}:{function}:{line} ] {message}",
-            level="DEBUG" if verbosity > 0 else "INFO",
+            level="INFO",
             colorize=False,
             enqueue=True,
             rotation=lambda _, __: next(rotate_once),

src/exo/shared/models/model_cards.py

@@ -30,42 +30,30 @@ from exo.utils.pydantic_ext import CamelCaseModel
 # kinda ugly...
 # TODO: load search path from config.toml
 _custom_cards_dir = Path(str(EXO_CUSTOM_MODEL_CARDS_DIR))
-_BUILTIN_CARD_DIRS = [
+CARD_SEARCH_PATH = [
     Path(RESOURCES_DIR) / "inference_model_cards",
     Path(RESOURCES_DIR) / "image_model_cards",
+    _custom_cards_dir,
 ]
 
 _card_cache: dict[ModelId, "ModelCard"] = {}
 
 
-async def _load_cards_from_dir(directory: Path, *, is_custom: bool) -> None:
-    """Load all TOML model cards from a directory into the cache."""
-    async for toml_file in directory.rglob("*.toml"):
-        try:
-            card = await ModelCard.load_from_path(toml_file)
-            if is_custom:
-                card = card.model_copy(update={"is_custom": True})
-            if card.model_id not in _card_cache:
-                _card_cache[card.model_id] = card
-        except (ValidationError, TOMLKitError):
-            pass
-
-
-async def _refresh_card_cache() -> None:
-    for path in _BUILTIN_CARD_DIRS:
-        await _load_cards_from_dir(path, is_custom=False)
-    await _load_cards_from_dir(_custom_cards_dir, is_custom=True)
+async def _refresh_card_cache():
+    for path in CARD_SEARCH_PATH:
+        async for toml_file in path.rglob("*.toml"):
+            try:
+                card = await ModelCard.load_from_path(toml_file)
+                if card.model_id not in _card_cache:
+                    _card_cache[card.model_id] = card
+            except (ValidationError, TOMLKitError):
+                pass
 
 
 def _is_image_card(card: "ModelCard") -> bool:
     return any(t in (ModelTask.TextToImage, ModelTask.ImageToImage) for t in card.tasks)
 
 
-def get_card(model_id: ModelId) -> "ModelCard | None":
-    """Look up a single model card from the cache by ID."""
-    return _card_cache.get(model_id)
-
-
 async def get_model_cards() -> list["ModelCard"]:
     if len(_card_cache) == 0:
         await _refresh_card_cache()
@@ -104,7 +92,6 @@ class ModelCard(CamelCaseModel):
     capabilities: list[str] = []
     uses_cfg: bool = False
     trust_remote_code: bool = True
-    is_custom: bool = False
 
     @field_validator("tasks", mode="before")
     @classmethod
@@ -113,7 +100,7 @@ class ModelCard(CamelCaseModel):
 
     async def save(self, path: Path) -> None:
         async with await open_file(path, "w") as f:
-            py = self.model_dump(exclude_none=True, exclude={"is_custom"})
+            py = self.model_dump(exclude_none=True)
             data = tomlkit.dumps(py)  # pyright: ignore[reportUnknownMemberType]
             await f.write(data)
 
@@ -135,24 +122,17 @@ class ModelCard(CamelCaseModel):
         if (mc := _card_cache.get(model_id)) is not None:
             return mc
 
-        mc = await ModelCard.fetch_from_hf(model_id)
-        await mc.save_to_custom_dir()
-        _card_cache[model_id] = mc
-        return mc
+        return await ModelCard.fetch_from_hf(model_id)
 
     @staticmethod
     async def fetch_from_hf(model_id: ModelId) -> "ModelCard":
-        """Fetches storage size and number of layers for a Hugging Face model, returns Pydantic ModelMeta.
-
-        This is a pure fetch — it does NOT save to disk or update the cache.
-        Persistence is handled by the event-sourcing layer (worker event handler).
-        """
+        """Fetches storage size and number of layers for a Hugging Face model, returns Pydantic ModelMeta."""
         # TODO: failure if files do not exist
         config_data = await fetch_config_data(model_id)
         num_layers = config_data.layer_count
         mem_size_bytes = await fetch_safetensors_size(model_id)
 
-        return ModelCard(
+        mc = ModelCard(
             model_id=ModelId(model_id),
             storage_size=mem_size_bytes,
             n_layers=num_layers,
@@ -161,13 +141,10 @@ class ModelCard(CamelCaseModel):
             num_key_value_heads=config_data.num_key_value_heads,
             tasks=[ModelTask.TextGeneration],
             trust_remote_code=False,
-            is_custom=True,
         )
-
-
-def add_to_card_cache(card: "ModelCard") -> None:
-    """Add or update a model card in the in-memory cache."""
-    _card_cache[card.model_id] = card
+        await mc.save_to_custom_dir()
+        _card_cache[model_id] = mc
+        return mc
 
 
 async def delete_custom_card(model_id: ModelId) -> bool:
@@ -180,6 +157,16 @@ async def delete_custom_card(model_id: ModelId) -> bool:
     return False
 
 
+def is_custom_card(model_id: ModelId) -> bool:
+    """Check if a model card exists in the custom cards directory."""
+    import os
+
+    card_path = Path(str(EXO_CUSTOM_MODEL_CARDS_DIR)) / (
+        ModelId(model_id).normalize() + ".toml"
+    )
+    return os.path.isfile(str(card_path))
+
+
 class ConfigData(BaseModel):
     model_config = {"extra": "ignore"}  # Allow unknown fields
 

src/exo/shared/types/chunks.py

@@ -1,18 +1,18 @@
 from collections.abc import Generator
 from typing import Any, Literal
 
-from exo.api.types import (
-    FinishReason,
+from exo.shared.models.model_cards import ModelId
+from exo.shared.types.api import (
     GenerationStats,
     ImageGenerationStats,
-    ToolCallItem,
     TopLogprobItem,
     Usage,
 )
-from exo.shared.models.model_cards import ModelId
 from exo.utils.pydantic_ext import TaggedModel
 
+from .api import FinishReason
 from .common import CommandId
+from .worker.runner_response import ToolCallItem
 
 
 class BaseChunk(TaggedModel):

src/exo/shared/types/commands.py

@@ -1,10 +1,10 @@
 from pydantic import Field
 
-from exo.api.types import (
+from exo.shared.models.model_cards import ModelCard, ModelId
+from exo.shared.types.api import (
     ImageEditsTaskParams,
     ImageGenerationTaskParams,
 )
-from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.chunks import InputImageChunk
 from exo.shared.types.common import CommandId, NodeId, SystemId
 from exo.shared.types.text_generation import TextGenerationTaskParams
@@ -81,14 +81,6 @@ class CancelDownload(BaseCommand):
     model_id: ModelId
 
 
-class AddCustomModelCard(BaseCommand):
-    model_card: ModelCard
-
-
-class DeleteCustomModelCard(BaseCommand):
-    model_id: ModelId
-
-
 DownloadCommand = StartDownload | DeleteDownload | CancelDownload
 
 
@@ -104,8 +96,6 @@ Command = (
     | TaskCancelled
     | TaskFinished
     | SendInputChunk
-    | AddCustomModelCard
-    | DeleteCustomModelCard
 )
 
 

src/exo/shared/types/events.py

@@ -3,10 +3,9 @@ from typing import final
 
 from pydantic import Field
 
-from exo.shared.models.model_cards import ModelCard
 from exo.shared.topology import Connection
 from exo.shared.types.chunks import GenerationChunk, InputImageChunk
-from exo.shared.types.common import CommandId, Id, ModelId, NodeId, SessionId, SystemId
+from exo.shared.types.common import CommandId, Id, NodeId, SessionId, SystemId
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
@@ -107,14 +106,6 @@ class TopologyEdgeDeleted(BaseEvent):
     conn: Connection
 
 
-class CustomModelCardAdded(BaseEvent):
-    model_card: ModelCard
-
-
-class CustomModelCardDeleted(BaseEvent):
-    model_id: ModelId
-
-
 @final
 class TraceEventData(FrozenModel):
     name: str
@@ -156,8 +147,6 @@ Event = (
     | TopologyEdgeDeleted
     | TracesCollected
     | TracesMerged
-    | CustomModelCardAdded
-    | CustomModelCardDeleted
 )
 
 

src/exo/shared/types/tasks.py

@@ -2,7 +2,7 @@ from enum import Enum
 
 from pydantic import Field
 
-from exo.api.types import (
+from exo.shared.types.api import (
     ImageEditsTaskParams,
     ImageGenerationTaskParams,
 )

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

@@ -1,7 +1,7 @@
 from collections.abc import Generator
 from typing import Any, Literal
 
-from exo.api.types import (
+from exo.shared.types.api import (
     FinishReason,
     GenerationStats,
     ImageGenerationStats,

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

@@ -63,7 +63,7 @@ class RunnerShutdown(BaseRunnerStatus):
 
 
 class RunnerFailed(BaseRunnerStatus):
-    error_message: str
+    error_message: str | None = None
 
 
 RunnerStatus = (

src/exo/utils/info_gatherer/info_gatherer.py

@@ -10,6 +10,7 @@ from typing import Self, cast
 import anyio
 from anyio import fail_after, open_process, to_thread
 from anyio.streams.buffered import BufferedByteReceiveStream
+from anyio.streams.text import TextReceiveStream
 from loguru import logger
 from pydantic import ValidationError
 
@@ -287,7 +288,7 @@ class ThunderboltBridgeInfo(TaggedModel):
                 )
             )
         except Exception as e:
-            logger.opt(exception=e).warning("Failed to gather Thunderbolt Bridge info")
+            logger.warning(f"Failed to gather Thunderbolt Bridge info: {e}")
             return None
 
 
@@ -381,61 +382,18 @@ class InfoGatherer:
     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)
-    _psutil_memory_fallback_enabled: bool = field(init=False, default=False)
-
-    def _get_macmon_path(self) -> str | None:
-        return os.getenv("EXO_MACMON_PATH") or shutil.which("macmon")
-
-    async def _can_read_macmon_metrics(self, macmon_path: str) -> bool:
-        try:
-            with fail_after(5):
-                proc = await anyio.run_process(
-                    [macmon_path, "pipe", "--samples", "1", "--interval", "100"],
-                    check=False,
-                )
-        except Exception as e:
-            logger.opt(exception=e).warning(
-                f"Failed to validate macmon at {macmon_path}"
-            )
-            return False
-
-        if proc.returncode != 0:
-            stderr = proc.stderr.decode("utf-8", errors="replace").strip()
-            logger.warning(
-                f"macmon preflight failed with return code {proc.returncode}: "
-                f"{stderr or 'no stderr'}"
-            )
-            return False
-
-        stdout = proc.stdout.decode("utf-8", errors="replace").strip()
-        if not stdout:
-            logger.warning("macmon preflight returned no metrics")
-            return False
-
-        try:
-            MacmonMetrics.from_raw_json(stdout.splitlines()[0])
-        except ValidationError as e:
-            logger.opt(exception=e).warning(
-                "macmon preflight returned unexpected metrics JSON"
-            )
-            return False
-
-        return True
 
     async def run(self):
         async with self._tg as tg:
             if IS_DARWIN:
-                if (macmon_path := self._get_macmon_path()) is not None:
-                    if await self._can_read_macmon_metrics(macmon_path):
-                        tg.start_soon(self._monitor_macmon, macmon_path)
-                    else:
-                        logger.warning(
-                            f"macmon at {macmon_path} is unusable, falling back to psutil memory monitoring"
-                        )
+                if (macmon_path := shutil.which("macmon")) is not None:
+                    tg.start_soon(self._monitor_macmon, macmon_path)
                 else:
+                    # macmon not installed — fall back to psutil for memory
                     logger.warning(
                         "macmon not found, falling back to psutil for memory monitoring"
                     )
+                    self.memory_poll_rate = 1
                 tg.start_soon(self._monitor_system_profiler_thunderbolt_data)
                 tg.start_soon(self._monitor_thunderbolt_bridge_status)
                 tg.start_soon(self._monitor_rdma_ctl_status)
@@ -460,7 +418,7 @@ class InfoGatherer:
                 with fail_after(30):
                     await self.info_sender.send(await StaticNodeInformation.gather())
             except Exception as e:
-                logger.opt(exception=e).warning("Error gathering static node info")
+                logger.warning(f"Error gathering static node info: {e}")
             await anyio.sleep(self.static_info_poll_interval)
 
     async def _monitor_misc(self):
@@ -471,7 +429,7 @@ class InfoGatherer:
                 with fail_after(10):
                     await self.info_sender.send(await MiscData.gather())
             except Exception as e:
-                logger.opt(exception=e).warning("Error gathering misc data")
+                logger.warning(f"Error gathering misc data: {e}")
             await anyio.sleep(self.misc_poll_interval)
 
     async def _monitor_system_profiler_thunderbolt_data(self):
@@ -498,7 +456,7 @@ class InfoGatherer:
                     conns = [it for i in data if (it := i.conn()) is not None]
                     await self.info_sender.send(MacThunderboltConnections(conns=conns))
             except Exception as e:
-                logger.opt(exception=e).warning("Error gathering Thunderbolt data")
+                logger.warning(f"Error gathering Thunderbolt data: {e}")
             await anyio.sleep(self.system_profiler_interval)
 
     async def _monitor_memory_usage(self):
@@ -516,7 +474,7 @@ class InfoGatherer:
                     MemoryUsage.from_psutil(override_memory=override_memory)
                 )
             except Exception as e:
-                logger.opt(exception=e).warning("Error gathering memory usage")
+                logger.warning(f"Error gathering memory usage: {e}")
             await anyio.sleep(self.memory_poll_rate)
 
     async def _watch_system_info(self):
@@ -528,7 +486,7 @@ class InfoGatherer:
                     nics = await get_network_interfaces()
                     await self.info_sender.send(NodeNetworkInterfaces(ifaces=nics))
             except Exception as e:
-                logger.opt(exception=e).warning("Error gathering network interfaces")
+                logger.warning(f"Error gathering network interfaces: {e}")
             await anyio.sleep(self.interface_watcher_interval)
 
     async def _monitor_thunderbolt_bridge_status(self):
@@ -541,9 +499,7 @@ class InfoGatherer:
                     if curr is not None:
                         await self.info_sender.send(curr)
             except Exception as e:
-                logger.opt(exception=e).warning(
-                    "Error gathering Thunderbolt Bridge status"
-                )
+                logger.warning(f"Error gathering Thunderbolt Bridge status: {e}")
             await anyio.sleep(self.thunderbolt_bridge_poll_interval)
 
     async def _monitor_rdma_ctl_status(self):
@@ -555,7 +511,7 @@ class InfoGatherer:
                 if curr is not None:
                     await self.info_sender.send(curr)
             except Exception as e:
-                logger.opt(exception=e).warning("Error gathering RDMA ctl status")
+                logger.warning(f"Error gathering RDMA ctl status: {e}")
             await anyio.sleep(self.rdma_ctl_poll_interval)
 
     async def _monitor_disk_usage(self):
@@ -566,7 +522,7 @@ class InfoGatherer:
                 with fail_after(5):
                     await self.info_sender.send(await NodeDiskUsage.gather())
             except Exception as e:
-                logger.opt(exception=e).warning("Error gathering disk usage")
+                logger.warning(f"Error gathering disk usage: {e}")
             await anyio.sleep(self.disk_poll_interval)
 
     async def _monitor_macmon(self, macmon_path: str):
@@ -589,21 +545,15 @@ class InfoGatherer:
                     if not p.stdout:
                         logger.critical("MacMon closed stdout")
                         return
-                    stream = BufferedByteReceiveStream(p.stdout)
+                    stream = TextReceiveStream(BufferedByteReceiveStream(p.stdout))
                     while True:
                         with fail_after(read_timeout):
-                            data = await stream.receive_until(
-                                delimiter=b"\n", max_bytes=8 * 1024
-                            )
-                            text = data.decode("utf-8", errors="replace").strip()
-                            metrics = MacmonMetrics.from_raw_json(text)
-                        await self.info_sender.send(metrics)
+                            text = await stream.receive()
+                        await self.info_sender.send(MacmonMetrics.from_raw_json(text))
             except TimeoutError:
                 logger.warning(
                     f"MacMon produced no output for {read_timeout}s, restarting"
                 )
-                self.memory_poll_rate = 1
-                self._tg.start_soon(self._monitor_memory_usage)
             except CalledProcessError as e:
                 stderr_msg = "no stderr"
                 stderr_output = cast(bytes | str | None, e.stderr)
@@ -616,10 +566,6 @@ class InfoGatherer:
                 logger.warning(
                     f"MacMon failed with return code {e.returncode}: {stderr_msg}"
                 )
-                self.memory_poll_rate = 1
-                self._tg.start_soon(self._monitor_memory_usage)
             except Exception as e:
-                logger.opt(exception=e).warning("Error in macmon monitor")
-                self.memory_poll_rate = 1
-                self._tg.start_soon(self._monitor_memory_usage)
+                logger.warning(f"Error in macmon monitor: {e}")
             await anyio.sleep(self.macmon_interval)

src/exo/utils/power_sampler.py

@@ -5,7 +5,7 @@ from typing import final
 
 import anyio
 
-from exo.api.types import NodePowerStats, PowerUsage
+from exo.shared.types.api import NodePowerStats, PowerUsage
 from exo.shared.types.common import NodeId
 from exo.shared.types.profiling import SystemPerformanceProfile
 

src/exo/utils/tests/test_power_sampler.py

@@ -3,7 +3,7 @@ from collections.abc import Mapping
 import anyio
 import pytest
 
-from exo.api.types import PowerUsage
+from exo.shared.types.api import PowerUsage
 from exo.shared.types.common import NodeId
 from exo.shared.types.profiling import SystemPerformanceProfile
 from exo.utils.power_sampler import PowerSampler

src/exo/worker/engines/image/distributed_model.py

@@ -6,8 +6,8 @@ import mlx.core as mx
 from mflux.models.common.config.config import Config
 from PIL import Image
 
-from exo.api.types import AdvancedImageParams
 from exo.download.download_utils import build_model_path
+from exo.shared.types.api import AdvancedImageParams
 from exo.shared.types.worker.instances import BoundInstance
 from exo.shared.types.worker.shards import CfgShardMetadata, PipelineShardMetadata
 from exo.worker.engines.image.config import ImageModelConfig

src/exo/worker/engines/image/generate.py

@@ -9,7 +9,7 @@ from typing import Generator, Literal
 import mlx.core as mx
 from PIL import Image
 
-from exo.api.types import (
+from exo.shared.types.api import (
     AdvancedImageParams,
     ImageEditsTaskParams,
     ImageGenerationStats,

src/exo/worker/engines/image/models/qwen/wrappers.py

@@ -56,10 +56,10 @@ class QwenJointBlockWrapper(JointBlockWrapper[QwenTransformerBlock]):
         attn = self.block.attn
 
         img_mod_params = self.block.img_mod_linear(
-            self.block.img_mod_silu(text_embeddings)
+            self.block.img_mod_silu(text_embeddings)  # pyright: ignore[reportUnknownArgumentType]
         )
         txt_mod_params = self.block.txt_mod_linear(
-            self.block.txt_mod_silu(text_embeddings)
+            self.block.txt_mod_silu(text_embeddings)  # pyright: ignore[reportUnknownArgumentType]
         )
 
         img_mod1, img_mod2 = mx.split(img_mod_params, 2, axis=-1)

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

@@ -57,8 +57,8 @@ 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 exo.shared.logging import logger
 from exo.shared.types.worker.shards import PipelineShardMetadata
-from exo.worker.runner.bootstrap import logger
 
 if TYPE_CHECKING:
     from mlx_lm.models.cache import Cache
@@ -480,7 +480,7 @@ def patch_tensor_model[T](model: T) -> T:
             last = cache[-1]  # pyright: ignore[reportAny]
             dep_cache = last[0] if hasattr(last, "caches") else last  # pyright: ignore[reportAny]
             if hasattr(dep_cache, "keys"):  # type: ignore
-                dep_cache.keys = mx.depends(dep_cache.keys, logits)  # pyright: ignore[reportAny]
+                dep_cache.keys = mx.depends(dep_cache.keys, logits)  # pyright: ignore[reportAny,reportUnknownMemberType]
 
         return logits
 

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

@@ -4,7 +4,7 @@ from typing import Any
 
 from mlx_lm.chat_templates import deepseek_v32
 
-from exo.api.types import ToolCallItem
+from exo.shared.types.api import ToolCallItem
 
 BOS_TOKEN: str = deepseek_v32.bos_token
 EOS_TOKEN: str = deepseek_v32.eos_token
@@ -15,28 +15,7 @@ USER_TOKEN = "<\uff5cUser\uff5c>"
 ASSISTANT_TOKEN = "<\uff5cAssistant\uff5c>"
 TOOL_CALLS_START = f"<{DSML_TOKEN}function_calls>"
 TOOL_CALLS_END = f"</{DSML_TOKEN}function_calls>"
-_ORPHAN_THINK_END = ASSISTANT_TOKEN + THINKING_END
-_FIXED_THINK_BLOCK = ASSISTANT_TOKEN + THINKING_START + "\n" + THINKING_END
-
-
-def encode_messages(
-    messages: list[dict[str, Any]],
-    thinking_mode: str = "thinking",
-    context: list[dict[str, Any]] | None = None,
-    drop_thinking: bool = True,
-    add_default_bos_token: bool = True,
-    tools: Any = None,  # pyright: ignore[reportAny]
-) -> str:
-    prompt: str = deepseek_v32.encode_messages(
-        messages,
-        thinking_mode=thinking_mode,
-        context=context,
-        drop_thinking=drop_thinking,
-        add_default_bos_token=add_default_bos_token,
-        tools=tools,
-    )
-    return prompt.replace(_ORPHAN_THINK_END, _FIXED_THINK_BLOCK)
-
+encode_messages = deepseek_v32.encode_messages
 
 _INVOKE_PATTERN = re.compile(
     rf"<{re.escape(DSML_TOKEN)}invoke\s+name=\"([^\"]+)\">"

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

@@ -6,14 +6,11 @@ import mlx.core as mx
 from mlx_lm.generate import (
     BatchGenerator as MlxBatchGenerator,
 )
-from mlx_lm.generate import (
-    generation_stream,
-)
 from mlx_lm.models.cache import RotatingKVCache
 from mlx_lm.sample_utils import make_logits_processors, make_sampler
 from mlx_lm.tokenizer_utils import StreamingDetokenizer, TokenizerWrapper
 
-from exo.api.types import (
+from exo.shared.types.api import (
     CompletionTokensDetails,
     FinishReason,
     GenerationStats,
@@ -66,7 +63,6 @@ class _EngineTask:
     potential_stop_sequence_text: str = ""
     completion_tokens: int = 0
     generation_start_time: float = 0.0
-    generation_time_at_start: float = 0.0
     in_thinking: bool = False
     reasoning_tokens: int = 0
     prefill_tps: float = 0.0
@@ -79,23 +75,22 @@ class ExoBatchGenerator:
     group: mx.distributed.Group | None
     kv_prefix_cache: KVPrefixCache | None
 
-    _mlx_gen: MlxBatchGenerator = field(init=False)
+    _exo_gen: MlxBatchGenerator = field(init=False)
     _active_tasks: dict[int, _EngineTask] = field(default_factory=dict, init=False)
 
     def __post_init__(self) -> None:
-        self._mlx_gen = MlxBatchGenerator(
+        self._exo_gen = MlxBatchGenerator(
             model=self.model,
             stop_tokens=set(eos_ids_from_tokenizer(self.tokenizer)),
             prefill_step_size=4096,
         )
-        self._mlx_gen._needs_topk = False  # pyright: ignore[reportAttributeAccessIssue]
 
     @property
     def has_work(self) -> bool:
         return (
             bool(self._active_tasks)
-            or bool(self._mlx_gen.unprocessed_prompts)
-            or self._mlx_gen.active_batch is not None
+            or bool(self._exo_gen.unprocessed_prompts)
+            or self._exo_gen.active_batch is not None
         )
 
     def submit(
@@ -193,7 +188,7 @@ class ExoBatchGenerator:
 
         max_tokens = task_params.max_output_tokens or MAX_TOKENS
 
-        uids = self._mlx_gen.insert(
+        uids = self._exo_gen.insert(
             prompts=[last_tokens.tolist()],
             max_tokens=[max_tokens],
             caches=[list(cache)],
@@ -216,7 +211,6 @@ class ExoBatchGenerator:
             on_generation_token=on_generation_token,
             generation_start_time=time.perf_counter(),
             prefill_tps=_prefill_tps,
-            generation_time_at_start=self._mlx_gen._stats.generation_time,
         )
 
         return uid
@@ -225,12 +219,7 @@ class ExoBatchGenerator:
         if not self.has_work:
             return []
 
-        self._mlx_gen._needs_topk = any(  # pyright: ignore[reportAttributeAccessIssue]
-            t.task_params.logprobs for t in self._active_tasks.values()
-        )
-        _step_tic = time.perf_counter()
-        responses = self._mlx_gen.next()
-        _next_elapsed = time.perf_counter() - _step_tic
+        responses = self._exo_gen.next()
 
         results: list[tuple[int, GenerationResponse]] = []
 
@@ -288,31 +277,28 @@ class ExoBatchGenerator:
             logprob: float | None = None
             top_logprobs: list[TopLogprobItem] | None = None
             if task_params.logprobs:
-                with mx.stream(generation_stream):
-                    logprob, top_logprobs = extract_top_logprobs(
-                        logprobs=response.logprobs,
-                        tokenizer=self.tokenizer,
-                        top_logprobs=task_params.top_logprobs or DEFAULT_TOP_LOGPROBS,
-                        selected_token=response.token,
-                        precomputed_indices=getattr(response, "_topk_indices", None),
-                        precomputed_values=getattr(response, "_topk_values", None),
-                        precomputed_selected=getattr(
-                            response, "_selected_logprob", None
-                        ),
-                    )
+                logprob, top_logprobs = extract_top_logprobs(
+                    logprobs=response.logprobs,
+                    tokenizer=self.tokenizer,
+                    top_logprobs=task_params.top_logprobs or DEFAULT_TOP_LOGPROBS,
+                    selected_token=response.token,
+                )
 
             stats: GenerationStats | None = None
             usage: Usage | None = None
             if is_done:
-                gen_time_delta = (
-                    self._mlx_gen._stats.generation_time
-                    - state.generation_time_at_start
-                )
-                generation_tps = (
-                    state.completion_tokens / gen_time_delta
-                    if gen_time_delta > 0
-                    else 0.0
-                )
+                try:
+                    mlx_stats = self._exo_gen.stats()
+                    generation_tps = mlx_stats.generation_tps
+                except ZeroDivisionError:
+                    generation_elapsed = (
+                        time.perf_counter() - state.generation_start_time
+                    )
+                    generation_tps = (
+                        state.completion_tokens / generation_elapsed
+                        if generation_elapsed > 0
+                        else 0.0
+                    )
 
                 stats = GenerationStats(
                     prompt_tps=state.prefill_tps,
@@ -359,22 +345,15 @@ class ExoBatchGenerator:
                     -max_stop_len:
                 ]
 
-        _step_elapsed = time.perf_counter() - _step_tic
-        _overhead = _step_elapsed - _next_elapsed
-        if self._mlx_gen._next_count % 64 == 0 and responses:
-            logger.debug(
-                f"step overhead: {_overhead * 1000:.2f}ms (next={_next_elapsed * 1000:.2f}ms total={_step_elapsed * 1000:.2f}ms)"
-            )
-
         return results
 
     def cancel(self, uids: list[int]) -> None:
-        self._mlx_gen.remove(uids)
+        self._exo_gen.remove(uids)
         for uid in uids:
             self._active_tasks.pop(uid, None)
 
     def close(self) -> None:
-        self._mlx_gen.close()
+        self._exo_gen.close()
 
     def _save_prefix_cache(
         self,
@@ -393,8 +372,9 @@ class ExoBatchGenerator:
                 if len(all_prompt_tokens) > 0
                 else 0.0
             )
-            if matched_index is not None and (
-                prefix_hit_length > 1000 or hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE
+            if (
+                matched_index is not None
+                and hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE
             ):
                 self.kv_prefix_cache.update_kv_cache(
                     matched_index,

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

@@ -13,7 +13,7 @@ from mlx_lm.models.cache import ArraysCache, RotatingKVCache
 from mlx_lm.sample_utils import make_logits_processors, make_sampler
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
-from exo.api.types import (
+from exo.shared.types.api import (
     CompletionTokensDetails,
     FinishReason,
     GenerationStats,
@@ -179,8 +179,7 @@ def pipeline_parallel_prefill(
         flush_prefill_sends()
 
     assert _prompt_cache is not None
-    with mx.stream(generation_stream):
-        mx.eval([c.state for c in _prompt_cache])  # type: ignore
+    mx.eval([c.state for c in _prompt_cache])  # type: ignore
 
     # Final callback matching generate_step
     prompt_progress_callback(total, total)
@@ -313,46 +312,52 @@ def warmup_inference(
     model_id: ModelId,
 ) -> int:
     logger.info(f"warming up inference for instance: {model_id}")
+    t = time.monotonic()
 
     content = "Prompt to warm up the inference engine. Repeat this."
 
-    warmup_task_params = TextGenerationTaskParams(
-        model=model_id,
-        input=[InputMessage(role="user", content=content)],
-        max_output_tokens=50,
-        temperature=0.0,
-    )
-
     warmup_prompt = apply_chat_template(
         tokenizer=tokenizer,
-        task_params=warmup_task_params,
+        task_params=TextGenerationTaskParams(
+            model=ModelId(""),
+            input=[InputMessage(role="user", content=content)],
+        ),
     )
 
     tokens_generated = 0
 
+    cache = make_kv_cache(
+        model=model,
+    )
+
+    # Use a default sampler for warmup
+    sampler = make_sampler(temp=0.0)
+
     mx_barrier(group)
 
     logger.info("Generating warmup tokens")
-
-    t = time.monotonic()
-
-    for _r in mlx_generate(
+    for _r in stream_generate(
         model=model,
         tokenizer=tokenizer,
-        task=warmup_task_params,
         prompt=warmup_prompt,
-        kv_prefix_cache=None,
-        group=group,
+        max_tokens=50,
+        sampler=sampler,
+        prompt_cache=cache,
+        prefill_step_size=2048,
+        kv_group_size=KV_GROUP_SIZE,
+        kv_bits=KV_BITS,
     ):
+        logger.info("Generated warmup token: " + str(_r.text))
         tokens_generated += 1
 
-    check_for_cancel_every = min(
-        math.ceil(tokens_generated / min(time.monotonic() - t, 0.001)), 100
-    )
+    logger.info("Generated ALL warmup tokens")
 
     mx_barrier(group)
 
     logger.info(f"warmed up by generating {tokens_generated} tokens")
+    check_for_cancel_every = min(
+        math.ceil(tokens_generated / min(time.monotonic() - t, 0.001)), 100
+    )
     if group is not None:
         check_for_cancel_every = int(
             mx.max(
@@ -393,44 +398,52 @@ def extract_top_logprobs(
     tokenizer: TokenizerWrapper,
     top_logprobs: int,
     selected_token: int,
-    precomputed_indices: list[int] | None = None,
-    precomputed_values: list[float] | None = None,
-    precomputed_selected: float | None = None,
 ) -> tuple[float, list[TopLogprobItem]]:
-    if (
-        precomputed_indices is not None
-        and precomputed_values is not None
-        and precomputed_selected is not None
-    ):
-        top_indices_list: list[int] = precomputed_indices[:top_logprobs]
-        top_values_list: list[float] = precomputed_values[:top_logprobs]
-        selected_logprob = precomputed_selected
-    else:
-        selected_logprob_arr = logprobs[selected_token]
-        top_logprobs = min(top_logprobs, logprobs.shape[0] - 1)
-        top_indices = mx.argpartition(-logprobs, top_logprobs)[:top_logprobs]
-        top_values = logprobs[top_indices]
-        sort_order = mx.argsort(-top_values)
-        top_indices = top_indices[sort_order]
-        top_values = top_values[sort_order]
-        mx.eval(selected_logprob_arr, top_indices, top_values)
-        selected_logprob = float(selected_logprob_arr.item())
-        top_indices_list = top_indices.tolist()  # type: ignore
-        top_values_list = top_values.tolist()  # type: ignore
+    """Extract the selected token's logprob and top alternative tokens.
+
+    Args:
+        logprobs: Full vocabulary logprobs array from MLX
+        tokenizer: Tokenizer for decoding token IDs to strings
+        top_logprobs: Number of top alternatives to return
+        selected_token: The token ID that was actually sampled
+
+    Returns:
+        Tuple of (selected_token_logprob, list of TopLogprobItem for top alternatives)
+    """
+    # Get the logprob of the selected token
+    selected_logprob = float(logprobs[selected_token].item())
+
+    # Get top indices (most probable tokens)
+    # mx.argpartition gives indices that would partition the array
+    # We negate logprobs since argpartition finds smallest, and we want largest
+    top_logprobs = min(top_logprobs, logprobs.shape[0])  # Don't exceed vocab size
+    top_indices = mx.argpartition(-logprobs, top_logprobs)[:top_logprobs]
+
+    # Get the actual logprob values for these indices
+    top_values = logprobs[top_indices]
+
+    # Sort by logprob (descending) for consistent ordering
+    sort_order = mx.argsort(-top_values)
+    top_indices = top_indices[sort_order]
+    top_values = top_values[sort_order]
 
     # Convert to list of TopLogprobItem
     top_logprob_items: list[TopLogprobItem] = []
-    for token_id, token_logprob in zip(top_indices_list, top_values_list, strict=True):
+    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
+        token_bytes = list(token_str.encode("utf-8"))
         top_logprob_items.append(
             TopLogprobItem(
                 token=token_str,
                 logprob=token_logprob,
-                bytes=list(token_str.encode("utf-8")),
+                bytes=token_bytes,
             )
         )
 
@@ -611,13 +624,12 @@ def mlx_generate(
         logprob: float | None = None
         top_logprobs: list[TopLogprobItem] | None = None
         if task.logprobs:
-            with mx.stream(generation_stream):
-                logprob, top_logprobs = extract_top_logprobs(
-                    logprobs=out.logprobs,
-                    tokenizer=tokenizer,
-                    top_logprobs=task.top_logprobs or DEFAULT_TOP_LOGPROBS,
-                    selected_token=out.token,
-                )
+            logprob, top_logprobs = extract_top_logprobs(
+                logprobs=out.logprobs,
+                tokenizer=tokenizer,
+                top_logprobs=task.top_logprobs or DEFAULT_TOP_LOGPROBS,
+                selected_token=out.token,
+            )
 
         if is_done:
             # Log generation stats
@@ -645,9 +657,9 @@ def mlx_generate(
                     if len(all_prompt_tokens) > 0
                     else 0.0
                 )
-                if matched_index is not None and (
-                    prefix_hit_length > 1000
-                    or hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE
+                if (
+                    matched_index is not None
+                    and hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE
                 ):
                     kv_prefix_cache.update_kv_cache(
                         matched_index,

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

@@ -1,14 +0,0 @@
-from exo.worker.engines.mlx.patches.opt_batch_gen import apply_batch_gen_patch
-from exo.worker.engines.mlx.patches.standard_yarn_rope import patch_yarn_rope
-
-_applied = False
-
-
-def apply_mlx_patches() -> None:
-    global _applied
-    if _applied:
-        return
-    _applied = True
-    patch_yarn_rope()
-    # patch_gdn_softplus()
-    apply_batch_gen_patch()

src/exo/worker/engines/mlx/patches/opt_batch_gen.py

@@ -1,173 +0,0 @@
-import time
-from typing import Any, cast
-
-import mlx.core as mx
-from mlx_lm.generate import BatchGenerator, generation_stream
-
-_PRECOMPUTE_TOP_K = 20
-
-_original_public_next = BatchGenerator.next
-
-_pending_topk_idx: mx.array | None = None
-_pending_topk_val: mx.array | None = None
-_pending_selected_lps: mx.array | None = None
-
-
-def _fast_next(self: BatchGenerator) -> list[BatchGenerator.Response]:
-    tic = time.perf_counter()
-    batch = self.active_batch
-    assert batch is not None
-    batch_size = len(batch)
-
-    prev_tokens = batch.y
-    prev_logprobs = batch.logprobs
-
-    has_processors = any(p for ps in batch.logits_processors for p in ps)
-    if has_processors:
-        for i, toks in enumerate(batch.tokens):
-            batch.tokens[i] = mx.concatenate([toks, prev_tokens[i : i + 1]])
-
-    logits = self.model(prev_tokens[:, None], cache=batch.cache)
-    logits = logits[:, -1, :]
-
-    if has_processors:
-        processed_logits: list[mx.array] = []
-        for e in range(batch_size):
-            sample_logits: mx.array = logits[e : e + 1]
-            for processor in batch.logits_processors[e]:
-                sample_logits = processor(batch.tokens[e], sample_logits)
-            processed_logits.append(sample_logits)
-        logits = mx.concatenate(processed_logits, axis=0)
-
-    logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
-
-    if (
-        batch_size == 1
-        or any(batch.samplers)
-        and all(s is batch.samplers[0] for s in batch.samplers)
-    ):
-        sampler = batch.samplers[0] or self.sampler
-        batch.y = sampler(logprobs)
-    elif any(batch.samplers):
-        all_samples: list[mx.array] = []
-        for e in range(batch_size):
-            s = batch.samplers[e] or self.sampler
-            all_samples.append(s(logprobs[e : e + 1]))
-        batch.y = mx.concatenate(all_samples, axis=0)
-    else:
-        batch.y = self.sampler(logprobs)
-    batch.logprobs = list(logprobs)
-
-    global _pending_topk_idx, _pending_topk_val, _pending_selected_lps
-
-    emit_topk_indices: list[list[int]] = (
-        cast(list[list[int]], _pending_topk_idx.tolist())
-        if _pending_topk_idx is not None
-        else []
-    )
-    emit_topk_values: list[list[float]] = (
-        cast(list[list[float]], _pending_topk_val.tolist())
-        if _pending_topk_val is not None
-        else []
-    )
-    emit_selected_lps: list[float] = (
-        cast(list[float], _pending_selected_lps.tolist())
-        if _pending_selected_lps is not None
-        else []
-    )
-
-    needs_topk: bool = getattr(self, "_needs_topk", False)
-    if needs_topk:
-        k = min(_PRECOMPUTE_TOP_K, logprobs.shape[1])
-        _pending_topk_idx = mx.argpartition(-logprobs, k, axis=1)[:, :k]
-        _pending_topk_val = mx.take_along_axis(logprobs, _pending_topk_idx, axis=1)
-        sort_order = mx.argsort(-_pending_topk_val, axis=1)
-        _pending_topk_idx = mx.take_along_axis(_pending_topk_idx, sort_order, axis=1)
-        _pending_topk_val = mx.take_along_axis(_pending_topk_val, sort_order, axis=1)
-        _pending_selected_lps = logprobs[mx.arange(batch_size), batch.y]
-        mx.async_eval(
-            batch.y,
-            *batch.logprobs,
-            *batch.tokens,
-            _pending_topk_idx,
-            _pending_topk_val,
-            _pending_selected_lps,
-        )
-    else:
-        _pending_topk_idx = None
-        _pending_topk_val = None
-        _pending_selected_lps = None
-        mx.async_eval(batch.y, *batch.logprobs, *batch.tokens)
-
-    prev_token_list: list[int] = cast(list[int], prev_tokens.tolist())
-
-    toc = time.perf_counter()
-    self._stats.generation_time += toc - tic
-
-    keep_idx: list[int] = []
-    end_idx: list[int] = []
-    responses: list[Any] = []
-    stop_tokens = self.stop_tokens
-
-    for e in range(batch_size):
-        t = prev_token_list[e]
-        uid = batch.uids[e]
-        num_tok = batch.num_tokens[e] + 1
-        batch.num_tokens[e] = num_tok
-
-        if t in stop_tokens:
-            finish_reason = "stop"
-            end_idx.append(e)
-        elif num_tok >= batch.max_tokens[e]:
-            finish_reason = "length"
-            end_idx.append(e)
-        else:
-            finish_reason = None
-            keep_idx.append(e)
-
-        cache = None
-        if finish_reason is not None:
-            cache = batch.extract_cache(e)
-        response = self.Response(uid, t, prev_logprobs[e], finish_reason, cache)
-        if emit_topk_indices and e < len(emit_topk_indices):
-            response._topk_indices = emit_topk_indices[e]  # pyright: ignore[reportAttributeAccessIssue]
-            response._topk_values = emit_topk_values[e]  # pyright: ignore[reportAttributeAccessIssue]
-            response._selected_logprob = emit_selected_lps[e]  # pyright: ignore[reportAttributeAccessIssue]
-        responses.append(response)
-
-    if end_idx:
-        if keep_idx:
-            batch.filter(keep_idx)
-            if (
-                _pending_topk_idx is not None
-                and _pending_topk_val is not None
-                and _pending_selected_lps is not None
-            ):
-                ki = mx.array(keep_idx)
-                _pending_topk_idx = _pending_topk_idx[ki]
-                _pending_topk_val = _pending_topk_val[ki]
-                _pending_selected_lps = _pending_selected_lps[ki]
-        else:
-            self.active_batch = None
-            _pending_topk_idx = None
-            _pending_topk_val = None
-            _pending_selected_lps = None
-
-    self._next_count += 1
-    if self._next_count % 512 == 0:
-        mx.clear_cache()
-    self._stats.generation_tokens += len(responses)
-    return responses
-
-
-def _patched_public_next(self: BatchGenerator) -> list[BatchGenerator.Response]:
-    batch = self.active_batch
-    # Only do decode with fast_next
-    if batch is not None and not self.unprocessed_prompts:
-        with mx.stream(generation_stream):
-            return _fast_next(self)
-    return _original_public_next(self)
-
-
-def apply_batch_gen_patch() -> None:
-    BatchGenerator.next = _patched_public_next

src/exo/worker/engines/mlx/patches/standard_yarn_rope.py

@@ -1,118 +0,0 @@
-import math
-
-import mlx.core as mx
-from mlx_lm.models import rope_utils
-
-_original_YarnRoPE_init = rope_utils.YarnRoPE.__init__  # noqa: N816
-_original_initialize_rope = rope_utils.initialize_rope
-
-
-def _patched_yarn_init(
-    self: rope_utils.YarnRoPE,
-    dims: int,
-    traditional: bool = False,
-    max_position_embeddings: int = 2048,
-    base: float = 10000,
-    scaling_factor: float = 1.0,
-    original_max_position_embeddings: int = 4096,
-    beta_fast: float = 32,
-    beta_slow: float = 1,
-    mscale: float = 1,
-    mscale_all_dim: float = 0,
-    truncate: bool = True,
-) -> None:
-    """Patch mlx_lm's YarnRoPE to match vLLM's inverse-frequency blending formula for compatability."""
-
-    super(rope_utils.YarnRoPE, self).__init__()
-
-    def yarn_find_correction_dim(num_rotations: float) -> float:
-        return (
-            dims
-            * math.log(original_max_position_embeddings / (num_rotations * 2 * math.pi))
-        ) / (2 * math.log(base))
-
-    def yarn_find_correction_range() -> tuple[float, float]:
-        low: float = yarn_find_correction_dim(beta_fast)
-        high: float = yarn_find_correction_dim(beta_slow)
-        if truncate:
-            low = math.floor(low)
-            high = math.ceil(high)
-        return max(low, 0), min(high, dims - 1)
-
-    def yarn_get_mscale(scale: float = 1, ms: float = 1) -> float:
-        if scale <= 1:
-            return 1.0
-        return 0.1 * ms * math.log(scale) + 1.0
-
-    def yarn_linear_ramp_mask(min_val: float, max_val: float, dim: int) -> mx.array:
-        if min_val == max_val:
-            max_val += 0.001
-        linear_func = (mx.arange(dim, dtype=mx.float32) - min_val) / (max_val - min_val)
-        return mx.clip(linear_func, 0, 1)
-
-    self.mscale = yarn_get_mscale(scaling_factor, mscale) / yarn_get_mscale(
-        scaling_factor, mscale_all_dim
-    )
-    pos_freqs = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
-    inv_freq_extrapolation = 1.0 / pos_freqs
-    inv_freq_interpolation = 1.0 / (scaling_factor * pos_freqs)
-    low, high = yarn_find_correction_range()
-    inv_freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dims // 2)
-    inv_freq = (
-        inv_freq_interpolation * (1 - inv_freq_mask)
-        + inv_freq_extrapolation * inv_freq_mask
-    )
-    self._freqs = 1.0 / inv_freq
-    self.dims = dims
-    self.traditional = traditional
-
-
-def _patched_initialize_rope(
-    dims: int,
-    base: float,
-    traditional: bool,
-    scaling_config: dict[str, str | int | float | bool] | None = None,
-    max_position_embeddings: int | None = None,
-) -> object:
-    rope_type = "default"
-    if scaling_config is not None:
-        rope_type = str(
-            scaling_config.get("type") or scaling_config.get("rope_type", "default")
-        )
-
-    # All the yarn rope types supported in mlx lm
-    if rope_type in ("yarn", "deepseek_yarn"):
-        assert scaling_config is not None
-        cfg = scaling_config
-
-        def _float(key: str, default: float) -> float:
-            v = cfg.get(key)
-            return float(v) if v is not None else default
-
-        def _int(key: str, default: int) -> int:
-            v = cfg.get(key)
-            return int(v) if v is not None else default
-
-        return rope_utils.YarnRoPE(
-            dims=dims,
-            max_position_embeddings=max_position_embeddings or 2048,
-            traditional=traditional,
-            scaling_factor=_float("factor", 1.0),
-            base=base,
-            original_max_position_embeddings=_int(
-                "original_max_position_embeddings", 4096
-            ),
-            beta_fast=_float("beta_fast", 32),
-            beta_slow=_float("beta_slow", 1),
-            mscale=_float("mscale", 1),
-            mscale_all_dim=_float("mscale_all_dim", 0),
-        )
-
-    return _original_initialize_rope(
-        dims, base, traditional, scaling_config, max_position_embeddings
-    )
-
-
-def patch_yarn_rope() -> None:
-    rope_utils.YarnRoPE.__init__ = _patched_yarn_init
-    rope_utils.initialize_rope = _patched_initialize_rope

src/exo/worker/engines/mlx/tests/test_extract_top_logprobs.py

@@ -1,290 +0,0 @@
-# type: ignore
-import math
-from unittest.mock import MagicMock
-
-import mlx.core as mx
-import mlx.nn as nn
-import pytest
-from mlx_lm.generate import BatchGenerator
-
-from exo.worker.engines.mlx.generator.generate import extract_top_logprobs
-from exo.worker.engines.mlx.patches.opt_batch_gen import (
-    _PRECOMPUTE_TOP_K,
-    apply_batch_gen_patch,
-)
-
-
-def _mock_tokenizer() -> MagicMock:
-    tok = MagicMock()
-    tok.decode = lambda ids: f"tok_{ids[0]}"
-    return tok
-
-
-def _make_logprobs(values: list[float]) -> mx.array:
-    arr = mx.array(values, dtype=mx.float32)
-    mx.eval(arr)
-    return arr
-
-
-class TestExtractTopLogprobsFallback:
-    def test_returns_correct_selected_logprob(self) -> None:
-        lp = _make_logprobs([-1.0, -2.0, -0.5, -3.0, -4.0])
-        selected, _ = extract_top_logprobs(
-            lp, _mock_tokenizer(), top_logprobs=3, selected_token=2
-        )
-        assert selected == pytest.approx(-0.5)
-
-    def test_returns_top_k_sorted_descending(self) -> None:
-        lp = _make_logprobs([-1.0, -2.0, -0.5, -3.0, -4.0])
-        _, items = extract_top_logprobs(
-            lp, _mock_tokenizer(), top_logprobs=3, selected_token=0
-        )
-        logprob_values = [item.logprob for item in items]
-        assert logprob_values == sorted(logprob_values, reverse=True)
-        assert len(items) == 3
-
-    def test_top_tokens_are_most_probable(self) -> None:
-        lp = _make_logprobs([-5.0, -1.0, -3.0, -0.1, -2.0])
-        _, items = extract_top_logprobs(
-            lp, _mock_tokenizer(), top_logprobs=2, selected_token=0
-        )
-        token_ids = [int(item.token.split("_")[1]) for item in items]
-        assert 3 in token_ids
-        assert 1 in token_ids
-
-    def test_top_logprobs_clamped_to_vocab_size(self) -> None:
-        lp = _make_logprobs([-1.0, -2.0, -3.0, -4.0, -5.0])
-        _, items = extract_top_logprobs(
-            lp, _mock_tokenizer(), top_logprobs=10, selected_token=0
-        )
-        assert len(items) == 4
-
-    def test_nan_logprobs_filtered(self) -> None:
-        lp = _make_logprobs([-1.0, float("nan"), -0.5])
-        _, items = extract_top_logprobs(
-            lp, _mock_tokenizer(), top_logprobs=3, selected_token=0
-        )
-        for item in items:
-            assert not math.isnan(item.logprob)
-
-    def test_token_bytes_correct(self) -> None:
-        tok = MagicMock()
-        tok.decode = lambda ids: "hello"
-        lp = _make_logprobs([-1.0, -2.0])
-        _, items = extract_top_logprobs(lp, tok, top_logprobs=2, selected_token=0)
-        assert items[0].bytes == list("hello".encode("utf-8"))
-
-
-class TestExtractTopLogprobsPrecomputed:
-    def test_uses_precomputed_data(self) -> None:
-        lp = _make_logprobs([-99.0])
-        selected, items = extract_top_logprobs(
-            lp,
-            _mock_tokenizer(),
-            top_logprobs=2,
-            selected_token=0,
-            precomputed_indices=[3, 1, 0],
-            precomputed_values=[-0.1, -1.0, -5.0],
-            precomputed_selected=-0.1,
-        )
-        assert selected == pytest.approx(-0.1)
-        assert len(items) == 2
-        assert items[0].token == "tok_3"
-        assert items[0].logprob == pytest.approx(-0.1)
-        assert items[1].token == "tok_1"
-        assert items[1].logprob == pytest.approx(-1.0)
-
-    def test_slices_precomputed_to_requested_k(self) -> None:
-        lp = _make_logprobs([-99.0])
-        _, items = extract_top_logprobs(
-            lp,
-            _mock_tokenizer(),
-            top_logprobs=1,
-            selected_token=0,
-            precomputed_indices=[3, 1, 0, 2, 4],
-            precomputed_values=[-0.1, -1.0, -2.0, -3.0, -4.0],
-            precomputed_selected=-0.1,
-        )
-        assert len(items) == 1
-        assert items[0].token == "tok_3"
-
-    def test_falls_back_when_precomputed_partial(self) -> None:
-        lp = _make_logprobs([-1.0, -2.0, -0.5])
-        selected, items = extract_top_logprobs(
-            lp,
-            _mock_tokenizer(),
-            top_logprobs=2,
-            selected_token=2,
-            precomputed_indices=[0, 2],
-            precomputed_values=None,
-            precomputed_selected=None,
-        )
-        assert selected == pytest.approx(-0.5)
-        assert len(items) == 2
-
-    def test_precomputed_matches_fallback(self) -> None:
-        lp = _make_logprobs([-1.0, -0.3, -2.5, -0.1, -4.0, -0.8, -3.0, -1.5])
-        tok = _mock_tokenizer()
-
-        selected_fb, items_fb = extract_top_logprobs(
-            lp, tok, top_logprobs=5, selected_token=1
-        )
-
-        pre_indices = [item.token.split("_")[1] for item in items_fb]
-        pre_indices_int = [int(x) for x in pre_indices]
-        pre_values = [item.logprob for item in items_fb]
-
-        selected_pc, items_pc = extract_top_logprobs(
-            lp,
-            tok,
-            top_logprobs=5,
-            selected_token=1,
-            precomputed_indices=pre_indices_int,
-            precomputed_values=pre_values,
-            precomputed_selected=selected_fb,
-        )
-
-        assert selected_pc == pytest.approx(selected_fb)
-        assert len(items_pc) == len(items_fb)
-        for a, b in zip(items_pc, items_fb, strict=True):
-            assert a.token == b.token
-            assert a.logprob == pytest.approx(b.logprob)
-
-
-def _tiny_model() -> nn.Module:
-    from mlx_lm.models.llama import Model, ModelArgs
-
-    mx.random.seed(42)
-    args = ModelArgs(
-        model_type="llama",
-        hidden_size=64,
-        num_hidden_layers=2,
-        intermediate_size=128,
-        num_attention_heads=2,
-        num_key_value_heads=1,
-        rms_norm_eps=1e-6,
-        vocab_size=256,
-        rope_theta=10000.0,
-        tie_word_embeddings=True,
-    )
-    model = Model(args)
-    mx.eval(model.parameters())
-    return model
-
-
-@pytest.mark.slow
-class TestBatchedTopKPrecompute:
-    @pytest.fixture(autouse=True)
-    def _reset_globals(self) -> None:
-        import exo.worker.engines.mlx.patches.opt_batch_gen as _mod
-
-        _mod._pending_topk_idx = None
-        _mod._pending_topk_val = None
-        _mod._pending_selected_lps = None
-
-    def _run_generator(
-        self, model: nn.Module, prompts: list[list[int]], steps: int, needs_topk: bool
-    ) -> list[list[BatchGenerator.Response]]:
-        apply_batch_gen_patch()
-        gen = BatchGenerator(model=model, stop_tokens=set(), prefill_step_size=512)
-        gen._needs_topk = needs_topk
-        gen.insert(prompts)
-        all_responses: list[list[BatchGenerator.Response]] = []
-        for _ in range(steps + len(prompts)):
-            responses = gen.next()
-            if responses:
-                all_responses.append(responses)
-            if gen.active_batch is None and not gen.unprocessed_prompts:
-                break
-        gen.close()
-        return all_responses
-
-    def test_precomputed_topk_attached_to_responses(self) -> None:
-        model = _tiny_model()
-        steps = self._run_generator(model, [[1, 2, 3]], 5, needs_topk=True)
-        found_precomputed = False
-        for step_responses in steps:
-            for resp in step_responses:
-                if hasattr(resp, "_topk_indices"):
-                    found_precomputed = True
-                    assert hasattr(resp, "_topk_values"), (
-                        "Response missing _topk_values"
-                    )
-                    assert hasattr(resp, "_selected_logprob"), (
-                        "Response missing _selected_logprob"
-                    )
-                    assert len(resp._topk_indices) == _PRECOMPUTE_TOP_K
-                    assert len(resp._topk_values) == _PRECOMPUTE_TOP_K
-        assert found_precomputed, "No responses had precomputed topk"
-
-    def test_no_topk_when_not_needed(self) -> None:
-        model = _tiny_model()
-        steps = self._run_generator(model, [[1, 2, 3]], 5, needs_topk=False)
-        for step_responses in steps:
-            for resp in step_responses:
-                assert not hasattr(resp, "_topk_indices")
-
-    def test_precomputed_matches_fallback_in_batch(self) -> None:
-        model = _tiny_model()
-        tok = _mock_tokenizer()
-        steps = self._run_generator(model, [[1, 2, 3]], 10, needs_topk=True)
-        for step_responses in steps[1:]:
-            for resp in step_responses:
-                if not hasattr(resp, "_topk_indices"):
-                    continue
-                selected_fb, items_fb = extract_top_logprobs(
-                    resp.logprobs, tok, top_logprobs=5, selected_token=resp.token
-                )
-                selected_pc, items_pc = extract_top_logprobs(
-                    resp.logprobs,
-                    tok,
-                    top_logprobs=5,
-                    selected_token=resp.token,
-                    precomputed_indices=resp._topk_indices,
-                    precomputed_values=resp._topk_values,
-                    precomputed_selected=resp._selected_logprob,
-                )
-                assert selected_pc == pytest.approx(selected_fb, abs=1e-5)
-                for a, b in zip(items_pc, items_fb, strict=True):
-                    assert a.token == b.token
-                    assert a.logprob == pytest.approx(b.logprob, abs=1e-5)
-
-    def test_topk_correct_after_batch_shrink(self) -> None:
-        model = _tiny_model()
-        tok = _mock_tokenizer()
-        apply_batch_gen_patch()
-        gen = BatchGenerator(
-            model=model, stop_tokens={0}, prefill_step_size=512, max_tokens=3
-        )
-        gen._needs_topk = True
-        gen.insert([[1, 2, 3], [4, 5, 6]], max_tokens=[3, 20])
-
-        seen_shrink = False
-        for _ in range(30):
-            responses = gen.next()
-            for resp in responses:
-                if resp.finish_reason is not None:
-                    seen_shrink = True
-                    continue
-                if not hasattr(resp, "_topk_indices"):
-                    continue
-                selected_fb, items_fb = extract_top_logprobs(
-                    resp.logprobs, tok, top_logprobs=5, selected_token=resp.token
-                )
-                selected_pc, _ = extract_top_logprobs(
-                    resp.logprobs,
-                    tok,
-                    top_logprobs=5,
-                    selected_token=resp.token,
-                    precomputed_indices=resp._topk_indices,
-                    precomputed_values=resp._topk_values,
-                    precomputed_selected=resp._selected_logprob,
-                )
-                assert selected_pc == pytest.approx(selected_fb, abs=1e-5), (
-                    f"Mismatch after batch shrink: precomputed={selected_pc}, fallback={selected_fb}"
-                )
-            if gen.active_batch is None and not gen.unprocessed_prompts:
-                break
-
-        gen.close()
-        assert seen_shrink, "Expected at least one request to finish (batch shrink)"

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

@@ -486,7 +486,16 @@ def _patch_lossy_chat_template(template: str) -> str | None:
 
 
 def _needs_dsml_encoding(task_params: TextGenerationTaskParams) -> bool:
-    return "deepseek-v3.2" in task_params.model.lower()
+    if "deepseek-v3.2" not in task_params.model.lower():
+        return False
+    # Use DSML encoding when tools are provided or tool results are in the conversation
+    if task_params.tools:
+        return True
+    if task_params.chat_template_messages:
+        return any(
+            msg.get("role") == "tool" for msg in task_params.chat_template_messages
+        )
+    return False
 
 
 def apply_chat_template(
@@ -505,6 +514,8 @@ def apply_chat_template(
     if task_params.chat_template_messages is not None:
         # Use pre-formatted messages that preserve tool_calls, thinking, etc.
         formatted_messages = list(task_params.chat_template_messages)
+        for msg in formatted_messages:
+            _normalize_tool_calls(msg)
     else:
         # Add system message (instructions) if present
         if task_params.instructions:
@@ -530,10 +541,7 @@ def apply_chat_template(
 
         prompt = encode_messages(
             messages=formatted_messages,
-            # Only use chat mode if enable thinking is explicitly Fakse.
-            thinking_mode="chat"
-            if task_params.enable_thinking is False
-            else "thinking",
+            thinking_mode="thinking" if task_params.enable_thinking else "chat",
             tools=task_params.tools,
         )
         if partial_assistant_content:
@@ -541,9 +549,6 @@ def apply_chat_template(
         logger.info(prompt)
         return prompt
 
-    for msg in formatted_messages:
-        _normalize_tool_calls(msg)
-
     extra_kwargs: dict[str, Any] = {}
     if task_params.enable_thinking is not None:
         # Qwen3 and GLM use "enable_thinking"; DeepSeek uses "thinking".
@@ -633,7 +638,6 @@ class NullKVCache(KVCache):
     @property
     def state(self) -> tuple[mx.array, mx.array]:
         # matches what mx.save_safetensors / mx.eval expect
-        assert self.keys is not None and self.values is not None
         return self.keys, self.values
 
     @state.setter
@@ -735,9 +739,12 @@ def _parse_kimi_tool_calls(text: str):
         if func_args_match is None:
             raise ValueError("No tool call arguments found.")
         func_args = func_args_match.group(1)
-        arg_dct = json.loads(func_args)  # pyright: ignore[reportAny]
+        try:
+            arg_dct = json.loads(func_args)  # pyright: ignore[reportAny]
+        except Exception:
+            arg_dct = None
 
-        return dict(id=tool_call_id, name=func_name, arguments=arg_dct)  # pyright: ignore[reportAny]
+        return dict(id=tool_call_id, name=func_name, arguments=arg_dct)
 
     tool_matches = _tool_call_split_regex.findall(text)
     if tool_matches:

src/exo/worker/main.py

@@ -5,10 +5,10 @@ import anyio
 from anyio import fail_after
 from loguru import logger
 
-from exo.api.types import ImageEditsTaskParams
 from exo.download.download_utils import resolve_model_in_path
 from exo.shared.apply import apply
-from exo.shared.models.model_cards import ModelId, add_to_card_cache, delete_custom_card
+from exo.shared.models.model_cards import ModelId
+from exo.shared.types.api import ImageEditsTaskParams
 from exo.shared.types.commands import (
     ForwarderCommand,
     ForwarderDownloadCommand,
@@ -16,8 +16,6 @@ from exo.shared.types.commands import (
 )
 from exo.shared.types.common import CommandId, NodeId, SystemId
 from exo.shared.types.events import (
-    CustomModelCardAdded,
-    CustomModelCardDeleted,
     Event,
     IndexedEvent,
     InputChunkReceived,
@@ -132,13 +130,6 @@ class Worker:
                         event.chunk.data
                     )
 
-                if isinstance(event, CustomModelCardAdded):
-                    await event.model_card.save_to_custom_dir()
-                    add_to_card_cache(event.model_card)
-
-                if isinstance(event, CustomModelCardDeleted):
-                    await delete_custom_card(event.model_id)
-
     async def plan_step(self):
         while True:
             await anyio.sleep(0.1)

src/exo/worker/plan.py

@@ -55,7 +55,7 @@ def plan(
     # Python short circuiting OR logic should evaluate these sequentially.
     return (
         _cancel_tasks(runners, tasks)
-        or _kill_runner(runners, instances)
+        or _kill_runner(runners, all_runners, instances)
         or _create_runner(node_id, runners, instances)
         or _model_needs_download(node_id, runners, global_download_status)
         or _init_distributed_backend(runners, all_runners)
@@ -67,14 +67,25 @@ def plan(
 
 def _kill_runner(
     runners: Mapping[RunnerId, RunnerSupervisor],
+    all_runners: Mapping[RunnerId, RunnerStatus],
     instances: Mapping[InstanceId, Instance],
 ) -> Shutdown | None:
     for runner in runners.values():
         runner_id = runner.bound_instance.bound_runner_id
         if (instance_id := runner.bound_instance.instance.instance_id) not in instances:
             return Shutdown(instance_id=instance_id, runner_id=runner_id)
-        if isinstance(runner.status, RunnerFailed):
-            return Shutdown(instance_id=instance_id, runner_id=runner_id)
+
+        for (
+            global_runner_id
+        ) in runner.bound_instance.instance.shard_assignments.node_to_runner.values():
+            if runner_id == global_runner_id:
+                continue
+
+            if isinstance(all_runners.get(global_runner_id, None), RunnerFailed):
+                return Shutdown(
+                    instance_id=instance_id,
+                    runner_id=runner_id,
+                )
 
 
 def _create_runner(

src/exo/worker/runner/bootstrap.py

@@ -8,7 +8,6 @@ from exo.shared.types.tasks import Task, TaskId
 from exo.shared.types.worker.instances import BoundInstance
 from exo.shared.types.worker.runners import RunnerFailed
 from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender
-from exo.worker.engines.mlx.patches import apply_mlx_patches
 
 logger: "loguru.Logger" = loguru.logger
 
@@ -46,8 +45,6 @@ def entrypoint(
         else:
             from exo.worker.runner.llm_inference.runner import Runner
 
-            apply_mlx_patches()
-
             runner = Runner(
                 bound_instance, event_sender, task_receiver, cancel_receiver
             )

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

@@ -4,10 +4,10 @@ from typing import TYPE_CHECKING, Literal
 
 import mlx.core as mx
 
-from exo.api.types import ImageGenerationStats
 from exo.shared.constants import EXO_MAX_CHUNK_SIZE, EXO_TRACING_ENABLED
 from exo.shared.models.model_cards import ModelTask
 from exo.shared.tracing import clear_trace_buffer, get_trace_buffer
+from exo.shared.types.api import ImageGenerationStats
 from exo.shared.types.chunks import ErrorChunk, ImageChunk
 from exo.shared.types.common import CommandId, ModelId
 from exo.shared.types.events import (
@@ -39,6 +39,7 @@ from exo.shared.types.worker.runner_response import (
 from exo.shared.types.worker.runners import (
     RunnerConnected,
     RunnerConnecting,
+    RunnerFailed,
     RunnerIdle,
     RunnerLoaded,
     RunnerLoading,
@@ -255,7 +256,9 @@ class Runner:
 
     def handle_task(self, task: Task):
         match task:
-            case ConnectToGroup() if isinstance(self.current_status, RunnerIdle):
+            case ConnectToGroup() if isinstance(
+                self.current_status, (RunnerIdle, RunnerFailed)
+            ):
                 logger.info("runner connecting")
                 self.update_status(RunnerConnecting())
                 self.acknowledge_task(task)

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

@@ -195,29 +195,21 @@ class SequentialGenerator(InferenceGenerator):
         assert self._active is not None
 
         task, mlx_gen, queue, output_generator = self._active
-        output: list[
-            tuple[TaskId, GenerationResponse | ToolCallResponse | Cancelled | Finished]
-        ] = []
+        response = None
         try:
-            response = next(mlx_gen)
-            queue.push(response)
-            # drain potentially many responses every time
-            while (parsed := next(output_generator, None)) is not None:
-                output.append((task.task_id, parsed))
-
+            queue.push(next(mlx_gen))
+            response = next(output_generator)
         except (StopIteration, PrefillCancelled):
-            output.append((task.task_id, Finished()))
+            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(
-            output,
+            [] if response is None else [(task.task_id, response)],
             map(lambda task: (task, Cancelled()), self._cancelled_tasks),
         )
 
@@ -435,11 +427,11 @@ class BatchGenerator(InferenceGenerator):
 
             task, queue, output_generator = self._active_tasks[uid]
             queue.push(response)
-            # If a generator fails to parse for some reason and returns early, we should not crash
-            while (parsed := next(output_generator, None)) is not None:
+            parsed = next(output_generator)
+
+            if parsed is not None:
                 output.append((task.task_id, parsed))
 
-            # check if original response was terminal and append a Finished()
             if response.finish_reason is not None:
                 output.append((task.task_id, Finished()))
                 del self._active_tasks[uid]

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

@@ -13,7 +13,7 @@ from openai_harmony import (  # pyright: ignore[reportMissingTypeStubs]
     load_harmony_encoding,
 )
 
-from exo.api.types import ToolCallItem
+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
@@ -159,42 +159,11 @@ def parse_deepseek_v32(
     # Text accumulated during a tool call block
     tool_call_text = ""
 
-    def _try_parse_tool_call(
-        text: str, response: GenerationResponse
-    ) -> ToolCallResponse | GenerationResponse:
-        parsed = parse_dsml_output(text)
-        if parsed is not None:
-            return ToolCallResponse(
-                tool_calls=parsed, usage=response.usage, stats=response.stats
-            )
-        logger.warning(f"DSML tool call parsing failed for: {text}")
-        return response.model_copy(update={"text": text})
-
     for response in responses:
         if response is None:
             yield None
             continue
 
-        if response.finish_reason is not None:
-            yield from pending_buffer
-            pending_buffer.clear()
-            if in_tool_call:
-                tool_call_text += response.text
-                yield (
-                    _try_parse_tool_call(tool_call_text, response)
-                    if TOOL_CALLS_END in tool_call_text
-                    else response.model_copy(update={"text": tool_call_text})
-                )
-            elif TOOL_CALLS_START in response.text and TOOL_CALLS_END in response.text:
-                dsml_start = response.text.index(TOOL_CALLS_START)
-                before = response.text[:dsml_start]
-                if before:
-                    yield response.model_copy(update={"text": before})
-                yield _try_parse_tool_call(response.text[dsml_start:], response)
-            else:
-                yield response
-            break
-
         # ── Handle thinking tags ──
         if not thinking and THINKING_START in response.text:
             thinking = True
@@ -222,7 +191,28 @@ def parse_deepseek_v32(
         if in_tool_call:
             tool_call_text += response.text
             if TOOL_CALLS_END in tool_call_text:
-                yield _try_parse_tool_call(tool_call_text, response)
+                # 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
@@ -238,22 +228,33 @@ def parse_deepseek_v32(
             if pre_text:
                 # Flush pending buffer tokens that contributed text before the marker
                 for buf_resp in pending_buffer:
-                    if not pre_text:
-                        break
-                    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 = ""
+                    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:
-                yield _try_parse_tool_call(tool_call_text, response)
+                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
@@ -266,13 +267,15 @@ def parse_deepseek_v32(
             continue
 
         # No partial match — flush all pending tokens and the current one
-        yield from pending_buffer
-        pending_buffer.clear()
+        for buf_resp in pending_buffer:
+            yield buf_resp
+        pending_buffer = []
         accumulated = ""
         yield response
 
     # Flush any remaining pending buffer at generator end
-    yield from pending_buffer
+    for buf_resp in pending_buffer:
+        yield buf_resp
 
 
 def _could_be_dsml_prefix(text: str) -> bool:
@@ -355,10 +358,8 @@ def parse_tool_calls(
 
             if parsed is None:
                 logger.warning(f"tool call parsing failed for text {combined}")
-                yield response.model_copy(
-                    update={"text": combined, "token": 0, "finish_reason": "error"}
-                )
-                break
+                yield response.model_copy(update={"text": combined})
+                continue
 
             yield ToolCallResponse(
                 tool_calls=parsed, usage=response.usage, stats=response.stats
@@ -373,7 +374,6 @@ def parse_tool_calls(
                 update={
                     "text": "".join(tool_call_text_parts),
                     "token": 0,
-                    "finish_reason": "error",
                 }
             )
             yield response

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

@@ -146,7 +146,9 @@ class Runner:
         self.send_task_status(task.task_id, TaskStatus.Running)
 
         match task:
-            case ConnectToGroup() if isinstance(self.current_status, RunnerIdle):
+            case ConnectToGroup() if isinstance(
+                self.current_status, (RunnerIdle, RunnerFailed)
+            ):
                 assert isinstance(self.generator, Builder)
                 logger.info("runner connecting")
                 self.update_status(RunnerConnecting())
@@ -317,9 +319,7 @@ class Runner:
         return ExitCode.AllTasksComplete
 
     def send_response(
-        self,
-        response: GenerationResponse | ToolCallResponse,
-        command_id: CommandId,
+        self, response: GenerationResponse | ToolCallResponse, command_id: CommandId
     ):
         match response:
             case GenerationResponse():

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

@@ -3,7 +3,7 @@ import math
 from dataclasses import dataclass
 from typing import Any, Callable
 
-from exo.api.types import ToolCallItem
+from exo.shared.types.api import ToolCallItem
 
 
 @dataclass

src/exo/worker/runner/runner_supervisor.py

@@ -24,7 +24,6 @@ from exo.shared.types.tasks import (
     CANCEL_ALL_TASKS,
     ImageEdits,
     ImageGeneration,
-    Shutdown,
     Task,
     TaskId,
     TaskStatus,
@@ -111,45 +110,39 @@ class RunnerSupervisor:
 
     async def run(self):
         self.runner_process.start()
-        try:
-            async with self._tg as tg:
-                tg.start_soon(self._watch_runner)
-                tg.start_soon(self._forward_events)
-        finally:
-            logger.info("Runner supervisor shutting down")
-            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()
-            with contextlib.suppress(ClosedResourceError):
-                self._cancel_sender.send(CANCEL_ALL_TASKS)
-            with contextlib.suppress(ClosedResourceError):
-                self._cancel_sender.close()
-
-            await to_thread.run_sync(self.runner_process.join, 5)
-
-            if self.runner_process.is_alive():
-                logger.warning(
-                    "Runner process didn't shutdown succesfully, terminating"
-                )
-                self.runner_process.terminate()
-                self.runner_process.join(timeout=5)
-                # This is overkill but it's not technically bad, just unnecessary.
-                if self.runner_process.is_alive():
-                    logger.critical("Runner process didn't respond to SIGTERM, killing")
-                    self.runner_process.kill()
-                    self.runner_process.join(timeout=5)
-            else:
-                logger.info("Runner process succesfully terminated")
-
-            self.runner_process.close()
+        async with self._tg as tg:
+            tg.start_soon(self._watch_runner)
+            tg.start_soon(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()
+        with contextlib.suppress(ClosedResourceError):
+            self._cancel_sender.send(CANCEL_ALL_TASKS)
+        with contextlib.suppress(ClosedResourceError):
+            self._cancel_sender.close()
+        self.runner_process.join(5)
+        if not self.runner_process.is_alive():
+            logger.info("Runner process succesfully terminated")
+            return
+
+        # This is overkill but it's not technically bad, just unnecessary.
+        logger.warning("Runner process didn't shutdown succesfully, terminating")
+        self.runner_process.terminate()
+        self.runner_process.join(1)
+        if not self.runner_process.is_alive():
+            return
+
+        logger.critical("Runner process didn't respond to SIGTERM, killing")
+        self.runner_process.kill()
 
     async def start_task(self, task: Task):
         if task.task_id in self.pending:
@@ -170,8 +163,7 @@ class RunnerSupervisor:
             await self._task_sender.send_async(task)
         except ClosedResourceError:
             self.in_progress.pop(task.task_id, None)
-            if not isinstance(task, Shutdown):
-                logger.warning(f"Task {task} dropped, runner closed communication.")
+            logger.warning(f"Task {task} dropped, runner closed communication.")
             return
         await event.wait()
 
@@ -226,6 +218,12 @@ class RunnerSupervisor:
             for tid in self.pending:
                 self.pending[tid].set()
 
+    def __del__(self) -> None:
+        if self.runner_process.is_alive():
+            logger.critical("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:

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

@@ -43,8 +43,8 @@ def run_pipeline_device(
 
         def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
             for layer in self.layers:
-                x = layer(x, *args, **kwargs)
-            return x
+                x = layer(x, *args, **kwargs)  # pyright: ignore[reportUnknownVariableType]
+            return x  # pyright: ignore[reportUnknownVariableType]
 
     try:
         group = mx.distributed.init(backend="ring", strict=True)