flush old addrs

Mb premature auto merge

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

bench/exo_bench.py

@@ -496,20 +496,28 @@ def main() -> int:
                                 all_rows.append(row)
 
                             if batch_results:
-                                agg_gen_tps = sum(
+                                valid_gen_tps = [
                                     x["stats"]["generation_tps"]
                                     for x, _ in batch_results
                                     if x["stats"]["generation_tps"] > 0
+                                ]
+                                agg_gen_tps = (
+                                    mean(valid_gen_tps) if valid_gen_tps else 0.0
                                 )
+                                gen_tps = agg_gen_tps / concurrency
                                 logger.info(
                                     f"[concurrent {concurrency}x]  "
                                     f"agg_gen_tps={agg_gen_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)
-                        gen_tps = mean(x["stats"]["generation_tps"] for x in runs)
+                        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/ChatForm.svelte

@@ -1,9 +1,6 @@
 <script lang="ts">
   import {
     isLoading,
-    sendMessage,
-    generateImage,
-    editImage,
     editingImage,
     clearEditingImage,
     selectedChatModel,
@@ -28,7 +25,7 @@
     modelTasks?: Record<string, string[]>;
     modelCapabilities?: Record<string, string[]>;
     onSend?: () => void;
-    onAutoSend?: (
+    onAutoSend: (
       content: string,
       files?: {
         id: string;
@@ -216,49 +213,10 @@
     uploadedFiles = [];
     resetTextareaHeight();
 
-    // When onAutoSend is provided, the parent controls all send logic
-    // (including launching non-running models before sending)
-    if (onAutoSend) {
-      onAutoSend(content, files);
-      onSend?.();
-      setTimeout(() => textareaRef?.focus(), 10);
-      return;
-    }
-
-    // Use image editing if in edit mode
-    if (isEditMode && currentEditingImage && content) {
-      editImage(content, currentEditingImage.imageDataUrl);
-    }
-    // If user attached an image with an ImageToImage model, use edit endpoint
-    else if (
-      currentModel &&
-      modelSupportsImageEditing(currentModel) &&
-      files.length > 0 &&
-      content
-    ) {
-      // Use the first attached image for editing
-      const imageFile = files[0];
-      if (imageFile.preview) {
-        editImage(content, imageFile.preview);
-      }
-    } else if (
-      currentModel &&
-      modelSupportsTextToImage(currentModel) &&
-      content
-    ) {
-      // Use image generation for text-to-image models
-      generateImage(content);
-    } else {
-      sendMessage(
-        content,
-        files,
-        modelSupportsThinking() ? thinkingEnabled : null,
-      );
-    }
-
+    // Parent controls all send logic (including image routing,
+    // launching non-running models before sending, etc.)
+    onAutoSend(content, files);
     onSend?.();
-
-    // Refocus the textarea after sending
     setTimeout(() => textareaRef?.focus(), 10);
   }
 

dashboard/src/lib/components/FamilyLogos.svelte

@@ -82,6 +82,12 @@
       d="M12.025 1.13c-5.77 0-10.449 4.647-10.449 10.378 0 1.112.178 2.181.503 3.185.064-.222.203-.444.416-.577a.96.96 0 0 1 .524-.15c.293 0 .584.124.84.284.278.173.48.408.71.694.226.282.458.611.684.951v-.014c.017-.324.106-.622.264-.874s.403-.487.762-.543c.3-.047.596.06.787.203s.31.313.4.467c.15.257.212.468.233.542.01.026.653 1.552 1.657 2.54.616.605 1.01 1.223 1.082 1.912.055.537-.096 1.059-.38 1.572.637.121 1.294.187 1.967.187.657 0 1.298-.063 1.921-.178-.287-.517-.44-1.041-.384-1.581.07-.69.465-1.307 1.081-1.913 1.004-.987 1.647-2.513 1.657-2.539.021-.074.083-.285.233-.542.09-.154.208-.323.4-.467a1.08 1.08 0 0 1 .787-.203c.359.056.604.29.762.543s.247.55.265.874v.015c.225-.34.457-.67.683-.952.23-.286.432-.52.71-.694.257-.16.547-.284.84-.285a.97.97 0 0 1 .524.151c.228.143.373.388.43.625l.006.04a10.3 10.3 0 0 0 .534-3.273c0-5.731-4.678-10.378-10.449-10.378M8.327 6.583a1.5 1.5 0 0 1 .713.174 1.487 1.487 0 0 1 .617 2.013c-.183.343-.762-.214-1.102-.094-.38.134-.532.914-.917.71a1.487 1.487 0 0 1 .69-2.803m7.486 0a1.487 1.487 0 0 1 .689 2.803c-.385.204-.536-.576-.916-.71-.34-.12-.92.437-1.103.094a1.487 1.487 0 0 1 .617-2.013 1.5 1.5 0 0 1 .713-.174m-10.68 1.55a.96.96 0 1 1 0 1.921.96.96 0 0 1 0-1.92m13.838 0a.96.96 0 1 1 0 1.92.96.96 0 0 1 0-1.92M8.489 11.458c.588.01 1.965 1.157 3.572 1.164 1.607-.007 2.984-1.155 3.572-1.164.196-.003.305.12.305.454 0 .886-.424 2.328-1.563 3.202-.22-.756-1.396-1.366-1.63-1.32q-.011.001-.02.006l-.044.026-.01.008-.03.024q-.018.017-.035.036l-.032.04a1 1 0 0 0-.058.09l-.014.025q-.049.088-.11.19a1 1 0 0 1-.083.116 1.2 1.2 0 0 1-.173.18q-.035.029-.075.058a1.3 1.3 0 0 1-.251-.243 1 1 0 0 1-.076-.107c-.124-.193-.177-.363-.337-.444-.034-.016-.104-.008-.2.022q-.094.03-.216.087-.06.028-.125.063l-.13.074q-.067.04-.136.086a3 3 0 0 0-.135.096 3 3 0 0 0-.26.219 2 2 0 0 0-.12.121 2 2 0 0 0-.106.128l-.002.002a2 2 0 0 0-.09.132l-.001.001a1.2 1.2 0 0 0-.105.212q-.013.036-.024.073c-1.139-.875-1.563-2.317-1.563-3.203 0-.334.109-.457.305-.454m.836 10.354c.824-1.19.766-2.082-.365-3.194-1.13-1.112-1.789-2.738-1.789-2.738s-.246-.945-.806-.858-.97 1.499.202 2.362c1.173.864-.233 1.45-.685.64-.45-.812-1.683-2.896-2.322-3.295s-1.089-.175-.938.647 2.822 2.813 2.562 3.244-1.176-.506-1.176-.506-2.866-2.567-3.49-1.898.473 1.23 2.037 2.16c1.564.932 1.686 1.178 1.464 1.53s-3.675-2.511-4-1.297c-.323 1.214 3.524 1.567 3.287 2.405-.238.839-2.71-1.587-3.216-.642-.506.946 3.49 2.056 3.522 2.064 1.29.33 4.568 1.028 5.713-.624m5.349 0c-.824-1.19-.766-2.082.365-3.194 1.13-1.112 1.789-2.738 1.789-2.738s.246-.945.806-.858.97 1.499-.202 2.362c-1.173.864.233 1.45.685.64.451-.812 1.683-2.896 2.322-3.295s1.089-.175.938.647-2.822 2.813-2.562 3.244 1.176-.506 1.176-.506 2.866-2.567 3.49-1.898-.473 1.23-2.037 2.16c-1.564.932-1.686 1.178-1.464 1.53s3.675-2.511 4-1.297c.323 1.214-3.524 1.567-3.287 2.405.238.839 2.71-1.587 3.216-.642.506.946-3.49 2.056-3.522 2.064-1.29.33-4.568 1.028-5.713-.624"
     />
   </svg>
+{:else if family === "step"}
+  <svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
+    <path
+      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}
   <svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
     <path

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

@@ -1577,6 +1577,7 @@ class AppStore {
     // Remove messages after user message (including the user message for image requests
     // since generateImage/editImage will re-add it)
     this.messages = this.messages.slice(0, lastUserIndex);
+    this.updateActiveConversation();
 
     switch (requestType) {
       case "image-generation":

dashboard/src/routes/+page.svelte

@@ -42,6 +42,9 @@
     setSelectedChatModel,
     selectedChatModel,
     sendMessage,
+    generateImage,
+    editImage,
+    editingImage,
     messages,
     debugMode,
     toggleDebugMode,
@@ -834,6 +837,52 @@
     if (!model?.tasks) return false;
     return model.tasks.includes("ImageToImage");
   }
+
+  // Route a message to the correct endpoint based on model capabilities.
+  // Image models go to generateImage/editImage; text models go to sendMessage.
+  function routeMessage(
+    content: string,
+    files?: {
+      id: string;
+      name: string;
+      type: string;
+      textContent?: string;
+      preview?: string;
+    }[],
+  ) {
+    const model = selectedChatModel();
+    if (!model) {
+      sendMessage(content, files, null);
+      return;
+    }
+
+    const currentEditImage = editingImage();
+
+    // Image editing mode (explicit edit or attached image with ImageToImage model)
+    if (currentEditImage && content && modelSupportsImageEditing(model)) {
+      editImage(content, currentEditImage.imageDataUrl);
+      return;
+    }
+    if (
+      modelSupportsImageEditing(model) &&
+      files?.length &&
+      files[0].preview &&
+      content
+    ) {
+      editImage(content, files[0].preview);
+      return;
+    }
+
+    // Text-to-image generation
+    if (modelSupportsImageGeneration(model) && content) {
+      generateImage(content);
+      return;
+    }
+
+    // Default: text chat
+    sendMessage(content, files, null);
+  }
+
   let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline");
   type InstanceMeta = "MlxRing" | "MlxJaccl";
 
@@ -2851,7 +2900,7 @@
         // Running model is same or better tier — use it directly
         setSelectedChatModel(bestRunning.id);
         if (!chatStarted) createConversation();
-        sendMessage(content, files);
+        routeMessage(content, files);
         return;
       }
     }
@@ -2868,7 +2917,7 @@
     if (hasRunningInstance(autoModel.id)) {
       setSelectedChatModel(autoModel.id);
       if (!chatStarted) createConversation();
-      sendMessage(content, files);
+      routeMessage(content, files);
       return;
     }
 
@@ -3021,7 +3070,7 @@
       if (pendingAutoMessage) {
         const msg = pendingAutoMessage;
         pendingAutoMessage = null;
-        sendMessage(msg.content, msg.files);
+        routeMessage(msg.content, msg.files);
       }
       return;
     }
@@ -3100,7 +3149,7 @@
     // Model is selected and running — send directly
     if (model && hasRunningInstance(model)) {
       chatLaunchState = "ready";
-      sendMessage(content, files, null);
+      routeMessage(content, files);
       return;
     }
 

flake.nix

@@ -112,17 +112,20 @@
             };
           };
 
-          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);
               uvLockMlxVersion = mlxPackage.version;
+              uvLockMlxRev = builtins.elemAt (builtins.split "#" mlxPackage.source.git) 2;
             in
             {
               metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { };
               mlx = pkgs.callPackage ./nix/mlx.nix {
                 inherit (self'.packages) metal-toolchain;
-                inherit uvLockMlxVersion;
+                inherit uvLockMlxVersion uvLockMlxRev;
               };
               default = self'.packages.exo;
             }
@@ -156,9 +159,6 @@
                 just
                 jq
               ]
-              ++ lib.optionals stdenv.isLinux [
-                unixtools.ifconfig
-              ]
               ++ lib.optionals stdenv.isDarwin [
                 macmon
               ];

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=''";
+}
+

nix/mlx.nix

@@ -11,6 +11,7 @@
 , fmt
 , python313Packages
 , uvLockMlxVersion
+, uvLockMlxRev
 }:
 
 assert stdenv.isDarwin;
@@ -41,15 +42,13 @@ let
 
   mlx = stdenv.mkDerivation rec {
     pname = "mlx";
-    version = let v = "0.30.7.dev20260225+257d5692"; in
-      assert v == uvLockMlxVersion || throw "MLX version mismatch: nix/mlx.nix has ${v} but uv.lock has ${uvLockMlxVersion}. Update both the version and hash in nix/mlx.nix.";
-      v;
+    version = uvLockMlxVersion;
     pyproject = true;
 
     src = fetchFromGitHub {
       owner = "rltakashige";
       repo = "mlx-jaccl-fix-small-recv";
-      rev = "257d5692fc7af6bba3b8afaeb63c549b7d1e43d5";
+      rev = uvLockMlxRev;
       hash = "sha256-GosFIWxIB48Egb1MqJrR3xhsUsQeWdRk5rV93USY6wQ=";
     };
 

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 = {

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/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/download/coordinator.py

@@ -1,4 +1,3 @@
-import asyncio
 from dataclasses import dataclass, field
 
 import anyio
@@ -47,7 +46,7 @@ class DownloadCoordinator:
 
     # Local state
     download_status: dict[ModelId, DownloadProgress] = field(default_factory=dict)
-    active_downloads: dict[ModelId, asyncio.Task[None]] = field(default_factory=dict)
+    active_downloads: dict[ModelId, anyio.CancelScope] = field(default_factory=dict)
 
     _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
 
@@ -77,8 +76,6 @@ class DownloadCoordinator:
             await self.event_sender.send(
                 NodeDownloadProgress(download_progress=completed)
             )
-            if model_id in self.active_downloads:
-                del self.active_downloads[model_id]
             self._last_progress_time.pop(model_id, None)
         elif (
             progress.status == "in_progress"
@@ -103,13 +100,9 @@ class DownloadCoordinator:
         logger.info(
             f"Starting DownloadCoordinator{' (offline mode)' if self.offline else ''}"
         )
-        try:
-            async with self._tg as tg:
-                tg.start_soon(self._command_processor)
-                tg.start_soon(self._emit_existing_download_progress)
-        finally:
-            for task in self.active_downloads.values():
-                task.cancel()
+        async with self._tg as tg:
+            tg.start_soon(self._command_processor)
+            tg.start_soon(self._emit_existing_download_progress)
 
     def shutdown(self) -> None:
         self._tg.cancel_tasks()
@@ -132,7 +125,7 @@ class DownloadCoordinator:
     async def _cancel_download(self, model_id: ModelId) -> None:
         if model_id in self.active_downloads and model_id in self.download_status:
             logger.info(f"Cancelling download for {model_id}")
-            self.active_downloads.pop(model_id).cancel()
+            self.active_downloads[model_id].cancel()
             current_status = self.download_status[model_id]
             pending = DownloadPending(
                 shard_metadata=current_status.shard_metadata,
@@ -236,9 +229,10 @@ class DownloadCoordinator:
         self.download_status[model_id] = status
         self.event_sender.send_nowait(NodeDownloadProgress(download_progress=status))
 
-        async def download_wrapper() -> None:
+        async def download_wrapper(cancel_scope: anyio.CancelScope) -> None:
             try:
-                await self.shard_downloader.ensure_shard(shard)
+                with cancel_scope:
+                    await self.shard_downloader.ensure_shard(shard)
             except Exception as e:
                 logger.error(f"Download failed for {model_id}: {e}")
                 failed = DownloadFailed(
@@ -251,12 +245,15 @@ class DownloadCoordinator:
                 await self.event_sender.send(
                     NodeDownloadProgress(download_progress=failed)
                 )
+            except anyio.get_cancelled_exc_class():
+                # ignore cancellation - let cleanup do its thing
+                pass
             finally:
-                if model_id in self.active_downloads:
-                    del self.active_downloads[model_id]
+                self.active_downloads.pop(model_id, None)
 
-        task = asyncio.create_task(download_wrapper())
-        self.active_downloads[model_id] = task
+        scope = anyio.CancelScope()
+        self._tg.start_soon(download_wrapper, scope)
+        self.active_downloads[model_id] = scope
 
     async def _delete_download(self, model_id: ModelId) -> None:
         # Protect read-only models (from EXO_MODELS_PATH) from deletion
@@ -272,7 +269,6 @@ class DownloadCoordinator:
         if model_id in self.active_downloads:
             logger.info(f"Cancelling active download for {model_id} before deletion")
             self.active_downloads[model_id].cancel()
-            del self.active_downloads[model_id]
 
         # Delete from disk
         logger.info(f"Deleting model files for {model_id}")

src/exo/utils/info_gatherer/info_gatherer.py

@@ -529,6 +529,9 @@ class InfoGatherer:
         if self.macmon_interval is None:
             return
         # macmon pipe --interval [interval in ms]
+        # Timeout: if macmon produces no output for this many seconds, restart it.
+        # macmon writes every macmon_interval seconds, so 10x that is generous.
+        read_timeout = max(self.macmon_interval * 10, 30)
         while True:
             try:
                 async with await open_process(
@@ -542,10 +545,15 @@ class InfoGatherer:
                     if not p.stdout:
                         logger.critical("MacMon closed stdout")
                         return
-                    async for text in TextReceiveStream(
-                        BufferedByteReceiveStream(p.stdout)
-                    ):
+                    stream = TextReceiveStream(BufferedByteReceiveStream(p.stdout))
+                    while True:
+                        with fail_after(read_timeout):
+                            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"
+                )
             except CalledProcessError as e:
                 stderr_msg = "no stderr"
                 stderr_output = cast(bytes | str | None, e.stderr)

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

@@ -65,6 +65,7 @@ class _EngineTask:
     generation_start_time: float = 0.0
     in_thinking: bool = False
     reasoning_tokens: int = 0
+    prefill_tps: float = 0.0
 
 
 @dataclass(eq=False)
@@ -209,6 +210,7 @@ class ExoBatchGenerator:
             detokenizer=self.tokenizer.detokenizer,
             on_generation_token=on_generation_token,
             generation_start_time=time.perf_counter(),
+            prefill_tps=_prefill_tps,
         )
 
         return uid
@@ -285,21 +287,22 @@ class ExoBatchGenerator:
             stats: GenerationStats | None = None
             usage: Usage | None = None
             if is_done:
-                generation_elapsed = time.perf_counter() - state.generation_start_time
-                generation_tps = (
-                    state.completion_tokens / generation_elapsed
-                    if generation_elapsed > 0
-                    else 0.0
-                )
                 try:
                     mlx_stats = self._exo_gen.stats()
+                    generation_tps = mlx_stats.generation_tps
                 except ZeroDivisionError:
-                    mlx_stats = None
+                    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=float(mlx_stats.prompt_tps)
-                    if mlx_stats is not None and mlx_stats.prompt_time > 0
-                    else 0.0,
-                    generation_tps=float(generation_tps),
+                    prompt_tps=state.prefill_tps,
+                    generation_tps=generation_tps,
                     prompt_tokens=len(state.all_prompt_tokens),
                     generation_tokens=state.completion_tokens,
                     peak_memory_usage=Memory.from_gb(mx.get_peak_memory() / 1e9),