implement engine interface for mlx and mflux

.mlx_typings/mlx_lm/generate.pyi

@@ -383,12 +383,11 @@ class GenerationBatch:
     state_machines: List[SequenceStateMachine]
     max_tokens: List[int]
     _current_tokens: Optional[mx.array]
-    _current_logprobs: mx.array | List[mx.array]
-    _next_tokens: Optional[mx.array]
-    _next_logprobs: mx.array | List[mx.array]
-    _token_context: List[Any]
+    _current_logprobs: List[mx.array]
+    _next_tokens: mx.array
+    _next_logprobs: List[mx.array]
+    _token_context: List[mx.array]
     _num_tokens: List[int]
-    _matcher_states: List[Any]
 
     def __init__(
         self,

.mlx_typings/mlx_lm/models/deepseek_v4.pyi

@@ -1,280 +0,0 @@
-"""Type stubs for mlx_lm.models.deepseek_v4"""
-
-from dataclasses import dataclass
-from typing import Any, Dict, List, Optional
-
-import mlx.core as mx
-import mlx.nn as nn
-
-from .base import BaseModelArgs
-from .cache import ArraysCache, RotatingKVCache
-from .switch_layers import SwitchGLU
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    model_type: str
-    vocab_size: int
-    hidden_size: int
-    intermediate_size: int
-    moe_intermediate_size: int
-    num_hidden_layers: int
-    num_attention_heads: int
-    num_key_value_heads: int
-    n_shared_experts: Optional[int]
-    n_routed_experts: int
-    num_experts_per_tok: int
-    head_dim: int
-    qk_rope_head_dim: int
-    q_lora_rank: int
-    o_lora_rank: int
-    o_groups: int
-    sliding_window: int
-    hc_mult: int
-    hc_sinkhorn_iters: int
-    hc_eps: float
-    compress_ratios: Optional[List[int]]
-    compress_rope_theta: float
-    rope_theta: float
-    rope_scaling: Optional[Dict[str, Any]]
-    rms_norm_eps: float
-    swiglu_limit: float
-    attention_bias: bool
-    max_position_embeddings: int
-
-class DeepseekV4RoPE(nn.Module):
-    dims: int
-    freqs: mx.array
-
-    def __init__(
-        self,
-        dims: int,
-        base: float,
-        scaling_config: Optional[Dict[str, Any]] = None,
-    ) -> None: ...
-    def __call__(
-        self,
-        x: mx.array,
-        offset: int = 0,
-        inverse: bool = False,
-    ) -> mx.array: ...
-
-class HyperConnection(nn.Module):
-    dim: int
-    hc_mult: int
-    norm_eps: float
-
-    def __init__(
-        self,
-        dim: int,
-        hc_mult: int,
-        norm_eps: float,
-        sinkhorn_iters: int,
-        hc_eps: float,
-    ) -> None: ...
-
-class HyperHead(nn.Module):
-    dim: int
-    hc_mult: int
-
-    def __init__(
-        self,
-        dim: int,
-        hc_mult: int,
-        norm_eps: float,
-        hc_eps: float,
-    ) -> None: ...
-    def __call__(self, x: mx.array) -> mx.array: ...
-
-class Compressor(nn.Module):
-    dim: int
-    head_dim: int
-    rope_head_dim: int
-    compress_ratio: int
-    overlap: bool
-    wkv_gate: nn.Linear
-    ape: mx.array
-    norm: nn.RMSNorm
-    rope: DeepseekV4RoPE
-
-    def __init__(
-        self,
-        dim: int,
-        compress_ratio: int,
-        head_dim: int,
-        rope_head_dim: int,
-        rms_norm_eps: float,
-        rope: DeepseekV4RoPE,
-    ) -> None: ...
-    def __call__(
-        self,
-        x: mx.array,
-        cache: "DeepseekV4Cache",
-        offset: Any,
-        key: str = ...,
-    ) -> mx.array: ...
-
-class Indexer(nn.Module):
-    def __init__(
-        self,
-        args: ModelArgs,
-        compress_ratio: int,
-        rope: DeepseekV4RoPE,
-    ) -> None: ...
-
-class _CompressorBranch:
-    buffer_kv: Optional[mx.array]
-    buffer_gate: Optional[mx.array]
-    prev_kv: Optional[mx.array]
-    prev_gate: Optional[mx.array]
-    pool: Optional[mx.array]
-    buffer_lengths: Optional[List[int]]
-    pool_lengths: Optional[List[int]]
-    buffer_count: int
-    _new_pool_lengths: Optional[List[int]]
-
-    def __init__(self) -> None: ...
-
-class DeepseekV4Cache:
-    local: RotatingKVCache
-    offset: int
-    keys: Optional[mx.array]
-    values: Optional[mx.array]
-    state: Any
-    meta_state: Any
-    nbytes: int
-    _branches: Dict[str, _CompressorBranch]
-    _pending_lengths: Optional[List[int]]
-
-    def __init__(self, sliding_window: int) -> None: ...
-    def update_and_fetch(
-        self, keys: mx.array, values: mx.array
-    ) -> tuple[mx.array, mx.array]: ...
-    def is_trimmable(self) -> bool: ...
-    def trim(self, n: int) -> int: ...
-    def empty(self) -> bool: ...
-    def size(self) -> int: ...
-    def prepare(
-        self,
-        *,
-        left_padding: Optional[List[int]] = None,
-        lengths: Optional[List[int]] = None,
-        right_padding: Optional[List[int]] = None,
-    ) -> None: ...
-    def finalize(self) -> None: ...
-    def filter(self, batch_indices: mx.array) -> None: ...
-    def extend(self, other: "DeepseekV4Cache") -> None: ...
-    def extract(self, idx: int) -> "DeepseekV4Cache": ...
-    @classmethod
-    def merge(cls, caches: List["DeepseekV4Cache"]) -> "DeepseekV4Cache": ...
-
-class V4Attention(nn.Module):
-    args: ModelArgs
-    layer_id: int
-    dim: int
-    n_heads: int
-    head_dim: int
-    rope_head_dim: int
-    nope_head_dim: int
-    n_groups: int
-    q_lora_rank: int
-    o_lora_rank: int
-    window: int
-    eps: float
-    scale: float
-    compress_ratio: int
-    wqkv_a: nn.Linear
-    q_norm: nn.RMSNorm
-    wq_b: nn.Linear
-    kv_norm: nn.RMSNorm
-    attn_sink: mx.array
-    wo_a: nn.Linear
-    wo_b: nn.Linear
-    rope: DeepseekV4RoPE
-    compressor: Compressor
-    indexer: Indexer
-
-    def __init__(self, args: ModelArgs, layer_id: int) -> None: ...
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-    ) -> mx.array: ...
-
-class DeepseekV4MLP(nn.Module):
-    gate_proj: nn.Linear
-    up_proj: nn.Linear
-    down_proj: nn.Linear
-
-    def __init__(
-        self,
-        hidden_size: int,
-        intermediate_size: int,
-        swiglu_limit: float = 0.0,
-    ) -> None: ...
-    def __call__(self, x: mx.array) -> mx.array: ...
-
-class MoEGate(nn.Module):
-    weight: mx.array
-
-    def __init__(self, args: ModelArgs, layer_id: int) -> None: ...
-    def __call__(
-        self, x: mx.array, input_ids: mx.array
-    ) -> tuple[mx.array, mx.array]: ...
-
-class DeepseekV4MoE(nn.Module):
-    num_experts_per_tok: int
-    switch_mlp: SwitchGLU
-    gate: MoEGate
-    shared_experts: DeepseekV4MLP
-
-    def __init__(self, args: ModelArgs, layer_id: int) -> None: ...
-    def __call__(self, x: mx.array, input_ids: mx.array) -> mx.array: ...
-
-class DeepseekV4Block(nn.Module):
-    attn_norm: nn.RMSNorm
-    attn: V4Attention
-    hc_attn: HyperConnection
-    ffn_norm: nn.RMSNorm
-    ffn: DeepseekV4MoE
-    hc_ffn: HyperConnection
-
-    def __init__(self, args: ModelArgs, layer_id: int) -> None: ...
-    def __call__(
-        self,
-        h: mx.array,
-        cache: Optional[Any],
-        input_ids: mx.array,
-    ) -> mx.array: ...
-
-class DeepseekV4Model(nn.Module):
-    args: ModelArgs
-    vocab_size: int
-    embed_tokens: nn.Embedding
-    layers: list[DeepseekV4Block]
-    norm: nn.RMSNorm
-    hc_head: HyperHead
-
-    def __init__(self, args: ModelArgs) -> None: ...
-    def __call__(
-        self,
-        inputs: mx.array,
-        cache: Optional[List[Any]] = None,
-    ) -> mx.array: ...
-
-class Model(nn.Module):
-    args: ModelArgs
-    model_type: str
-    model: DeepseekV4Model
-    lm_head: nn.Linear
-
-    def __init__(self, args: ModelArgs) -> None: ...
-    def __call__(
-        self,
-        inputs: mx.array,
-        cache: Optional[List[Any]] = None,
-    ) -> mx.array: ...
-    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
-    def make_cache(self) -> list[RotatingKVCache | DeepseekV4Cache]: ...
-    @property
-    def layers(self) -> list[DeepseekV4Block]: ...

Cargo.toml

@@ -1,6 +1,6 @@
 [workspace]
 resolver = "3"
-members = ["rust/exo_pyo3_bindings", "rust/networking"]
+members = ["rust/networking", "rust/exo_pyo3_bindings", "rust/util"]
 
 [workspace.package]
 version = "0.0.1"
@@ -20,19 +20,30 @@ opt-level = 3
 [workspace.dependencies]
 ## Crate members as common dependencies
 networking = { path = "rust/networking" }
+util = { path = "rust/util" }
 
 # Macro dependecies
 extend = "1.2"
+delegate = "0.13"
+
+# Utility dependencies
+keccak-const = "0.2"
+
+# Async dependencies
+async-stream = "0.3"
 tokio = "1.46"
 futures-lite = "2.6.1"
+futures-timer = "3.0"
+
+# Data structures
+either = "1.15"
 
 # Tracing/logging
 log = "0.4"
-env_logger = "0.11.10"
 
 # networking
-zenoh = "=1.9.0"
-zerompk = "0.4.2"
+libp2p = "0.56"
+libp2p-tcp = "0.44"
 
 [workspace.lints.rust]
 static_mut_refs = "warn"      # Or use "warn" instead of deny

app/EXO/EXO/ContentView.swift

@@ -16,13 +16,22 @@ struct ContentView: View {
     @EnvironmentObject private var updater: SparkleUpdater
     @EnvironmentObject private var thunderboltBridgeService: ThunderboltBridgeService
     @EnvironmentObject private var settingsWindowController: SettingsWindowController
-    @EnvironmentObject private var bugReportWindowController: BugReportWindowController
     @State private var focusedNode: NodeViewModel?
     @State private var deletingInstanceIDs: Set<String> = []
     @State private var showAllNodes = false
     @State private var showAllInstances = false
     @State private var baseURLCopied = false
     @State private var showAdvanced = false
+    @State private var showDebugInfo = false
+    private enum BugReportPhase: Equatable {
+        case idle
+        case prompting
+        case sending(String)
+        case success(String)
+        case failure(String)
+    }
+    @State private var bugReportPhase: BugReportPhase = .idle
+    @State private var bugReportUserDescription: String = ""
     @State private var uninstallInProgress = false
     @State private var pendingNamespace: String = ""
     @State private var pendingHFToken: String = ""
@@ -285,13 +294,6 @@ struct ContentView: View {
             ) {
                 updater.checkForUpdates()
             }
-            HoverButton(
-                title: "Share Bug Report…",
-                tint: .primary,
-                trailingSystemImage: "ladybug"
-            ) {
-                bugReportWindowController.open()
-            }
             .padding(.bottom, 8)
             HoverButton(title: "Quit", tint: .secondary) {
                 controller.stop()
@@ -475,6 +477,40 @@ struct ContentView: View {
         }
     }
 
+    private var debugSection: some View {
+        VStack(alignment: .leading, spacing: 4) {
+            HoverButton(
+                title: "Debug Info",
+                tint: .primary,
+                trailingSystemImage: showDebugInfo ? "chevron.up" : "chevron.down",
+                small: true
+            ) {
+                showDebugInfo.toggle()
+            }
+            if showDebugInfo {
+                VStack(alignment: .leading, spacing: 4) {
+                    Text("Version: \(buildTag)")
+                        .font(.caption2)
+                        .foregroundColor(.secondary)
+                    Text("Commit: \(buildCommit)")
+                        .font(.caption2)
+                        .foregroundColor(.secondary)
+                    Text(thunderboltStatusText)
+                        .font(.caption2)
+                        .foregroundColor(thunderboltStatusColor)
+                    clusterThunderboltBridgeView
+                    interfaceIpList
+                    rdmaStatusView
+                    sendBugReportButton
+                        .padding(.top, 6)
+                }
+                .padding(.leading, 8)
+                .transition(.opacity)
+            }
+        }
+        .animation(.easeInOut(duration: 0.25), value: showDebugInfo)
+    }
+
     private var rdmaStatusView: some View {
         let rdmaStatuses = stateService.latestSnapshot?.nodeRdmaCtl ?? [:]
         let localNodeId = stateService.localNodeId
@@ -523,6 +559,127 @@ struct ContentView: View {
         }
     }
 
+    private var sendBugReportButton: some View {
+        VStack(alignment: .leading, spacing: 6) {
+            switch bugReportPhase {
+            case .idle:
+                Button {
+                    bugReportPhase = .prompting
+                    bugReportUserDescription = ""
+                } label: {
+                    HStack {
+                        Text("Send Bug Report")
+                            .font(.caption)
+                            .fontWeight(.semibold)
+                        Spacer()
+                    }
+                    .padding(.vertical, 6)
+                    .padding(.horizontal, 8)
+                    .background(
+                        RoundedRectangle(cornerRadius: 6)
+                            .fill(Color.accentColor.opacity(0.12))
+                    )
+                }
+                .buttonStyle(.plain)
+
+            case .prompting:
+                VStack(alignment: .leading, spacing: 6) {
+                    VStack(alignment: .leading, spacing: 2) {
+                        Text("Tell us what went wrong (optional)")
+                            .font(.caption2)
+                            .foregroundColor(.secondary)
+                        Text(
+                            "A quick description of what you were doing and what happened helps us track down the bug for you."
+                        )
+                        .font(.caption2)
+                        .foregroundColor(.secondary)
+                        .opacity(0.8)
+                        .fixedSize(horizontal: false, vertical: true)
+                    }
+                    TextEditor(text: $bugReportUserDescription)
+                        .font(.caption2)
+                        .frame(height: 60)
+                        .overlay(
+                            RoundedRectangle(cornerRadius: 4)
+                                .stroke(Color.secondary.opacity(0.3), lineWidth: 1)
+                        )
+                    HStack(spacing: 8) {
+                        Button("Send") {
+                            Task {
+                                await sendBugReport()
+                            }
+                        }
+                        .font(.caption2)
+                        .buttonStyle(.borderedProminent)
+                        .controlSize(.small)
+                        Button("Cancel") {
+                            bugReportPhase = .idle
+                        }
+                        .font(.caption2)
+                        .buttonStyle(.bordered)
+                        .controlSize(.small)
+                    }
+                }
+                .padding(8)
+                .background(
+                    RoundedRectangle(cornerRadius: 6)
+                        .fill(Color.accentColor.opacity(0.06))
+                )
+
+            case .sending(let message):
+                HStack(spacing: 6) {
+                    ProgressView()
+                        .scaleEffect(0.6)
+                    Text(message)
+                        .font(.caption2)
+                        .foregroundColor(.secondary)
+                }
+
+            case .success(let message):
+                VStack(alignment: .leading, spacing: 6) {
+                    Text(message)
+                        .font(.caption2)
+                        .foregroundColor(.secondary)
+                        .fixedSize(horizontal: false, vertical: true)
+                    Button {
+                        openGitHubIssue()
+                    } label: {
+                        HStack(spacing: 4) {
+                            Image(systemName: "arrow.up.right.square")
+                                .imageScale(.small)
+                            Text("Create GitHub Issue")
+                                .font(.caption2)
+                        }
+                    }
+                    .buttonStyle(.bordered)
+                    .controlSize(.small)
+                    Button("Done") {
+                        bugReportPhase = .idle
+                        bugReportUserDescription = ""
+                    }
+                    .font(.caption2)
+                    .buttonStyle(.plain)
+                    .foregroundColor(.secondary)
+                }
+
+            case .failure(let message):
+                VStack(alignment: .leading, spacing: 4) {
+                    Text(message)
+                        .font(.caption2)
+                        .foregroundColor(.red)
+                        .fixedSize(horizontal: false, vertical: true)
+                    Button("Dismiss") {
+                        bugReportPhase = .idle
+                    }
+                    .font(.caption2)
+                    .buttonStyle(.plain)
+                    .foregroundColor(.secondary)
+                }
+            }
+        }
+        .animation(.easeInOut(duration: 0.2), value: bugReportPhase)
+    }
+
     private var processToggleBinding: Binding<Bool> {
         Binding(
             get: {
@@ -563,6 +720,61 @@ struct ContentView: View {
         )
     }
 
+    private func sendBugReport() async {
+        bugReportPhase = .sending("Collecting logs...")
+        let service = BugReportService()
+        let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines)
+        do {
+            let outcome = try await service.sendReport(
+                isManual: true,
+                userDescription: description.isEmpty ? nil : description
+            )
+            if outcome.success {
+                bugReportPhase = .success(outcome.message)
+            } else {
+                bugReportPhase = .failure(outcome.message)
+            }
+        } catch {
+            bugReportPhase = .failure(error.localizedDescription)
+        }
+    }
+
+    private func openGitHubIssue() {
+        let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines)
+
+        var bodyParts: [String] = []
+        bodyParts.append("## Describe the bug")
+        bodyParts.append("")
+        if !description.isEmpty {
+            bodyParts.append(description)
+        } else {
+            bodyParts.append("A clear and concise description of what the bug is.")
+        }
+        bodyParts.append("")
+        bodyParts.append("## Environment")
+        bodyParts.append("")
+        bodyParts.append("- macOS Version: \(ProcessInfo.processInfo.operatingSystemVersionString)")
+        bodyParts.append("- EXO Version: \(buildTag) (\(buildCommit))")
+        bodyParts.append("")
+        bodyParts.append("## Additional context")
+        bodyParts.append("")
+        bodyParts.append("A bug report with diagnostic logs was submitted via the app.")
+
+        let body = bodyParts.joined(separator: "\n")
+
+        var components = URLComponents(string: "https://github.com/exo-explore/exo/issues/new")!
+        components.queryItems = [
+            URLQueryItem(name: "template", value: "bug_report.md"),
+            URLQueryItem(name: "title", value: "[BUG] "),
+            URLQueryItem(name: "body", value: body),
+            URLQueryItem(name: "labels", value: "bug"),
+        ]
+
+        if let url = components.url {
+            NSWorkspace.shared.open(url)
+        }
+    }
+
     private func showUninstallConfirmationAlert() {
         let alert = NSAlert()
         alert.messageText = "Uninstall EXO"
@@ -645,6 +857,13 @@ struct ContentView: View {
         }
     }
 
+    private var buildTag: String {
+        Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown"
+    }
+
+    private var buildCommit: String {
+        Bundle.main.infoDictionary?["EXOBuildCommit"] as? String ?? "unknown"
+    }
 }
 
 private struct HoverButton: View {

app/EXO/EXO/EXOApp.swift

@@ -22,7 +22,6 @@ struct EXOApp: App {
     @StateObject private var updater: SparkleUpdater
     @StateObject private var thunderboltBridgeService: ThunderboltBridgeService
     @StateObject private var settingsWindowController: SettingsWindowController
-    @StateObject private var bugReportWindowController: BugReportWindowController
     private let terminationObserver: TerminationObserver
     private let firstLaunchPopout = FirstLaunchPopout()
     private let ciContext = CIContext(options: nil)
@@ -47,7 +46,6 @@ struct EXOApp: App {
         let thunderboltBridge = ThunderboltBridgeService(clusterStateService: service)
         _thunderboltBridgeService = StateObject(wrappedValue: thunderboltBridge)
         _settingsWindowController = StateObject(wrappedValue: SettingsWindowController())
-        _bugReportWindowController = StateObject(wrappedValue: BugReportWindowController())
         enableLaunchAtLoginIfNeeded()
         // Install LaunchDaemon to disable Thunderbolt Bridge on startup (prevents network loops)
         NetworkSetupHelper.promptAndInstallIfNeeded()
@@ -68,7 +66,6 @@ struct EXOApp: App {
                 .environmentObject(updater)
                 .environmentObject(thunderboltBridgeService)
                 .environmentObject(settingsWindowController)
-                .environmentObject(bugReportWindowController)
         } label: {
             menuBarIcon
                 .onReceive(controller.$isFirstLaunchReady) { ready in

app/EXO/EXO/Services/ClusterStateService.swift

@@ -17,7 +17,7 @@ final class ClusterStateService: ObservableObject {
 
     init(
         baseURL: URL = URL(string: "http://127.0.0.1:52415")!,
-        session: URLSession = ClusterStateService.makeNonCachingSession()
+        session: URLSession = .shared
     ) {
         self.baseURL = baseURL
         self.endpoint = baseURL.appendingPathComponent("state")
@@ -27,23 +27,6 @@ final class ClusterStateService: ObservableObject {
         self.decoder = decoder
     }
 
-    /// `URLSession.shared` carries an on-disk `URLCache` that persists every
-    /// response body under `~/Library/Caches/exolabs.EXO/`. We poll `/state`
-    /// at 2 Hz from `startPolling`, so leaving the shared cache attached
-    /// dirties ~500–620 KB/sec of file-backed memory and trips macOS's
-    /// per-process `disk writes` resource limit (microstackshot reports
-    /// observed on M3 Ultra producing GBs of cached responses per hour).
-    /// Cluster-state polling responses are time-sensitive and small; they
-    /// gain nothing from being cached on disk. Use an ephemeral session
-    /// with `urlCache = nil` so neither response bodies nor metadata
-    /// touch disk.
-    private static func makeNonCachingSession() -> URLSession {
-        let config = URLSessionConfiguration.ephemeral
-        config.urlCache = nil
-        config.requestCachePolicy = .reloadIgnoringLocalCacheData
-        return URLSession(configuration: config)
-    }
-
     func startPolling(interval: TimeInterval = 0.5) {
         stopPolling()
         Task {

app/EXO/EXO/Views/BugReportWindowController.swift

@@ -1,242 +0,0 @@
-import AppKit
-import SwiftUI
-
-/// Manages a standalone window for the bug-report flow.
-/// Ensures only one instance exists and brings it to front on repeated opens.
-@MainActor
-final class BugReportWindowController: ObservableObject {
-    private var window: NSWindow?
-
-    func open() {
-        if let existing = window, existing.isVisible {
-            existing.makeKeyAndOrderFront(nil)
-            NSApp.activate()
-            return
-        }
-
-        let view = BugReportView(onDismiss: { [weak self] in
-            self?.window?.close()
-        })
-
-        let hostingController = NSHostingController(rootView: view)
-        hostingController.sizingOptions = [.preferredContentSize, .minSize]
-
-        let newWindow = NSWindow(contentViewController: hostingController)
-        newWindow.styleMask = [.titled, .closable, .resizable]
-        newWindow.title = "Send a Bug Report"
-        newWindow.center()
-        newWindow.setFrameAutosaveName("ExoBugReportWindow")
-        newWindow.isReleasedWhenClosed = false
-        newWindow.makeKeyAndOrderFront(nil)
-        NSApp.activate()
-
-        window = newWindow
-    }
-}
-
-private struct BugReportView: View {
-    fileprivate enum Phase: Equatable {
-        case prompting
-        case sending(String)
-        case success(String)
-        case failure(String)
-    }
-
-    let onDismiss: () -> Void
-
-    @State private var phase: Phase = .prompting
-    @State private var userDescription: String = ""
-    @FocusState private var descriptionFocused: Bool
-
-    var body: some View {
-        VStack(alignment: .leading, spacing: 12) {
-            switch phase {
-            case .prompting:
-                promptingView
-            case .sending(let message):
-                sendingView(message: message)
-            case .success(let message):
-                successView(message: message)
-            case .failure(let message):
-                failureView(message: message)
-            }
-        }
-        .padding(16)
-        .frame(minWidth: 380)
-        .animation(.easeInOut(duration: 0.2), value: phase)
-        .onAppear { descriptionFocused = true }
-    }
-
-    private var promptingView: some View {
-        VStack(alignment: .leading, spacing: 8) {
-            Text("Description (optional)")
-                .font(.subheadline)
-                .foregroundColor(.secondary)
-            ZStack(alignment: .topLeading) {
-                if userDescription.isEmpty {
-                    Text("What were you doing when it broke?")
-                        .font(.body)
-                        .foregroundColor(Color(nsColor: .placeholderTextColor))
-                        .padding(.horizontal, 10)
-                        .padding(.vertical, 8)
-                        .allowsHitTesting(false)
-                }
-                TextEditor(text: $userDescription)
-                    .font(.body)
-                    .scrollContentBackground(.hidden)
-                    .padding(4)
-                    .frame(height: 72)
-                    .focused($descriptionFocused)
-            }
-            .background(
-                RoundedRectangle(cornerRadius: 6)
-                    .fill(Color(nsColor: .textBackgroundColor))
-            )
-            .overlay(
-                RoundedRectangle(cornerRadius: 6)
-                    .strokeBorder(Color(nsColor: .separatorColor), lineWidth: 1)
-            )
-
-            Text("Diagnostic logs will be uploaded with your report.")
-                .font(.caption)
-                .foregroundColor(.secondary)
-
-            HStack {
-                Spacer()
-                Button("Cancel") { onDismiss() }
-                    .keyboardShortcut(.cancelAction)
-                Button("Send") {
-                    Task { await send() }
-                }
-                .keyboardShortcut(.defaultAction)
-            }
-            .padding(.top, 4)
-        }
-    }
-
-    private func sendingView(message: String) -> some View {
-        VStack(alignment: .leading, spacing: 12) {
-            HStack(spacing: 10) {
-                ProgressView().controlSize(.small)
-                Text(message)
-                    .foregroundColor(.secondary)
-            }
-            HStack {
-                Spacer()
-                Button("Cancel") { onDismiss() }
-                    .keyboardShortcut(.cancelAction)
-                    .disabled(true)
-                Button("Send") {}
-                    .disabled(true)
-            }
-        }
-    }
-
-    private func successView(message: String) -> some View {
-        VStack(alignment: .leading, spacing: 12) {
-            HStack(alignment: .top, spacing: 10) {
-                Image(systemName: "checkmark.circle.fill")
-                    .foregroundColor(.green)
-                    .font(.title2)
-                Text(message)
-                    .fixedSize(horizontal: false, vertical: true)
-            }
-            HStack {
-                Button {
-                    openGitHubIssue()
-                } label: {
-                    HStack(spacing: 4) {
-                        Image(systemName: "arrow.up.right.square")
-                        Text("Open GitHub Issue")
-                    }
-                }
-                Spacer()
-                Button("Done") { onDismiss() }
-                    .keyboardShortcut(.defaultAction)
-            }
-        }
-    }
-
-    private func failureView(message: String) -> some View {
-        VStack(alignment: .leading, spacing: 12) {
-            HStack(alignment: .top, spacing: 10) {
-                Image(systemName: "exclamationmark.triangle.fill")
-                    .foregroundColor(.orange)
-                    .font(.title2)
-                Text(message)
-                    .fixedSize(horizontal: false, vertical: true)
-            }
-            HStack {
-                Spacer()
-                Button("Try Again") {
-                    phase = .prompting
-                }
-                Button("Close") { onDismiss() }
-                    .keyboardShortcut(.defaultAction)
-            }
-        }
-    }
-
-    private func send() async {
-        phase = .sending("Collecting logs and uploading…")
-        let service = BugReportService()
-        let description = userDescription.trimmingCharacters(in: .whitespacesAndNewlines)
-        do {
-            let outcome = try await service.sendReport(
-                isManual: true,
-                userDescription: description.isEmpty ? nil : description
-            )
-            if outcome.success {
-                phase = .success(outcome.message)
-            } else {
-                phase = .failure(outcome.message)
-            }
-        } catch {
-            phase = .failure(error.localizedDescription)
-        }
-    }
-
-    private func openGitHubIssue() {
-        let description = userDescription.trimmingCharacters(in: .whitespacesAndNewlines)
-
-        var bodyParts: [String] = []
-        bodyParts.append("## Describe the bug")
-        bodyParts.append("")
-        if !description.isEmpty {
-            bodyParts.append(description)
-        } else {
-            bodyParts.append("A clear and concise description of what the bug is.")
-        }
-        bodyParts.append("")
-        bodyParts.append("## Environment")
-        bodyParts.append("")
-        bodyParts.append("- macOS Version: \(ProcessInfo.processInfo.operatingSystemVersionString)")
-        bodyParts.append("- EXO Version: \(buildTag) (\(buildCommit))")
-        bodyParts.append("")
-        bodyParts.append("## Additional context")
-        bodyParts.append("")
-        bodyParts.append("A bug report with diagnostic logs was submitted via the app.")
-
-        let body = bodyParts.joined(separator: "\n")
-
-        var components = URLComponents(string: "https://github.com/exo-explore/exo/issues/new")!
-        components.queryItems = [
-            URLQueryItem(name: "template", value: "bug_report.md"),
-            URLQueryItem(name: "title", value: "[BUG] "),
-            URLQueryItem(name: "body", value: body),
-            URLQueryItem(name: "labels", value: "bug"),
-        ]
-
-        if let url = components.url {
-            NSWorkspace.shared.open(url)
-        }
-    }
-
-    private var buildTag: String {
-        Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown"
-    }
-
-    private var buildCommit: String {
-        Bundle.main.infoDictionary?["EXOBuildCommit"] as? String ?? "unknown"
-    }
-}

app/EXO/EXO/Views/SettingsView.swift

@@ -21,6 +21,8 @@ struct SettingsView: View {
     @State private var pendingReadOnlyModelsDirs: String = ""
     @State private var pendingCustomEnvironmentVariables: [CustomEnvironmentVariable] = []
     @State private var needsRestart = false
+    @State private var bugReportInFlight = false
+    @State private var bugReportMessage: String?
     @State private var uninstallInProgress = false
 
     var body: some View {
@@ -200,6 +202,8 @@ struct SettingsView: View {
                 VStack(alignment: .leading, spacing: 2) {
                     rdmaStatusView
                 }
+
+                sendBugReportButton
             }
 
             Section("Danger Zone") {
@@ -500,30 +504,63 @@ struct SettingsView: View {
         }
     }
 
+    private var sendBugReportButton: some View {
+        VStack(alignment: .leading, spacing: 4) {
+            Button {
+                Task {
+                    await sendBugReport()
+                }
+            } label: {
+                HStack {
+                    if bugReportInFlight {
+                        ProgressView()
+                            .scaleEffect(0.6)
+                    }
+                    Text("Send Bug Report")
+                        .font(.caption)
+                        .fontWeight(.semibold)
+                    Spacer()
+                }
+            }
+            .disabled(bugReportInFlight)
+
+            if let message = bugReportMessage {
+                Text(message)
+                    .font(.caption2)
+                    .foregroundColor(.secondary)
+                    .fixedSize(horizontal: false, vertical: true)
+            }
+        }
+    }
+
     // MARK: - Actions
 
+    private func sendBugReport() async {
+        bugReportInFlight = true
+        bugReportMessage = "Collecting logs..."
+        let service = BugReportService()
+        do {
+            let outcome = try await service.sendReport(isManual: true)
+            bugReportMessage = outcome.message
+        } catch {
+            bugReportMessage = error.localizedDescription
+        }
+        bugReportInFlight = false
+    }
+
     private func showUninstallConfirmationAlert() {
         let alert = NSAlert()
         alert.messageText = "Uninstall EXO"
         alert.informativeText = """
-            This will remove EXO and all its components:
+            This will remove EXO and all its system components:
 
             • Network configuration daemon
             • Launch at login registration
             • EXO network location
-            • EXO data directory (~/.exo)
 
             The app will be moved to Trash.
             """
         alert.alertStyle = .warning
-
-        let checkbox = NSButton(
-            checkboxWithTitle: "Keep downloaded models (~/.exo/models)",
-            target: nil, action: nil)
-        checkbox.state = .off
-        checkbox.sizeToFit()
-        alert.accessoryView = checkbox
-
         alert.addButton(withTitle: "Uninstall")
         alert.addButton(withTitle: "Cancel")
 
@@ -533,11 +570,11 @@ struct SettingsView: View {
 
         let response = alert.runModal()
         if response == .alertFirstButtonReturn {
-            performUninstall(keepModels: checkbox.state == .on)
+            performUninstall()
         }
     }
 
-    private func performUninstall(keepModels: Bool) {
+    private func performUninstall() {
         uninstallInProgress = true
 
         controller.cancelPendingLaunch()
@@ -547,7 +584,6 @@ struct SettingsView: View {
         DispatchQueue.global(qos: .utility).async {
             do {
                 try NetworkSetupHelper.uninstall()
-                try Self.removeExoDirectory(keepModels: keepModels)
 
                 DispatchQueue.main.async {
                     LaunchAtLoginHelper.disable()
@@ -571,23 +607,6 @@ struct SettingsView: View {
         }
     }
 
-    private static func removeExoDirectory(keepModels: Bool) throws {
-        let fm = FileManager.default
-        let exoDir = ExoProcessController.exoDirectoryURL
-        guard fm.fileExists(atPath: exoDir.path) else { return }
-
-        if !keepModels {
-            try fm.removeItem(at: exoDir)
-            return
-        }
-
-        let contents = try fm.contentsOfDirectory(
-            at: exoDir, includingPropertiesForKeys: nil, options: [])
-        for entry in contents where entry.lastPathComponent != "models" {
-            try? fm.removeItem(at: entry)
-        }
-    }
-
     private func moveAppToTrash() {
         guard let appURL = Bundle.main.bundleURL as URL? else { return }
         do {

app/EXO/uninstall-exo.sh

@@ -3,55 +3,25 @@
 # EXO Uninstaller Script
 #
 # This script removes all EXO system components that persist after deleting the app.
-# Run with: sudo ./uninstall-exo.sh [--keep-models]
-#
-# Options:
-#   --keep-models   Preserve ~/.exo/models when removing the EXO data directory.
+# Run with: sudo ./uninstall-exo.sh
 #
 # Components removed:
 # - LaunchDaemon: /Library/LaunchDaemons/io.exo.networksetup.plist
 # - Network script: /Library/Application Support/EXO/
 # - Log files: /var/log/io.exo.networksetup.*
 # - Network location: "exo"
-# - EXO data directory: ~/.exo (or all of ~/.exo except models/ when --keep-models is set)
 # - Launch at login registration
 #
 
 set -euo pipefail
 
-KEEP_MODELS=0
-for arg in "$@"; do
-  case "$arg" in
-  --keep-models)
-    KEEP_MODELS=1
-    ;;
-  -h | --help)
-    echo "Usage: sudo ./uninstall-exo.sh [--keep-models]"
-    echo "  --keep-models   Preserve ~/.exo/models when removing the EXO data directory."
-    exit 0
-    ;;
-  *)
-    echo "Unknown argument: $arg" >&2
-    echo "Usage: sudo ./uninstall-exo.sh [--keep-models]" >&2
-    exit 2
-    ;;
-  esac
-done
-
 LABEL="io.exo.networksetup"
-# Current script path. Older installs used a different filename; keep the
-# legacy path here so a fresh uninstall still cleans up upgraded machines.
-CURRENT_SCRIPT_DEST="/Library/Application Support/EXO/disable_bridge.sh"
-LEGACY_SCRIPT_DEST="/Library/Application Support/EXO/disable_bridge_enable_dhcp.sh"
+SCRIPT_DEST="/Library/Application Support/EXO/disable_bridge_enable_dhcp.sh"
 PLIST_DEST="/Library/LaunchDaemons/io.exo.networksetup.plist"
 LOG_OUT="/var/log/${LABEL}.log"
 LOG_ERR="/var/log/${LABEL}.err.log"
 APP_BUNDLE_ID="io.exo.EXO"
 
-# Resolve the invoking user's home, even when run via sudo.
-USER_HOME="$(eval echo "~${SUDO_USER:-$USER}")"
-EXO_DIR="$USER_HOME/.exo"
-
 # Colors for output
 RED='\033[0;31m'
 GREEN='\033[0;32m'
@@ -99,17 +69,11 @@ else
   echo_warn "LaunchDaemon plist not found (already removed?)"
 fi
 
-# Remove the script (current and legacy filenames) — backwards-compatible:
-# tolerate either, both, or neither being present.
-removed_any_script=0
-for script in "$CURRENT_SCRIPT_DEST" "$LEGACY_SCRIPT_DEST"; do
-  if [[ -f $script ]]; then
-    rm -f "$script"
-    echo_info "Removed network setup script: $script"
-    removed_any_script=1
-  fi
-done
-if [[ $removed_any_script -eq 0 ]]; then
+# Remove the script and parent directory
+if [[ -f $SCRIPT_DEST ]]; then
+  rm -f "$SCRIPT_DEST"
+  echo_info "Removed network setup script"
+else
   echo_warn "Network setup script not found (already removed?)"
 fi
 
@@ -151,22 +115,6 @@ if networksetup -listnetworkservices 2>/dev/null | grep -q "Thunderbolt Bridge";
   echo_info "Re-enabled Thunderbolt Bridge"
 fi
 
-# Remove EXO data directory (~/.exo)
-EXO_DIR_REMOVED=""
-if [[ -d $EXO_DIR ]]; then
-  if [[ $KEEP_MODELS == "1" && -d "$EXO_DIR/models" ]]; then
-    find "$EXO_DIR" -mindepth 1 -maxdepth 1 ! -name models -exec rm -rf {} +
-    EXO_DIR_REMOVED="kept_models"
-    echo_info "Removed ~/.exo (preserved models/)"
-  else
-    rm -rf "$EXO_DIR"
-    EXO_DIR_REMOVED="full"
-    echo_info "Removed ~/.exo"
-  fi
-else
-  echo_warn "~/.exo not found (already removed?)"
-fi
-
 # Note about launch at login registration
 # SMAppService-based login items cannot be removed from a shell script.
 # They can only be unregistered from within the app itself or manually via System Settings.
@@ -196,10 +144,6 @@ echo "  • Network setup LaunchDaemon"
 echo "  • Network configuration script"
 echo "  • Log files"
 echo "  • 'exo' network location"
-case "$EXO_DIR_REMOVED" in
-full) echo "  • EXO data directory (~/.exo)" ;;
-kept_models) echo "  • EXO data directory (~/.exo, models preserved)" ;;
-esac
 echo ""
 echo "Your network has been restored to use the 'Automatic' location."
 echo "Thunderbolt Bridge has been re-enabled (if present)."

bench/eval_configs/models.toml

@@ -7,7 +7,7 @@
 #   name, patterns, reasoning
 #
 # Optional per-model overrides (CLI flags take priority over these):
-#   temperature, top_p, max_tokens, reasoning_effort, enable_thinking
+#   temperature, top_p, max_tokens, reasoning_effort
 #
 # Fallback defaults (when no per-model config):
 #   reasoning:     temperature=1.0, max_tokens=131072, reasoning_effort="high"
@@ -18,9 +18,10 @@
 
 # ─── Qwen3.5 (Feb 2026) ─────────────────────────────────────────────
 # Source: HuggingFace model cards (Qwen/Qwen3.5-*)
-# Model card recommends: temp=0.6, top_p=0.95, top_k=20
-# We omit top_k to match vllm eval (which doesn't set it).
-# max_tokens=121072 to match vllm eval (131072 context - 10000 safety margin).
+# 35B-A3B thinking general: temp=1.0, top_p=0.95, top_k=20
+# 397B thinking: temp=0.6, top_p=0.95, top_k=20
+# Non-thinking: temp=0.7, top_p=0.8, top_k=20
+# max_tokens: 32768 general, 81920 for complex math/code
 
 [[model]]
 name = "Qwen3.5 2B"
@@ -28,8 +29,7 @@ patterns = ["Qwen3.5-2B"]
 reasoning = true
 temperature = 0.6
 top_p = 0.95
-enable_thinking = true
-max_tokens = 121072
+max_tokens = 81920
 
 [[model]]
 name = "Qwen3.5 9B"
@@ -37,8 +37,7 @@ patterns = ["Qwen3.5-9B"]
 reasoning = true
 temperature = 0.6
 top_p = 0.95
-enable_thinking = true
-max_tokens = 121072
+max_tokens = 81920
 
 [[model]]
 name = "Qwen3.5 27B"
@@ -46,17 +45,15 @@ patterns = ["Qwen3.5-27B"]
 reasoning = true
 temperature = 0.6
 top_p = 0.95
-enable_thinking = true
-max_tokens = 121072
+max_tokens = 81920
 
 [[model]]
 name = "Qwen3.5 35B A3B"
 patterns = ["Qwen3.5-35B-A3B"]
 reasoning = true
-temperature = 0.6
+temperature = 1.0
 top_p = 0.95
-enable_thinking = true
-max_tokens = 121072
+max_tokens = 81920
 
 [[model]]
 name = "Qwen3.5 122B A10B"
@@ -64,8 +61,7 @@ patterns = ["Qwen3.5-122B-A10B"]
 reasoning = true
 temperature = 0.6
 top_p = 0.95
-enable_thinking = true
-max_tokens = 121072
+max_tokens = 81920
 
 [[model]]
 name = "Qwen3.5 397B A17B"
@@ -73,14 +69,12 @@ patterns = ["Qwen3.5-397B-A17B"]
 reasoning = true
 temperature = 0.6
 top_p = 0.95
-enable_thinking = true
-max_tokens = 121072
+max_tokens = 81920
 
 # ─── Qwen3 (Apr 2025) ───────────────────────────────────────────────
 # Source: HuggingFace model cards (Qwen/Qwen3-*)
-# Model card recommends: temp=0.6, top_p=0.95, top_k=20
-# We omit top_k to match vllm eval (which doesn't set it).
-# Non-thinking: temp=0.7, top_p=0.8
+# Thinking: temp=0.6, top_p=0.95, top_k=20
+# Non-thinking: temp=0.7, top_p=0.8, top_k=20
 # max_tokens: 32768 general, 38912 for complex math/code
 
 [[model]]
@@ -89,7 +83,6 @@ patterns = ["Qwen3-0.6B"]
 reasoning = true
 temperature = 0.6
 top_p = 0.95
-enable_thinking = true
 max_tokens = 38912
 
 [[model]]
@@ -98,7 +91,6 @@ patterns = ["Qwen3-30B-A3B"]
 reasoning = true
 temperature = 0.6
 top_p = 0.95
-enable_thinking = true
 max_tokens = 38912
 
 [[model]]
@@ -107,7 +99,6 @@ patterns = ["Qwen3-235B-A22B"]
 reasoning = true
 temperature = 0.6
 top_p = 0.95
-enable_thinking = true
 max_tokens = 38912
 
 [[model]]
@@ -116,7 +107,6 @@ patterns = ["Qwen3-Next-80B-A3B-Thinking"]
 reasoning = true
 temperature = 0.6
 top_p = 0.95
-enable_thinking = true
 max_tokens = 38912
 
 [[model]]
@@ -139,9 +129,9 @@ max_tokens = 16384
 name = "Qwen3 Coder Next"
 patterns = ["Qwen3-Coder-Next"]
 reasoning = false
-temperature = 1.0
-top_p = 0.95
-max_tokens = 121072
+temperature = 0.7
+top_p = 0.8
+max_tokens = 16384
 
 # ─── GPT-OSS (OpenAI) ───────────────────────────────────────────────
 # Source: OpenAI GitHub README + HuggingFace discussion #21
@@ -175,38 +165,10 @@ patterns = ["DeepSeek-V3.1"]
 reasoning = true
 temperature = 0.0
 
-[[model]]
-name = "DeepSeek V3.2"
-patterns = ["DeepSeek-V3.2"]
-reasoning = true
-temperature = 1.0
-top_p = 0.95
-enable_thinking = true
-
-# ─── NVIDIA Nemotron ───────────────────────────────────────────────────
-# Source: HuggingFace model cards
-# All variants: temp=1.0, top_p=0.95, enable_thinking=true
-
-[[model]]
-name = "Nemotron Cascade 2 30B A3B"
-patterns = ["Nemotron-Cascade-2-30B-A3B"]
-reasoning = true
-temperature = 1.0
-top_p = 0.95
-enable_thinking = true
-
-[[model]]
-name = "Nemotron 3 Super 120B A12B"
-patterns = ["Nemotron-3-Super-120B-A12B", "NVIDIA-Nemotron-3-Super-120B-A12B"]
-reasoning = true
-temperature = 1.0
-top_p = 0.95
-enable_thinking = true
-
 # ─── GLM (ZhipuAI / THUDM) ──────────────────────────────────────────
 # Source: HuggingFace model cards + generation_config.json + docs.z.ai
 # GLM 4.5+: temp=1.0, top_p=0.95
-# max_tokens=121072 to match vllm eval (131072 context - 10000 safety margin)
+# Reasoning tasks: 131072 max_tokens; coding/SWE tasks: temp=0.7
 
 [[model]]
 name = "GLM-5"
@@ -214,8 +176,7 @@ patterns = ["GLM-5"]
 reasoning = true
 temperature = 1.0
 top_p = 0.95
-enable_thinking = true
-max_tokens = 121072
+max_tokens = 131072
 
 [[model]]
 name = "GLM 4.5 Air"
@@ -230,8 +191,7 @@ patterns = ["GLM-4.7-"]
 reasoning = true
 temperature = 1.0
 top_p = 0.95
-enable_thinking = true
-max_tokens = 121072
+max_tokens = 131072
 # Note: matches both GLM-4.7 and GLM-4.7-Flash
 
 # ─── Kimi (Moonshot AI) ─────────────────────────────────────────────
@@ -253,8 +213,7 @@ patterns = ["Kimi-K2.5"]
 reasoning = true
 temperature = 1.0
 top_p = 0.95
-enable_thinking = true
-max_tokens = 121072
+max_tokens = 131072
 
 [[model]]
 name = "Kimi K2 Instruct"
@@ -264,17 +223,7 @@ temperature = 0.6
 
 # ─── MiniMax ─────────────────────────────────────────────────────────
 # Source: HuggingFace model cards + generation_config.json
-# All models: temp=1.0, top_p=0.95
-# max_tokens=90000 to match vllm eval (100000 context - 10000 safety margin)
-
-[[model]]
-name = "MiniMax M2.7"
-patterns = ["MiniMax-M2.7"]
-reasoning = true
-temperature = 1.0
-top_p = 0.95
-enable_thinking = true
-max_tokens = 90000
+# All models: temp=1.0, top_p=0.95, top_k=40
 
 [[model]]
 name = "MiniMax M2.5"
@@ -282,8 +231,6 @@ patterns = ["MiniMax-M2.5"]
 reasoning = true
 temperature = 1.0
 top_p = 0.95
-enable_thinking = true
-max_tokens = 90000
 
 [[model]]
 name = "MiniMax M2.1"
@@ -304,8 +251,6 @@ patterns = ["Step-3.5-Flash"]
 reasoning = true
 temperature = 1.0
 top_p = 0.95
-enable_thinking = true
-max_tokens = 121072
 
 # ─── Llama (Meta) ───────────────────────────────────────────────────
 # Source: generation_config.json + meta-llama/llama-models generation.py

bench/exo_bench.py

@@ -35,7 +35,6 @@ from harness import (
     ExoHttpError,
     add_common_instance_args,
     capture_cluster_snapshot,
-    find_existing_instance,
     instance_id_from_instance,
     node_ids_from_instance,
     nodes_used_in_instance,
@@ -80,7 +79,7 @@ def load_tokenizer_for_bench(model_id: str) -> Any:
         model_path = Path(
             snapshot_download(
                 model_id,
-                allow_patterns=["*.json", "*.py", "*.tiktoken", "*.model", "*.jinja"],
+                allow_patterns=["*.json", "*.py", "*.tiktoken", "*.model"],
             )
         )
 
@@ -123,48 +122,8 @@ def load_tokenizer_for_bench(model_id: str) -> Any:
 
         return hf_tokenizer
 
-    # TODO: Change back to using only transformers
-    try:
-        return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-    except (AttributeError, ValueError):
-        from huggingface_hub import snapshot_download
-        from transformers import PretrainedConfig
-
-        model_path = Path(
-            snapshot_download(
-                model_id,
-                allow_patterns=[
-                    "*.json",
-                    "*.py",
-                    "tokenizer.model",
-                    "*.tiktoken",
-                    "tiktoken.model",
-                    "*.txt",
-                    "*.jsonl",
-                    "*.jinja",
-                ],
-            )
-        )
-        stub_kwargs: dict[str, Any] = {}
-        config_file = model_path / "config.json"
-        if config_file.exists():
-            with open(config_file) as f:
-                raw = json.load(f)
-            for key in (
-                "model_type",
-                "max_position_embeddings",
-                "vocab_size",
-                "bos_token_id",
-                "eos_token_id",
-                "pad_token_id",
-            ):
-                if key in raw:
-                    stub_kwargs[key] = raw[key]
-        return AutoTokenizer.from_pretrained(
-            str(model_path),
-            config=PretrainedConfig(**stub_kwargs),
-            trust_remote_code=True,
-        )
+    # Default: use AutoTokenizer
+    return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
 
 
 def format_peak_memory(b: float) -> str:
@@ -278,72 +237,28 @@ def run_one_completion(
     prompt_sizer: PromptSizer,
     *,
     use_prefix_cache: bool = False,
-    stream: bool = False,
 ) -> tuple[dict[str, Any], int]:
     content, pp_tokens = prompt_sizer.build(pp_hint)
     payload: dict[str, Any] = {
         "model": model_id,
         "messages": [{"role": "user", "content": content}],
+        "stream": False,
         "max_tokens": tg,
         "logprobs": False,
         "use_prefix_cache": use_prefix_cache,
     }
 
-    if not stream:
-        payload["stream"] = False
-        t0 = time.perf_counter()
-        out = client.post_bench_chat_completions(payload)
-        elapsed = time.perf_counter() - t0
+    t0 = time.perf_counter()
+    out = client.post_bench_chat_completions(payload)
+    elapsed = time.perf_counter() - t0
 
-        stats = out.get("generation_stats")
-        choices = out.get("choices") or [{}]
-        message = choices[0].get("message", {}) if choices else {}
-        content = message.get("content") or ""
-        preview = content[:200] if content else ""
-    else:
-        tokens = 0
-        first_token_time = None
-        t0 = time.perf_counter()
-        text_parts: list[str] = []
-        stats = None
+    stats = out.get("generation_stats")
 
-        for raw_line in client.stream_bench_chat_completions(payload):
-            line = raw_line.strip()
-            if line.startswith(": generation_stats "):
-                with contextlib.suppress(json.JSONDecodeError):
-                    stats = json.loads(line[len(": generation_stats ") :])
-                continue
-            if not line.startswith("data: "):
-                continue
-            data = line[6:]
-            if data == "[DONE]":
-                break
-            try:
-                chunk = json.loads(data)
-                delta = chunk.get("choices", [{}])[0].get("delta", {})
-                if delta.get("content"):
-                    if first_token_time is None:
-                        first_token_time = time.perf_counter()
-                    tokens += 1
-                    text_parts.append(delta["content"])
-            except json.JSONDecodeError:
-                pass
-
-        elapsed = time.perf_counter() - t0
-        preview = "".join(text_parts)[:200]
-
-        if not stats:
-            ttft = (first_token_time - t0) if first_token_time else elapsed
-            gen_time = elapsed - ttft if tokens > 1 else elapsed
-            gen_tps = (tokens - 1) / gen_time if tokens > 1 and gen_time > 0 else 0.0
-            prompt_tps = pp_tokens / ttft if ttft > 0 else 0.0
-            stats = {
-                "prompt_tokens": pp_tokens,
-                "generation_tokens": tokens,
-                "prompt_tps": round(prompt_tps, 2),
-                "generation_tps": round(gen_tps, 2),
-                "peak_memory_usage": {"inBytes": 0},
-            }
+    # Extract preview, handling None content (common for thinking models)
+    choices = out.get("choices") or [{}]
+    message = choices[0].get("message", {}) if choices else {}
+    content = message.get("content") or ""
+    preview = content[:200] if content else ""
 
     return {
         "elapsed_s": elapsed,
@@ -363,19 +278,9 @@ class PromptSizer:
     def _make_counter(tokenizer: Any) -> Callable[[str], int]:
         def count_fn(user_content: str) -> int:
             messages = [{"role": "user", "content": user_content}]
-            try:
-                ids = tokenizer.apply_chat_template(
-                    messages, tokenize=True, add_generation_prompt=True
-                )
-            except ValueError:
-                # Models without a Jinja chat template (e.g. DeepSeek V4 which
-                # ships its own Python encoder). Use the exo-side V4 encoder.
-                from exo.worker.engines.mlx.deepseek_v4_encoding import (
-                    encode_messages as encode_v4,
-                )
-
-                prompt = encode_v4(messages, thinking_mode="thinking")
-                ids = tokenizer.encode(prompt, add_special_tokens=False)
+            ids = tokenizer.apply_chat_template(
+                messages, tokenize=True, add_generation_prompt=True
+            )
             # Fix for transformers 5.x
             if hasattr(ids, "input_ids"):
                 ids = ids.input_ids
@@ -470,11 +375,6 @@ def main() -> int:
         action="store_true",
         help="Force all pp×tg combinations (cartesian product) even when lists have equal length.",
     )
-    ap.add_argument(
-        "--stream",
-        action="store_true",
-        help="Use /bench/chat/completions with streaming SSE response (bench=True still applies: no EOS detection, no KV cache).",
-    )
     ap.add_argument(
         "--no-system-metrics",
         action="store_true",
@@ -540,124 +440,81 @@ def main() -> int:
         logger.error("[exo-bench] tokenizer usable but prompt sizing failed")
         raise
 
-    # Optionally reuse a running instance for this model
-    reused_instance_id: str | None = None
-    if args.reuse_instance:
-        existing = find_existing_instance(client, full_model_id)
-        if existing:
-            reused_instance_id = existing
-            logger.info(f"Reusing existing instance {reused_instance_id}")
-        else:
-            logger.warning(
-                "--reuse-instance: no existing instance found, creating a new one"
-            )
+    selected = settle_and_fetch_placements(
+        client, full_model_id, args, settle_timeout=args.settle_timeout
+    )
 
-    if reused_instance_id is not None:
-        # Use the existing instance directly — skip placement iteration
-        selected = []
-        download_duration_s = None
+    if not selected:
+        logger.error("No valid placements matched your filters.")
+        return 1
+
+    selected.sort(
+        key=lambda p: (
+            str(p.get("instance_meta", "")),
+            str(p.get("sharding", "")),
+            -nodes_used_in_instance(p["instance"]),
+        ),
+        reverse=True,
+    )
+
+    logger.debug(f"exo-bench model: short_id={short_id} full_id={full_model_id}")
+    logger.info(f"placements: {len(selected)}")
+    for p in selected:
+        logger.info(
+            f"  - {p['sharding']} / {p['instance_meta']} / nodes={nodes_used_in_instance(p['instance'])}"
+        )
+
+    if args.dry_run:
+        return 0
+
+    settle_deadline = (
+        time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
+    )
+
+    logger.info("Planning phase: checking downloads...")
+    download_duration_s = run_planning_phase(
+        client,
+        full_model_id,
+        selected[0],
+        args.danger_delete_downloads,
+        args.timeout,
+        settle_deadline,
+    )
+    if download_duration_s is not None:
+        logger.info(f"Download: {download_duration_s:.1f}s (freshly downloaded)")
     else:
-        selected = settle_and_fetch_placements(
-            client, full_model_id, args, settle_timeout=args.settle_timeout
-        )
-
-        if not selected:
-            logger.error("No valid placements matched your filters.")
-            return 1
-
-        selected.sort(
-            key=lambda p: (
-                str(p.get("instance_meta", "")),
-                str(p.get("sharding", "")),
-                nodes_used_in_instance(p["instance"]),
-            ),
-            reverse=True,
-        )
-
-        logger.debug(f"exo-bench model: short_id={short_id} full_id={full_model_id}")
-        logger.info(f"placements: {len(selected)}")
-        for p in selected:
-            logger.info(
-                f"  - {p['sharding']} / {p['instance_meta']} / nodes={nodes_used_in_instance(p['instance'])}"
-            )
-
-        if args.dry_run:
-            return 0
-
-        settle_deadline = (
-            time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
-        )
-
-        logger.info("Planning phase: checking downloads...")
-        download_duration_s = run_planning_phase(
-            client,
-            full_model_id,
-            selected[0],
-            args.danger_delete_downloads,
-            args.timeout,
-            settle_deadline,
-        )
-        if download_duration_s is not None:
-            logger.info(f"Download: {download_duration_s:.1f}s (freshly downloaded)")
-        else:
-            logger.info("Download: model already cached")
+        logger.info("Download: model already cached")
 
     cluster_snapshot = capture_cluster_snapshot(client)
     all_rows: list[dict[str, Any]] = []
     all_system_metrics: dict[str, dict[str, dict[str, float]]] = {}
 
-    # If reusing an existing instance, run a single benchmark pass against it
-    if reused_instance_id is not None:
-        selected = [None]
-
     for preview in selected:
-        created_instance = False
-        if preview is not None:
-            instance = preview["instance"]
-            instance_id = instance_id_from_instance(instance)
+        instance = preview["instance"]
+        instance_id = instance_id_from_instance(instance)
 
-            sharding = str(preview["sharding"])
-            instance_meta = str(preview["instance_meta"])
-            n_nodes = nodes_used_in_instance(instance)
+        sharding = str(preview["sharding"])
+        instance_meta = str(preview["instance_meta"])
+        n_nodes = nodes_used_in_instance(instance)
 
-            logger.info("=" * 80)
-            logger.info(
-                f"PLACEMENT: {sharding} / {instance_meta} / nodes={n_nodes} / instance_id={instance_id}"
-            )
+        logger.info("=" * 80)
+        logger.info(
+            f"PLACEMENT: {sharding} / {instance_meta} / nodes={n_nodes} / instance_id={instance_id}"
+        )
 
-            # Delete any existing instances to free resources before placing
-            try:
-                state = client.request_json("GET", "/state")
-                for old_id in list(state.get("instances", {}).keys()):
-                    logger.info(f"Deleting stale instance {old_id}")
-                    with contextlib.suppress(ExoHttpError):
-                        client.request_json("DELETE", f"/instance/{old_id}")
-                if state.get("instances"):
-                    time.sleep(2)
-            except Exception as e:
-                logger.warning(f"Failed to clean up stale instances: {e}")
+        client.request_json("POST", "/instance", body={"instance": instance})
+        try:
+            wait_for_instance_ready(client, instance_id)
+        except (RuntimeError, TimeoutError) as e:
+            logger.error(f"Failed to initialize placement: {e}")
+            with contextlib.suppress(ExoHttpError):
+                client.request_json("DELETE", f"/instance/{instance_id}")
+            continue
 
-            client.request_json("POST", "/instance", body={"instance": instance})
-            try:
-                wait_for_instance_ready(client, instance_id)
-            except (RuntimeError, TimeoutError) as e:
-                logger.error(f"Failed to initialize placement: {e}")
-                with contextlib.suppress(ExoHttpError):
-                    client.request_json("DELETE", f"/instance/{instance_id}")
-                continue
-
-            time.sleep(1)
-            created_instance = True
-        else:
-            instance_id = reused_instance_id
-            sharding = "reused"
-            instance_meta = "reused"
-            n_nodes = 0
-            logger.info("=" * 80)
-            logger.info(f"Using existing instance {instance_id}")
+        time.sleep(1)
 
         sampler: SystemMetricsSampler | None = None
-        if not args.no_system_metrics and preview is not None:
+        if not args.no_system_metrics:
             nids = node_ids_from_instance(instance)
             sampler = SystemMetricsSampler(
                 ExoClient(args.host, args.port, timeout_s=30),
@@ -666,20 +523,16 @@ def main() -> int:
             )
             sampler.start()
 
-        def _do_one(c: ExoClient, pp: int, tg: int) -> tuple[dict[str, Any], int]:
-            return run_one_completion(
-                c,
-                full_model_id,
-                pp,
-                tg,
-                prompt_sizer,
-                use_prefix_cache=args.use_prefix_cache,
-                stream=args.stream,
-            )
-
         try:
             for i in range(args.warmup):
-                _do_one(client, pp_list[0], tg_list[0])
+                run_one_completion(
+                    client,
+                    full_model_id,
+                    pp_list[0],
+                    tg_list[0],
+                    prompt_sizer,
+                    use_prefix_cache=args.use_prefix_cache,
+                )
                 logger.debug(f"  warmup {i + 1}/{args.warmup} done")
 
             # If pp and tg lists have same length, run in tandem (zip)
@@ -701,7 +554,14 @@ def main() -> int:
                             # Sequential: single request
                             try:
                                 inf_t0 = time.monotonic()
-                                row, actual_pp_tokens = _do_one(client, pp, tg)
+                                row, actual_pp_tokens = run_one_completion(
+                                    client,
+                                    full_model_id,
+                                    pp,
+                                    tg,
+                                    prompt_sizer,
+                                    use_prefix_cache=args.use_prefix_cache,
+                                )
                                 inference_windows.append((inf_t0, time.monotonic()))
                             except Exception as e:
                                 logger.error(e)
@@ -850,12 +710,10 @@ def main() -> int:
                         gen_tps = per_req_tps * concurrency
                         ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
                         gtok = mean(x["stats"]["generation_tokens"] for x in runs)
+                        peak = mean(
+                            x["stats"]["peak_memory_usage"]["inBytes"] for x in runs
+                        )
 
-                        def _peak_bytes(s: dict[str, Any]) -> float:
-                            pm = s["peak_memory_usage"]
-                            return pm.get("inBytes") or pm.get("in_bytes", 0)
-
-                        peak = mean(_peak_bytes(x["stats"]) for x in runs)
                         summary = (
                             f"prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f}    "
                             f"prompt_tokens={ptok} gen_tokens={gtok}    "
@@ -880,16 +738,15 @@ def main() -> int:
                 if placement_metrics:
                     all_system_metrics.update(placement_metrics)
 
-            if created_instance and instance_id is not None:
-                try:
-                    client.request_json("DELETE", f"/instance/{instance_id}")
-                except ExoHttpError as e:
-                    if e.status != 404:
-                        raise
-                wait_for_instance_gone(client, instance_id)
-                logger.debug(f"Deleted instance {instance_id}")
+            try:
+                client.request_json("DELETE", f"/instance/{instance_id}")
+            except ExoHttpError as e:
+                if e.status != 404:
+                    raise
+            wait_for_instance_gone(client, instance_id)
+            logger.debug(f"Deleted instance {instance_id}")
 
-                time.sleep(5)
+            time.sleep(5)
 
     output: dict[str, Any] = {"runs": all_rows}
     if cluster_snapshot:

bench/exo_eval.py

@@ -47,7 +47,6 @@ from harness import (
     ExoHttpError,
     add_common_instance_args,
     capture_cluster_snapshot,
-    find_existing_instance,
     instance_id_from_instance,
     nodes_used_in_instance,
     resolve_model_short_id,
@@ -63,15 +62,6 @@ from loguru import logger
 # ---------------------------------------------------------------------------
 
 MAX_RETRIES = 30
-INSTANCE_HEALTH_CHECK_AFTER = (
-    3  # Check instance health after this many consecutive failures
-)
-
-
-class InstanceFailedError(RuntimeError):
-    """Raised when the exo instance is detected as failed/gone."""
-
-
 DEFAULT_MAX_TOKENS = 16_384
 REASONING_MAX_TOKENS = 131_072
 TEMPERATURE_NON_REASONING = 0.0
@@ -281,7 +271,7 @@ def run_humaneval_test(
 
 @dataclass
 class QuestionResult:
-    question_id: int | str
+    question_id: int
     prompt: str
     response: str
     extracted_answer: str | None
@@ -291,11 +281,7 @@ class QuestionResult:
     prompt_tokens: int = 0
     completion_tokens: int = 0
     reasoning_tokens: int = 0
-    reasoning_content: str = ""
-    finish_reason: str = ""
     elapsed_s: float = 0.0
-    power_watts: float = 0.0
-    energy_joules: float = 0.0
 
 
 @dataclass
@@ -531,10 +517,6 @@ class ApiResult:
     prompt_tokens: int
     completion_tokens: int
     reasoning_tokens: int
-    reasoning_content: str = ""
-    finish_reason: str = ""
-    power_watts: float = 0.0
-    energy_joules: float = 0.0
 
 
 async def _call_api(
@@ -548,9 +530,6 @@ async def _call_api(
     system_message: str | None = None,
     reasoning_effort: str | None = None,
     top_p: float | None = None,
-    top_k: int | None = None,
-    min_p: float | None = None,
-    enable_thinking: bool | None = None,
 ) -> ApiResult:
     messages = []
     if system_message:
@@ -567,12 +546,6 @@ async def _call_api(
         body["reasoning_effort"] = reasoning_effort
     if top_p is not None:
         body["top_p"] = top_p
-    if top_k is not None:
-        body["top_k"] = top_k
-    if min_p is not None:
-        body["min_p"] = min_p
-    if enable_thinking is not None:
-        body["enable_thinking"] = enable_thinking
 
     resp = await client.post(
         f"{base_url}/v1/chat/completions",
@@ -581,40 +554,19 @@ async def _call_api(
     )
     resp.raise_for_status()
     data = resp.json()
-    choice = data["choices"][0]
-    message = choice["message"]
-    content = message.get("content") or ""
-    reasoning_content = message.get("reasoning_content") or ""
-    finish_reason = choice.get("finish_reason") or ""
-
-    # For thinking models, empty content is expected when finish_reason is "length"
-    if not content.strip() and finish_reason != "length" and not reasoning_content:
+    content = data["choices"][0]["message"]["content"]
+    if not content or not content.strip():
         raise ValueError("Empty response from model")
     usage = data.get("usage", {})
     details = usage.get("completion_tokens_details", {})
-    power = data.get("power_usage") or {}
     return ApiResult(
         content=content,
         prompt_tokens=usage.get("prompt_tokens", 0),
         completion_tokens=usage.get("completion_tokens", 0),
         reasoning_tokens=details.get("reasoning_tokens", 0) if details else 0,
-        reasoning_content=reasoning_content,
-        finish_reason=finish_reason,
-        power_watts=power.get("total_avg_sys_power_watts", 0.0),
-        energy_joules=power.get("total_energy_joules", 0.0),
     )
 
 
-async def _check_instance_health(base_url: str) -> bool:
-    """Return True if the exo instance is still reachable."""
-    try:
-        async with httpx.AsyncClient() as c:
-            resp = await c.get(f"{base_url}/models", timeout=5.0)
-            return resp.status_code == 200
-    except Exception:
-        return False
-
-
 async def call_with_retries(
     client: httpx.AsyncClient,
     base_url: str,
@@ -626,14 +578,8 @@ async def call_with_retries(
     system_message: str | None = None,
     reasoning_effort: str | None = None,
     top_p: float | None = None,
-    top_k: int | None = None,
-    min_p: float | None = None,
-    enable_thinking: bool | None = None,
-    instance_failed: asyncio.Event | None = None,
 ) -> ApiResult | None:
     for attempt in range(MAX_RETRIES):
-        if instance_failed and instance_failed.is_set():
-            raise InstanceFailedError("Instance already marked as failed")
         try:
             return await _call_api(
                 client,
@@ -646,30 +592,8 @@ async def call_with_retries(
                 system_message,
                 reasoning_effort,
                 top_p,
-                top_k,
-                min_p,
-                enable_thinking,
             )
         except Exception as e:
-            is_conn_error = isinstance(
-                e,
-                (
-                    httpx.ConnectError,
-                    httpx.RemoteProtocolError,
-                    ConnectionRefusedError,
-                    OSError,
-                ),
-            )
-            if (
-                is_conn_error
-                and attempt >= INSTANCE_HEALTH_CHECK_AFTER
-                and not await _check_instance_health(base_url)
-            ):
-                if instance_failed:
-                    instance_failed.set()
-                raise InstanceFailedError(
-                    f"Instance is down after {attempt + 1} failures: {e}"
-                ) from e
             if attempt < MAX_RETRIES - 1:
                 wait = min(2**attempt, 60)
                 logger.warning(
@@ -694,16 +618,10 @@ async def evaluate_benchmark(
     max_tokens: int,
     concurrency: int = 1,
     limit: int | None = None,
-    offset: int = 0,
     timeout: float | None = None,
     reasoning_effort: str | None = None,
     top_p: float | None = None,
-    top_k: int | None = None,
-    min_p: float | None = None,
-    enable_thinking: bool | None = None,
     difficulty: str | None = None,
-    checkpoint_path: Path | None = None,
-    release_version: str | None = None,
 ) -> list[QuestionResult]:
     """Run a benchmark. Returns per-question results."""
     import datasets
@@ -734,21 +652,7 @@ async def evaluate_benchmark(
         ds = ds.filter(lambda x: x["difficulty"] == difficulty)
         logger.info(f"Filtered to {len(ds)} {difficulty} problems")
 
-    if release_version and "release_version" in ds.column_names:
-        ds = ds.filter(lambda x: x["release_version"] == release_version)
-        logger.info(
-            f"Filtered to {len(ds)} problems with release_version={release_version}"
-        )
-
-    # Sort by question_id to match LCB runner ordering (scenario_router.py:60).
-    # This ensures [offset:offset+limit] slices select the same problems as vllm.
-    if "question_id" in ds.column_names:
-        ds = ds.sort("question_id")
-
     total = len(ds)
-    if offset > 0:
-        ds = ds.select(range(min(offset, total), total))
-        total = len(ds)
     if limit and limit < total:
         ds = ds.select(range(limit))
         total = limit
@@ -756,13 +660,6 @@ async def evaluate_benchmark(
     logger.info(
         f"Evaluating {benchmark_name}: {total} questions, concurrency={concurrency}, "
         f"temperature={temperature}, max_tokens={max_tokens}"
-        + (f", top_k={top_k}" if top_k is not None else "")
-        + (f", min_p={min_p}" if min_p is not None else "")
-        + (
-            f", enable_thinking={enable_thinking}"
-            if enable_thinking is not None
-            else ""
-        )
     )
 
     if config.kind == "code":
@@ -770,64 +667,16 @@ async def evaluate_benchmark(
             "Code benchmarks execute model-generated code. Use a sandboxed environment."
         )
 
-    # Load checkpoint for resume
-    checkpoint_data: dict[str | int, dict[str, Any]] = {}
-    if checkpoint_path and checkpoint_path.exists():
-        with open(checkpoint_path) as f:
-            for line in f:
-                entry = json.loads(line)
-                checkpoint_data[entry["question_id"]] = entry
-        logger.info(f"Loaded {len(checkpoint_data)} checkpointed results")
-
     semaphore = asyncio.Semaphore(concurrency)
-    instance_failed = asyncio.Event()
     results: list[QuestionResult | None] = [None] * total
     completed = 0
     lock = asyncio.Lock()
 
-    def _get_question_id(idx: int, doc: dict) -> str | int:
-        """Get a stable question ID for checkpointing."""
-        if benchmark_name == "livecodebench":
-            return doc.get("question_id", idx)
-        elif benchmark_name == "humaneval":
-            return doc.get("task_id", idx)
-        return idx
-
     async def process_question(
         idx: int, doc: dict, http_client: httpx.AsyncClient
     ) -> None:
         nonlocal completed
         system_msg = None
-        question_id = _get_question_id(idx, doc)
-
-        # Bail out early if instance is already dead
-        if instance_failed.is_set():
-            return
-
-        # Check checkpoint
-        if question_id in checkpoint_data:
-            cached = checkpoint_data[question_id]
-            results[idx] = QuestionResult(
-                question_id=question_id,
-                prompt=cached.get("prompt", ""),
-                response=cached.get("response", ""),
-                extracted_answer=cached.get("extracted_answer"),
-                gold_answer=cached.get("gold_answer", ""),
-                correct=cached.get("correct", False),
-                error=cached.get("error"),
-                prompt_tokens=cached.get("prompt_tokens", 0),
-                completion_tokens=cached.get("completion_tokens", 0),
-                reasoning_tokens=cached.get("reasoning_tokens", 0),
-                reasoning_content=cached.get("reasoning_content", ""),
-                finish_reason=cached.get("finish_reason", ""),
-                elapsed_s=cached.get("elapsed_s", 0.0),
-                power_watts=cached.get("power_watts", 0.0),
-                energy_joules=cached.get("energy_joules", 0.0),
-            )
-            async with lock:
-                completed += 1
-            logger.info(f"  [{completed}/{total}] {question_id} (cached)")
-            return
 
         if benchmark_name == "gpqa_diamond":
             prompt, gold = format_gpqa_question(doc, idx)
@@ -848,50 +697,24 @@ async def evaluate_benchmark(
             raise ValueError(f"Unknown benchmark: {benchmark_name}")
 
         async with semaphore:
-            if instance_failed.is_set():
-                return
             t0 = time.monotonic()
-            try:
-                # Race the API call against the instance_failed event
-                api_task = asyncio.create_task(
-                    call_with_retries(
-                        http_client,
-                        base_url,
-                        model,
-                        prompt,
-                        temperature,
-                        max_tokens,
-                        timeout,
-                        system_message=system_msg,
-                        reasoning_effort=reasoning_effort,
-                        top_p=top_p,
-                        top_k=top_k,
-                        min_p=min_p,
-                        enable_thinking=enable_thinking,
-                        instance_failed=instance_failed,
-                    )
-                )
-                failed_waiter = asyncio.create_task(instance_failed.wait())
-                done, pending = await asyncio.wait(
-                    [api_task, failed_waiter],
-                    return_when=asyncio.FIRST_COMPLETED,
-                )
-                for p in pending:
-                    p.cancel()
-                    with contextlib.suppress(asyncio.CancelledError):
-                        await p
-                if instance_failed.is_set() and api_task not in done:
-                    logger.error(f"Instance failed, aborting {question_id}")
-                    return
-                api_result = api_task.result()
-            except InstanceFailedError:
-                logger.error(f"Instance failed, skipping {question_id}")
-                return
+            api_result = await call_with_retries(
+                http_client,
+                base_url,
+                model,
+                prompt,
+                temperature,
+                max_tokens,
+                timeout,
+                system_message=system_msg,
+                reasoning_effort=reasoning_effort,
+                top_p=top_p,
+            )
             elapsed = time.monotonic() - t0
 
         if api_result is None:
             result = QuestionResult(
-                question_id=question_id,
+                question_id=idx,
                 prompt=prompt,
                 response="",
                 extracted_answer=None,
@@ -906,17 +729,13 @@ async def evaluate_benchmark(
                 "prompt_tokens": api_result.prompt_tokens,
                 "completion_tokens": api_result.completion_tokens,
                 "reasoning_tokens": api_result.reasoning_tokens,
-                "reasoning_content": api_result.reasoning_content,
-                "finish_reason": api_result.finish_reason,
                 "elapsed_s": elapsed,
-                "power_watts": api_result.power_watts,
-                "energy_joules": api_result.energy_joules,
             }
 
             if config.kind == "mc":
                 extracted = extract_mc_answer(response, valid_letters)
                 result = QuestionResult(
-                    question_id=question_id,
+                    question_id=idx,
                     prompt=prompt,
                     response=response,
                     extracted_answer=extracted,
@@ -930,7 +749,7 @@ async def evaluate_benchmark(
                     check_aime_answer(extracted, int(gold)) if extracted else False
                 )
                 result = QuestionResult(
-                    question_id=question_id,
+                    question_id=idx,
                     prompt=prompt,
                     response=response,
                     extracted_answer=extracted,
@@ -944,7 +763,7 @@ async def evaluate_benchmark(
                 code = extract_code_block(response, preserve_indent=keep_indent)
                 if code is None:
                     result = QuestionResult(
-                        question_id=question_id,
+                        question_id=idx,
                         prompt=prompt,
                         response=response,
                         extracted_answer=None,
@@ -959,7 +778,7 @@ async def evaluate_benchmark(
                         code,
                     )
                     result = QuestionResult(
-                        question_id=question_id,
+                        question_id=idx,
                         prompt=prompt,
                         response=response,
                         extracted_answer="pass" if passed else "fail",
@@ -974,7 +793,7 @@ async def evaluate_benchmark(
                         exec_meta["sample"],
                     )
                     result = QuestionResult(
-                        question_id=question_id,
+                        question_id=idx,
                         prompt=prompt,
                         response=response,
                         extracted_answer="pass" if passed else "fail",
@@ -985,7 +804,7 @@ async def evaluate_benchmark(
                     )
                 else:
                     result = QuestionResult(
-                        question_id=question_id,
+                        question_id=idx,
                         prompt=prompt,
                         response=response,
                         extracted_answer=None,
@@ -996,7 +815,7 @@ async def evaluate_benchmark(
                     )
             else:
                 result = QuestionResult(
-                    question_id=question_id,
+                    question_id=idx,
                     prompt=prompt,
                     response=response,
                     extracted_answer=None,
@@ -1008,82 +827,24 @@ async def evaluate_benchmark(
 
         results[idx] = result
 
-        # Write checkpoint (skip infra failures so they get retried on resume,
-        # but keep wrong answers — they are legitimate results)
-        if checkpoint_path is not None and result.response:
-            _write_checkpoint(checkpoint_path, result)
-
         async with lock:
             completed += 1
             n = completed
-
-        # Log progress
-        thinking_info = ""
-        if result.reasoning_content:
-            thinking_info = f", {len(result.reasoning_content)} chars thinking"
-        logger.info(
-            f"  [{n}/{total}] {question_id}: {len(result.response)} chars{thinking_info}, "
-            f"tokens: {result.prompt_tokens}+{result.completion_tokens} "
-            f"[{result.finish_reason}]"
-            + (f" {result.extracted_answer}" if result.extracted_answer else "")
-        )
-
-    async def _health_monitor() -> None:
-        """Periodically check if the instance is still alive."""
-        # Wait a bit before first check to let things start
-        await asyncio.sleep(10)
-        while not instance_failed.is_set():
-            if not await _check_instance_health(base_url):
-                # Double-check to avoid false positives
-                await asyncio.sleep(2)
-                if not await _check_instance_health(base_url):
-                    logger.error("Health monitor: instance is down!")
-                    instance_failed.set()
-                    return
-            await asyncio.sleep(5)
+        if n % max(1, total // 20) == 0 or n == total:
+            correct_so_far = sum(1 for r in results if r is not None and r.correct)
+            answered = sum(1 for r in results if r is not None)
+            logger.info(
+                f"  [{n}/{total}] {correct_so_far}/{answered} correct "
+                f"({correct_so_far / max(answered, 1):.1%})"
+            )
 
     async with httpx.AsyncClient() as http_client:
-        monitor = asyncio.create_task(_health_monitor())
         tasks = [process_question(i, doc, http_client) for i, doc in enumerate(ds)]
         await asyncio.gather(*tasks)
-        monitor.cancel()
-        with contextlib.suppress(asyncio.CancelledError):
-            await monitor
-
-    if instance_failed.is_set():
-        completed_count = sum(1 for r in results if r is not None)
-        logger.error(
-            f"Instance failed! Completed {completed_count}/{total} problems. "
-            f"Checkpoint saved — restart to resume remaining problems."
-        )
-        raise InstanceFailedError("Instance failed during evaluation")
 
     return [r for r in results if r is not None]
 
 
-def _write_checkpoint(path: Path, result: QuestionResult) -> None:
-    """Append a single result to the JSONL checkpoint file."""
-    entry = {
-        "question_id": result.question_id,
-        "prompt": result.prompt,
-        "response": result.response,
-        "extracted_answer": result.extracted_answer,
-        "gold_answer": result.gold_answer,
-        "correct": result.correct,
-        "error": result.error,
-        "prompt_tokens": result.prompt_tokens,
-        "completion_tokens": result.completion_tokens,
-        "reasoning_tokens": result.reasoning_tokens,
-        "reasoning_content": result.reasoning_content,
-        "finish_reason": result.finish_reason,
-        "elapsed_s": round(result.elapsed_s, 2),
-        "power_watts": round(result.power_watts, 2),
-        "energy_joules": round(result.energy_joules, 2),
-    }
-    with open(path, "a") as f:
-        f.write(json.dumps(entry) + "\n")
-
-
 # ---------------------------------------------------------------------------
 # Results display
 # ---------------------------------------------------------------------------
@@ -1106,8 +867,6 @@ def print_results(
     total_elapsed = sum(r.elapsed_s for r in results)
     wall_clock = max(r.elapsed_s for r in results) if results else 0.0
     avg_gen_tps = total_completion_tokens / total_elapsed if total_elapsed > 0 else 0.0
-    total_energy = sum(r.energy_joules for r in results)
-    avg_power = sum(r.power_watts for r in results) / max(total, 1)
 
     label = f"[c={concurrency}] " if concurrency is not None else ""
     print(f"\n{label}{benchmark_name}: {correct}/{total} ({accuracy:.1%})")
@@ -1119,10 +878,6 @@ def print_results(
         f"  |  total time: {total_elapsed:.1f}s  wall clock: {wall_clock:.1f}s"
     )
     print(tok_line)
-    if total_energy > 0:
-        print(
-            f"  power: avg {avg_power:.1f}W  |  total energy: {total_energy:.1f}J ({total_energy / 3600:.2f}Wh)"
-        )
     if errors:
         print(f"  API errors: {errors}")
     if no_extract:
@@ -1141,8 +896,6 @@ def print_results(
         "total_elapsed_s": total_elapsed,
         "wall_clock_s": wall_clock,
         "avg_gen_tps": avg_gen_tps,
-        "avg_power_watts": avg_power,
-        "total_energy_joules": total_energy,
     }
 
 
@@ -1300,11 +1053,7 @@ def save_results(
                 "prompt_tokens": r.prompt_tokens,
                 "completion_tokens": r.completion_tokens,
                 "reasoning_tokens": r.reasoning_tokens,
-                "reasoning_content": r.reasoning_content,
-                "finish_reason": r.finish_reason,
                 "elapsed_s": round(r.elapsed_s, 2),
-                "power_watts": round(r.power_watts, 2),
-                "energy_joules": round(r.energy_joules, 2),
             }
             for r in results
         ],
@@ -1320,15 +1069,6 @@ def save_results(
 # ---------------------------------------------------------------------------
 
 
-def _checkpoint_path(
-    results_dir: str, benchmark: str, model: str, concurrency: int
-) -> Path:
-    """Return the JSONL checkpoint path for a benchmark run."""
-    out_dir = Path(results_dir) / model.replace("/", "_") / benchmark
-    out_dir.mkdir(parents=True, exist_ok=True)
-    return out_dir / f"c{concurrency}.checkpoint.jsonl"
-
-
 def parse_int_list(values: list[str]) -> list[int]:
     items: list[int] = []
     for v in values:
@@ -1356,12 +1096,6 @@ def main() -> int:
         default=None,
         help="Max questions per benchmark (for fast iteration).",
     )
-    ap.add_argument(
-        "--offset",
-        type=int,
-        default=0,
-        help="Skip first N questions (0-based).",
-    )
 
     reasoning_group = ap.add_mutually_exclusive_group()
     reasoning_group.add_argument(
@@ -1381,8 +1115,6 @@ def main() -> int:
         "--temperature", type=float, default=None, help="Override temperature."
     )
     ap.add_argument("--top-p", type=float, default=None, help="Override top_p.")
-    ap.add_argument("--top-k", type=int, default=None, help="Override top_k.")
-    ap.add_argument("--min-p", type=float, default=None, help="Override min_p.")
     ap.add_argument(
         "--max-tokens", type=int, default=None, help="Override max output tokens."
     )
@@ -1416,31 +1148,15 @@ def main() -> int:
         choices=["easy", "medium", "hard"],
         help="Filter by difficulty (livecodebench only). E.g. --difficulty hard",
     )
-    ap.add_argument(
-        "--release-version",
-        default=None,
-        help="LCB dataset release version (livecodebench only). E.g. release_v5",
-    )
     ap.add_argument(
         "--results-dir",
         default="eval_results",
         help="Directory for result JSON files (default: eval_results).",
     )
     ap.add_argument(
-        "--enable-thinking",
-        type=lambda v: v.lower() in ("true", "1", "yes"),
-        default=None,
-        help="Enable thinking mode for models that support it.",
-    )
-    ap.add_argument(
-        "--force",
+        "--skip-instance-setup",
         action="store_true",
-        help="Discard any existing checkpoint and run from scratch.",
-    )
-    ap.add_argument(
-        "--keep-instance",
-        action="store_true",
-        help="Skip deleting the instance after eval (for chaining runs).",
+        help="Skip exo instance management (assumes model is already running).",
     )
 
     args, _ = ap.parse_known_args()
@@ -1461,26 +1177,13 @@ def main() -> int:
     # Instance management
     client = ExoClient(args.host, args.port, timeout_s=args.timeout)
     instance_id: str | None = None
-    created_instance = False
 
-    _short_id, full_model_id = resolve_model_short_id(
-        client,
-        args.model,
-        force_download=args.force_download,
-    )
-
-    # Optionally reuse a running instance for this model
-    if args.reuse_instance:
-        existing = find_existing_instance(client, full_model_id)
-        if existing:
-            instance_id = existing
-            logger.info(f"Reusing existing instance {instance_id}")
-        else:
-            logger.warning(
-                "--reuse-instance: no existing instance found, creating a new one"
-            )
-
-    if instance_id is None:
+    if not args.skip_instance_setup:
+        short_id, full_model_id = resolve_model_short_id(
+            client,
+            args.model,
+            force_download=args.force_download,
+        )
         selected = settle_and_fetch_placements(
             client,
             full_model_id,
@@ -1495,7 +1198,7 @@ def main() -> int:
             key=lambda p: (
                 str(p.get("instance_meta", "")),
                 str(p.get("sharding", "")),
-                nodes_used_in_instance(p["instance"]),
+                -nodes_used_in_instance(p["instance"]),
             ),
             reverse=True,
         )
@@ -1522,18 +1225,6 @@ def main() -> int:
         if download_duration is not None:
             logger.info(f"Download: {download_duration:.1f}s")
 
-        # Delete any existing instances to free resources before placing
-        try:
-            state = client.request_json("GET", "/state")
-            for old_id in list(state.get("instances", {}).keys()):
-                logger.info(f"Deleting stale instance {old_id}")
-                with contextlib.suppress(ExoHttpError):
-                    client.request_json("DELETE", f"/instance/{old_id}")
-            if state.get("instances"):
-                time.sleep(2)
-        except Exception as e:
-            logger.warning(f"Failed to clean up stale instances: {e}")
-
         client.request_json("POST", "/instance", body={"instance": instance})
         try:
             wait_for_instance_ready(client, instance_id)
@@ -1543,9 +1234,10 @@ def main() -> int:
                 client.request_json("DELETE", f"/instance/{instance_id}")
             return 1
         time.sleep(1)
-        created_instance = True
-
-    cluster_snapshot = capture_cluster_snapshot(client)
+        cluster_snapshot = capture_cluster_snapshot(client)
+    else:
+        full_model_id = args.model
+        cluster_snapshot = None
 
     # Auto-detect reasoning from model config
     model_config = load_model_config(full_model_id)
@@ -1599,57 +1291,16 @@ def main() -> int:
         reasoning_effort = str(cfg["reasoning_effort"])
     else:
         reasoning_effort = "high" if is_reasoning else None
-
-    if args.top_k is not None:
-        top_k: int | None = args.top_k
-    elif "top_k" in cfg:
-        top_k = int(cfg["top_k"])
-    else:
-        top_k = None
-
-    if args.min_p is not None:
-        min_p: float | None = args.min_p
-    elif "min_p" in cfg:
-        min_p = float(cfg["min_p"])
-    else:
-        min_p = None
-
-    if args.enable_thinking is not None:
-        enable_thinking: bool | None = args.enable_thinking
-    elif "enable_thinking" in cfg:
-        enable_thinking = bool(cfg["enable_thinking"])
-    else:
-        enable_thinking = None
-
     base_url = f"http://{args.host}:{args.port}"
 
     logger.info(f"Model: {full_model_id}")
     logger.info(
         f"Settings: temperature={temperature}, max_tokens={max_tokens}, "
         + (f"top_p={top_p}, " if top_p is not None else "")
-        + (f"top_k={top_k}, " if top_k is not None else "")
-        + (f"min_p={min_p}, " if min_p is not None else "")
         + f"reasoning={'yes' if is_reasoning else 'no'}"
         + (f", reasoning_effort={reasoning_effort}" if reasoning_effort else "")
-        + (
-            f", enable_thinking={enable_thinking}"
-            if enable_thinking is not None
-            else ""
-        )
     )
 
-    # Common kwargs for evaluate_benchmark
-    eval_kwargs: dict[str, Any] = {
-        "reasoning_effort": reasoning_effort,
-        "top_p": top_p,
-        "top_k": top_k,
-        "min_p": min_p,
-        "enable_thinking": enable_thinking,
-        "difficulty": args.difficulty,
-        "offset": args.offset,
-        "release_version": args.release_version,
-    }
-
     try:
         if args.compare_concurrency:
             concurrency_levels = parse_int_list(args.compare_concurrency)
@@ -1658,11 +1309,6 @@ def main() -> int:
                 for c in concurrency_levels:
                     logger.info(f"\n{'=' * 50}")
                     logger.info(f"Running {task_name} at concurrency={c}")
-                    checkpoint_path = _checkpoint_path(
-                        args.results_dir, task_name, full_model_id, c
-                    )
-                    if args.force and checkpoint_path.exists():
-                        checkpoint_path.unlink()
                     results = asyncio.run(
                         evaluate_benchmark(
                             task_name,
@@ -1673,8 +1319,9 @@ def main() -> int:
                             concurrency=c,
                             limit=args.limit,
                             timeout=args.request_timeout,
-                            checkpoint_path=checkpoint_path,
-                            **eval_kwargs,
+                            reasoning_effort=reasoning_effort,
+                            top_p=top_p,
+                            difficulty=args.difficulty,
                         )
                     )
                     if results:
@@ -1689,18 +1336,10 @@ def main() -> int:
                             cluster=cluster_snapshot,
                         )
                         results_by_c[c] = results
-                    # Clean up checkpoint on success
-                    if checkpoint_path.exists():
-                        checkpoint_path.unlink()
                 if len(results_by_c) >= 2:
                     print_comparison(task_name, results_by_c)
         else:
             for task_name in task_names:
-                checkpoint_path = _checkpoint_path(
-                    args.results_dir, task_name, full_model_id, args.num_concurrent
-                )
-                if args.force and checkpoint_path.exists():
-                    checkpoint_path.unlink()
                 results = asyncio.run(
                     evaluate_benchmark(
                         task_name,
@@ -1711,8 +1350,9 @@ def main() -> int:
                         concurrency=args.num_concurrent,
                         limit=args.limit,
                         timeout=args.request_timeout,
-                        checkpoint_path=checkpoint_path,
-                        **eval_kwargs,
+                        reasoning_effort=reasoning_effort,
+                        top_p=top_p,
+                        difficulty=args.difficulty,
                     )
                 )
                 if results:
@@ -1726,25 +1366,14 @@ def main() -> int:
                         scores,
                         cluster=cluster_snapshot,
                     )
-                # Clean up checkpoint on success
-                if checkpoint_path.exists():
-                    checkpoint_path.unlink()
     finally:
-        if created_instance and instance_id is not None:
-            if args.keep_instance:
-                logger.info(f"Keeping instance {instance_id} (--keep-instance)")
-            else:
-                try:
-                    client.request_json("DELETE", f"/instance/{instance_id}")
-                except ExoHttpError as e:
-                    if e.status != 404:
-                        raise
-                try:
-                    wait_for_instance_gone(client, instance_id)
-                except TimeoutError:
-                    logger.warning(
-                        f"Timed out waiting for instance {instance_id} to be deleted"
-                    )
+        if instance_id is not None:
+            try:
+                client.request_json("DELETE", f"/instance/{instance_id}")
+            except ExoHttpError as e:
+                if e.status != 404:
+                    raise
+            wait_for_instance_gone(client, instance_id)
 
     return 0
 

bench/harness.py

@@ -6,7 +6,6 @@ import http.client
 import json
 import os
 import time
-from collections.abc import Iterator
 from typing import Any
 from urllib.parse import urlencode
 
@@ -70,30 +69,6 @@ class ExoClient:
     def post_bench_chat_completions(self, payload: dict[str, Any]) -> dict[str, Any]:
         return self.request_json("POST", "/bench/chat/completions", body=payload)
 
-    def stream_bench_chat_completions(self, payload: dict[str, Any]) -> Iterator[str]:
-        """POST /bench/chat/completions with stream=True, yielding raw SSE lines."""
-        payload = {**payload, "stream": True}
-        data = json.dumps(payload).encode("utf-8")
-        conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s)
-        try:
-            conn.request(
-                "POST",
-                "/bench/chat/completions",
-                body=data,
-                headers={
-                    "Content-Type": "application/json",
-                    "Accept": "text/event-stream",
-                },
-            )
-            resp = conn.getresponse()
-            if resp.status >= 400:
-                raw = resp.read().decode("utf-8", errors="replace")
-                raise ExoHttpError(resp.status, resp.reason, raw[:300])
-            for line in resp:
-                yield line.decode("utf-8", errors="replace")
-        finally:
-            conn.close()
-
     def get_state_path(self, path: str) -> Any:
         try:
             return self.request_json("GET", f"/state/{path}")
@@ -293,15 +268,11 @@ def sharding_filter(sharding: str, wanted: str) -> bool:
 
 
 def fetch_and_filter_placements(
-    client: ExoClient,
-    full_model_id: str,
-    args: argparse.Namespace,
-    node_id: str | None = None,
+    client: ExoClient, full_model_id: str, args: argparse.Namespace
 ) -> list[dict[str, Any]]:
-    params: dict[str, str] = {"model_id": full_model_id}
-    if node_id is not None:
-        params["node_ids"] = node_id
-    previews_resp = client.request_json("GET", "/instance/previews", params=params)
+    previews_resp = client.request_json(
+        "GET", "/instance/previews", params={"model_id": full_model_id}
+    )
     previews = previews_resp.get("previews") or []
 
     selected: list[dict[str, Any]] = []
@@ -361,9 +332,8 @@ def settle_and_fetch_placements(
     full_model_id: str,
     args: argparse.Namespace,
     settle_timeout: float = 0,
-    node_id: str | None = None,
 ) -> list[dict[str, Any]]:
-    selected = fetch_and_filter_placements(client, full_model_id, args, node_id=node_id)
+    selected = fetch_and_filter_placements(client, full_model_id, args)
 
     if not selected and settle_timeout > 0:
         backoff = _SETTLE_INITIAL_BACKOFF_S
@@ -376,9 +346,7 @@ def settle_and_fetch_placements(
             )
             time.sleep(min(backoff, remaining))
             backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S)
-            selected = fetch_and_filter_placements(
-                client, full_model_id, args, node_id=node_id
-            )
+            selected = fetch_and_filter_placements(client, full_model_id, args)
 
     return selected
 
@@ -494,8 +462,9 @@ def run_planning_phase(
         )
         logger.info(f"Started download on {node_id}")
 
-    # Wait for downloads (no timeout — poll until complete or failed)
-    while True:
+    # Wait for downloads
+    start = time.time()
+    while time.time() - start < timeout:
         all_done = True
         for node_id in node_ids:
             node_downloads = client.get_node_downloads(node_id) or []
@@ -545,24 +514,9 @@ def run_planning_phase(
             if download_t0 is not None:
                 return time.perf_counter() - download_t0
             return None
-        time.sleep(10)
+        time.sleep(1)
 
-
-def find_existing_instance(client: ExoClient, model_id: str) -> str | None:
-    """Find an existing running instance for the given model."""
-    try:
-        state = client.request_json("GET", "/state")
-    except Exception:
-        return None
-    for inst_id, inst in state.get("instances", {}).items():
-        # Instance structure is nested: {"MlxJacclInstance": {"shardAssignments": {"modelId": ...}}}
-        for _inst_type, inner in inst.items():
-            if not isinstance(inner, dict):
-                continue
-            sa = inner.get("shardAssignments", {})
-            if sa.get("modelId") == model_id:
-                return inst_id
-    return None
+    raise TimeoutError("Downloads did not complete in time")
 
 
 def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
@@ -618,8 +572,3 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
         action="store_true",
         help="Delete existing models from smallest to largest to make room for benchmark model.",
     )
-    ap.add_argument(
-        "--reuse-instance",
-        action="store_true",
-        help="Reuse an existing running instance for this model instead of creating a new one.",
-    )

bench/prefill-decode.toml

@@ -1,36 +0,0 @@
-# Prefill/Decode disaggregation benchmark config.
-#
-# Top-level keys are bench-wide. [prefill] and [decode] sections set per-side
-# placement filters and (optionally) per-side model.
-#
-# Example:
-#   uv run python bench/prefill_decode_bench.py --config bench/prefill-decode.toml
-
-host = "james"
-port = 52415
-timeout = 7200.0
-settle_timeout = 60.0
-
-# Workload
-pp = [4096]
-tg = [512]
-repeat = 1
-warmup = 0
-
-json_out = "bench/prefill_decode_results.json"
-
-[prefill]
-model = "mlx-community/gpt-oss-20b-MXFP4-Q8"
-node = "mike"
-instance_meta = "ring"
-sharding = "pipeline"
-min_nodes = 1
-max_nodes = 1
-
-[decode]
-model = "mlx-community/gpt-oss-20b-MXFP4-Q8"
-node = "james"
-instance_meta = "ring"
-sharding = "pipeline"
-min_nodes = 1
-max_nodes = 1

bench/prefill_decode_bench.py

@@ -1,785 +0,0 @@
-# type: ignore
-#!/usr/bin/env python3
-"""Disaggregated prefill-decode benchmark for exo (MLX → MLX).
-
-Spins up two MLX instances on the cluster, marks one as Prefill source and
-the other as Decode target via /v1/instance-links, then sends chat
-completions to the API. The master routes the request to the decode
-instance and stamps `prefill_endpoint` pointing at the prefill instance —
-the worker decides per-request whether to ship prefill remotely
-(uncached_count > REMOTE_PREFILL_MIN_TOKENS).
-
-Usage:
-    uv run python bench/prefill_decode_bench.py --model <id> --pp 2048,8192 --tg 128
-    uv run python bench/prefill_decode_bench.py --model <id> --pp 4096 --tg 128 --repeat 3
-    uv run python bench/prefill_decode_bench.py --model <id> --pp 2048 --tg 128 --dry-run
-"""
-
-from __future__ import annotations
-
-import argparse
-import contextlib
-import copy
-import itertools
-import json
-import sys
-import time
-import tomllib
-from pathlib import Path
-from statistics import mean
-from typing import Any
-
-from exo_bench import (
-    PromptSizer,
-    format_peak_memory,
-    load_tokenizer_for_bench,
-    parse_int_list,
-)
-from harness import (
-    ExoClient,
-    ExoHttpError,
-    add_common_instance_args,
-    instance_id_from_instance,
-    node_ids_from_instance,
-    nodes_used_in_instance,
-    resolve_model_short_id,
-    run_planning_phase,
-    settle_and_fetch_placements,
-    unwrap_instance,
-    wait_for_instance_gone,
-    wait_for_instance_ready,
-)
-from loguru import logger
-
-
-def _node_id_to_friendly(client: ExoClient) -> dict[str, str]:
-    identities = client.get_node_identities() or {}
-    out: dict[str, str] = {}
-    for node_id, identity in identities.items():
-        if isinstance(identity, dict):
-            name = identity.get("friendlyName") or identity.get("friendly_name")
-            if isinstance(name, str):
-                out[str(node_id)] = name
-    return out
-
-
-def _placement_node_friendly_names(
-    placement: dict[str, Any], id_to_friendly: dict[str, str]
-) -> list[str]:
-    instance = placement["instance"]
-    return [id_to_friendly.get(nid, nid) for nid in node_ids_from_instance(instance)]
-
-
-def _filter_by_node(
-    placements: list[dict[str, Any]],
-    friendly_name: str,
-    id_to_friendly: dict[str, str],
-) -> list[dict[str, Any]]:
-    target = friendly_name.lower()
-    matched: list[dict[str, Any]] = []
-    for p in placements:
-        names = [n.lower() for n in _placement_node_friendly_names(p, id_to_friendly)]
-        if any(target == n or target in n for n in names):
-            matched.append(p)
-    return matched
-
-
-def _node_id_by_friendly(id_to_friendly: dict[str, str], target: str) -> str | None:
-    target_lc = target.lower()
-    for nid, name in id_to_friendly.items():
-        if target_lc == name.lower() or target_lc in name.lower():
-            return nid
-    return None
-
-
-def _load_toml(path: str) -> dict[str, Any]:
-    with Path(path).open("rb") as f:
-        return tomllib.load(f)
-
-
-_TOP_LEVEL_TOML_KEYS = {
-    "host",
-    "port",
-    "timeout",
-    "settle_timeout",
-    "model",
-    "pp",
-    "tg",
-    "repeat",
-    "warmup",
-    "json_out",
-    "instance_meta",
-    "sharding",
-    "min_nodes",
-    "max_nodes",
-    "force_download",
-    "danger_delete_downloads",
-    "all_combinations",
-}
-
-
-def _inject_toml_into_argv() -> None:
-    """If --config X is in sys.argv, pre-load it and inject required CLI args
-    (--model, --pp, --tg) so argparse's required=True checks pass."""
-    argv = sys.argv
-    if "--config" not in argv:
-        return
-    idx = argv.index("--config")
-    if idx + 1 >= len(argv):
-        return
-    cfg_path = argv[idx + 1]
-    cfg = _load_toml(cfg_path)
-    decode = cfg.get("decode", {})
-
-    def _has(flag: str) -> bool:
-        return any(a == flag or a.startswith(flag + "=") for a in argv)
-
-    # --model: prefer top-level, then [decode].model
-    if not _has("--model"):
-        model = cfg.get("model") or decode.get("model")
-        if model:
-            argv += ["--model", str(model)]
-    if not _has("--pp"):
-        pp = cfg.get("pp")
-        if pp:
-            argv += (
-                ["--pp", *(str(x) for x in pp)]
-                if isinstance(pp, list)
-                else [
-                    "--pp",
-                    str(pp),
-                ]
-            )
-    if not _has("--tg"):
-        tg = cfg.get("tg")
-        if tg:
-            argv += (
-                ["--tg", *(str(x) for x in tg)]
-                if isinstance(tg, list)
-                else [
-                    "--tg",
-                    str(tg),
-                ]
-            )
-
-
-def _merge_toml_into_args(args: argparse.Namespace, cfg: dict[str, Any]) -> None:
-    """Apply top-level toml keys onto args namespace where args has a default."""
-    for key, value in cfg.items():
-        if key in {"prefill", "decode"}:
-            continue
-        if key not in _TOP_LEVEL_TOML_KEYS:
-            continue
-        attr = key
-        current = getattr(args, attr, None)
-        if current in (None, [], False):
-            setattr(args, attr, value)
-
-
-def _side_args(
-    base: argparse.Namespace, overrides: dict[str, Any]
-) -> argparse.Namespace:
-    out = copy.copy(base)
-    for k in (
-        "instance_meta",
-        "sharding",
-        "min_nodes",
-        "max_nodes",
-        "skip_pipeline_jaccl",
-        "skip_tensor_ring",
-    ):
-        if k in overrides:
-            setattr(out, k, overrides[k])
-    return out
-
-
-def _pick_two_distinct_placements(
-    placements: list[dict[str, Any]],
-) -> tuple[dict[str, Any], dict[str, Any]] | None:
-    if len(placements) < 2:
-        return None
-    seen_nodes: set[tuple[str, ...]] = set()
-    chosen: list[dict[str, Any]] = []
-    for p in placements:
-        nodes = tuple(sorted(str(n) for n in p.get("nodes", [])))
-        if nodes in seen_nodes:
-            continue
-        seen_nodes.add(nodes)
-        chosen.append(p)
-        if len(chosen) == 2:
-            return chosen[0], chosen[1]
-    return None
-
-
-def _create_instance_link(
-    client: ExoClient,
-    prefill_instance_id: str,
-    decode_instance_id: str,
-) -> str:
-    out = client.request_json(
-        "POST",
-        "/v1/instance-links",
-        body={
-            "prefill_instances": [prefill_instance_id],
-            "decode_instances": [decode_instance_id],
-        },
-    )
-    return str(out.get("commandId", ""))
-
-
-def _list_instance_links(client: ExoClient) -> list[dict[str, Any]]:
-    out = client.request_json("GET", "/v1/instance-links")
-    return out if isinstance(out, list) else []
-
-
-def _delete_instance_link(client: ExoClient, link_id: str) -> None:
-    client.request_json("DELETE", f"/v1/instance-links/{link_id}")
-
-
-def run_one(
-    client: ExoClient,
-    model_id: str,
-    pp_hint: int,
-    tg: int,
-    prompt_sizer: PromptSizer,
-) -> tuple[dict[str, Any], int]:
-    content, pp_tokens = prompt_sizer.build(pp_hint)
-    payload: dict[str, Any] = {
-        "model": model_id,
-        "messages": [{"role": "user", "content": content}],
-        "stream": False,
-        "max_tokens": tg,
-    }
-
-    t0 = time.perf_counter()
-    out = client.post_bench_chat_completions(payload)
-    elapsed = time.perf_counter() - t0
-
-    stats = out.get("generation_stats")
-    choices = out.get("choices") or [{}]
-    message = choices[0].get("message", {}) if choices else {}
-    text = message.get("content") or ""
-    preview = text[:200] if text else ""
-
-    return {
-        "elapsed_s": elapsed,
-        "output_text_preview": preview,
-        "stats": stats,
-    }, pp_tokens
-
-
-def _run_phase(
-    *,
-    client: ExoClient,
-    label: str,
-    pp_tg_pairs: list[tuple[int, int]],
-    model_id: str,
-    prompt_sizer: PromptSizer,
-    warmup: int,
-    repeat: int,
-    common_meta: dict[str, Any],
-) -> list[dict[str, Any]]:
-    logger.info(f"=== phase: {label} (model={model_id}) ===")
-    rows: list[dict[str, Any]] = []
-    for i in range(warmup):
-        run_one(client, model_id, pp_tg_pairs[0][0], pp_tg_pairs[0][1], prompt_sizer)
-        logger.debug(f"  warmup {i + 1}/{warmup} done")
-
-    for pp, tg in pp_tg_pairs:
-        logger.info(f"--- {label}: pp={pp} tg={tg} ---")
-        runs: list[dict[str, Any]] = []
-        for r in range(repeat):
-            time.sleep(2)
-            try:
-                row, actual_pp_tokens = run_one(client, model_id, pp, tg, prompt_sizer)
-            except Exception as e:
-                logger.error(e)
-                continue
-            row.update(common_meta)
-            row.update(
-                {
-                    "phase": label,
-                    "phase_model_id": model_id,
-                    "pp_tokens": actual_pp_tokens,
-                    "tg": tg,
-                    "repeat_index": r,
-                }
-            )
-            runs.append(row)
-            rows.append(row)
-
-        if runs:
-            prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
-            gen_tps = mean(x["stats"]["generation_tps"] 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(x["stats"]["peak_memory_usage"]["inBytes"] for x in runs)
-            avg_elapsed = mean(x["elapsed_s"] for x in runs)
-            logger.info(
-                f"[{label}] prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f}    "
-                f"prompt_tokens={ptok} gen_tokens={gtok}    "
-                f"peak_memory={format_peak_memory(peak)}    "
-                f"avg_elapsed={avg_elapsed:.2f}s"
-            )
-        time.sleep(2)
-    return rows
-
-
-def _summarise(rows: list[dict[str, Any]]) -> dict[tuple[int, int], dict[str, float]]:
-    grouped: dict[tuple[int, int], list[dict[str, Any]]] = {}
-    for r in rows:
-        key = (int(r["pp_tokens"]), int(r["tg"]))
-        grouped.setdefault(key, []).append(r)
-    out: dict[tuple[int, int], dict[str, float]] = {}
-    for key, runs in grouped.items():
-        out[key] = {
-            "prompt_tps": mean(x["stats"]["prompt_tps"] for x in runs),
-            "gen_tps": mean(x["stats"]["generation_tps"] for x in runs),
-            "elapsed_s": mean(x["elapsed_s"] for x in runs),
-        }
-    return out
-
-
-def _print_diff(
-    disagg_rows: list[dict[str, Any]],
-    decode_alone_rows: list[dict[str, Any]],
-    prefill_alone_rows: list[dict[str, Any]],
-) -> None:
-    disagg = _summarise(disagg_rows)
-    decode_alone = _summarise(decode_alone_rows)
-    prefill_alone = _summarise(prefill_alone_rows)
-    keys = set(disagg.keys()) | set(decode_alone.keys()) | set(prefill_alone.keys())
-
-    width = 64
-    for key in sorted(keys):
-        pp, tg = key
-        logger.info("─" * width)
-        logger.info(f"  pp={pp}  tg={tg}")
-        logger.info("─" * width)
-        logger.info(
-            f"  {'phase':<16} {'elapsed':>10}  {'prompt_tps':>11}  {'gen_tps':>9}"
-        )
-        for label, summary in (
-            ("disaggregated", disagg.get(key)),
-            ("decode_alone", decode_alone.get(key)),
-            ("prefill_alone", prefill_alone.get(key)),
-        ):
-            if summary is None:
-                logger.info(f"  {label:<16} {'—':>10}  {'—':>11}  {'—':>9}")
-                continue
-            logger.info(
-                f"  {label:<16} "
-                f"{summary['elapsed_s']:>9.2f}s  "
-                f"{summary['prompt_tps']:>11.1f}  "
-                f"{summary['gen_tps']:>9.2f}"
-            )
-
-        d = disagg.get(key)
-        da = decode_alone.get(key)
-        pa = prefill_alone.get(key)
-        if d and da and d["elapsed_s"] > 0:
-            logger.info(
-                f"  speedup vs decode_alone:  {da['elapsed_s'] / d['elapsed_s']:.2f}x"
-            )
-        if d and pa and d["elapsed_s"] > 0:
-            logger.info(
-                f"  speedup vs prefill_alone: {pa['elapsed_s'] / d['elapsed_s']:.2f}x"
-            )
-    logger.info("─" * width)
-
-
-def main() -> int:
-    _inject_toml_into_argv()
-    ap = argparse.ArgumentParser(
-        prog="prefill-decode-bench",
-        description="Benchmark MLX-MLX disaggregated prefill/decode via instance links.",
-    )
-    add_common_instance_args(ap)
-    ap.add_argument(
-        "--pp",
-        nargs="+",
-        required=True,
-        help="Prompt-size hints (ints, must be >1000). Accepts commas.",
-    )
-    ap.add_argument(
-        "--tg",
-        nargs="+",
-        required=True,
-        help="Generation lengths (ints). Accepts commas.",
-    )
-    ap.add_argument(
-        "--repeat", type=int, default=1, help="Repetitions per (pp,tg) pair."
-    )
-    ap.add_argument(
-        "--warmup",
-        type=int,
-        default=0,
-        help="Warmup runs (uses first pp/tg).",
-    )
-    ap.add_argument(
-        "--json-out",
-        default="bench/prefill_decode_results.json",
-        help="Write raw per-run results JSON to this path.",
-    )
-    ap.add_argument("--stdout", action="store_true", help="Write results to stdout")
-    ap.add_argument(
-        "--dry-run", action="store_true", help="List selected placements and exit."
-    )
-    ap.add_argument(
-        "--all-combinations",
-        action="store_true",
-        help="Force all pp×tg combinations even when lists have equal length.",
-    )
-    ap.add_argument(
-        "--prefill-model",
-        default=None,
-        help="Model id for the prefill instance. Defaults to --model.",
-    )
-    ap.add_argument(
-        "--prefill-node",
-        default=None,
-        help="friendly_name of the node hosting the prefill instance.",
-    )
-    ap.add_argument(
-        "--decode-node",
-        default=None,
-        help="friendly_name of the node hosting the decode instance.",
-    )
-    ap.add_argument(
-        "--config",
-        default=None,
-        help="TOML config file. CLI flags override toml values.",
-    )
-    ap.add_argument(
-        "--compare-baseline",
-        action="store_true",
-        help="Also run each (pp,tg) pair without the prefill/decode link "
-        "(decode instance does its own prefill) and report the diff.",
-    )
-    args = ap.parse_args()
-    cfg = _load_toml(args.config) if args.config else {}
-    _merge_toml_into_args(args, cfg)
-    prefill_overrides = cfg.get("prefill", {}) if cfg else {}
-    decode_overrides = cfg.get("decode", {}) if cfg else {}
-    if args.prefill_model is None and "model" in prefill_overrides:
-        args.prefill_model = prefill_overrides["model"]
-    if args.prefill_node is None and "node" in prefill_overrides:
-        args.prefill_node = prefill_overrides["node"]
-    if args.decode_node is None and "node" in decode_overrides:
-        args.decode_node = decode_overrides["node"]
-    if "model" in decode_overrides and not args.model:
-        args.model = decode_overrides["model"]
-
-    pp_list = parse_int_list(args.pp)
-    tg_list = parse_int_list(args.tg)
-    if not pp_list or not tg_list:
-        logger.error("pp and tg lists must be non-empty")
-        return 2
-    for pp in pp_list:
-        if pp <= 1000:
-            logger.error(
-                f"pp={pp} must be >1000 (remote prefill triggers when uncached >1000)"
-            )
-            return 2
-    if args.repeat <= 0:
-        logger.error("--repeat must be >= 1")
-        return 2
-
-    use_combinations = args.all_combinations or len(pp_list) != len(tg_list)
-    if use_combinations:
-        logger.info(
-            f"pp/tg mode: combinations (product) — {len(pp_list) * len(tg_list)} pairs"
-        )
-    else:
-        logger.info(f"pp/tg mode: tandem (zip) — {len(pp_list)} pairs")
-
-    client = ExoClient(args.host, args.port, timeout_s=args.timeout)
-
-    decode_short_id, decode_full_id = resolve_model_short_id(
-        client, args.model, force_download=args.force_download
-    )
-    if args.prefill_model:
-        prefill_short_id, prefill_full_id = resolve_model_short_id(
-            client, args.prefill_model, force_download=args.force_download
-        )
-    else:
-        prefill_short_id, prefill_full_id = decode_short_id, decode_full_id
-
-    tokenizer = load_tokenizer_for_bench(decode_full_id)
-    if tokenizer is None:
-        raise RuntimeError("[prefill-decode-bench] decode tokenizer load failed")
-    try:
-        decode_prompt_sizer = PromptSizer(tokenizer)
-    except Exception:
-        logger.error("[prefill-decode-bench] decode prompt sizing failed")
-        raise
-
-    if prefill_full_id == decode_full_id:
-        prefill_prompt_sizer = decode_prompt_sizer
-    else:
-        prefill_tokenizer = load_tokenizer_for_bench(prefill_full_id)
-        if prefill_tokenizer is None:
-            raise RuntimeError("[prefill-decode-bench] prefill tokenizer load failed")
-        prefill_prompt_sizer = PromptSizer(prefill_tokenizer)
-
-    id_to_friendly = _node_id_to_friendly(client)
-
-    prefill_args = _side_args(args, prefill_overrides)
-    decode_args = _side_args(args, decode_overrides)
-
-    if prefill_full_id == decode_full_id and prefill_overrides == decode_overrides:
-        placements = settle_and_fetch_placements(
-            client, decode_full_id, args, settle_timeout=args.settle_timeout
-        )
-        prefill_candidates = (
-            _filter_by_node(placements, args.prefill_node, id_to_friendly)
-            if args.prefill_node
-            else placements
-        )
-        decode_candidates = (
-            _filter_by_node(placements, args.decode_node, id_to_friendly)
-            if args.decode_node
-            else placements
-        )
-        if args.prefill_node and not prefill_candidates:
-            logger.error(f"No placement on prefill node {args.prefill_node!r}.")
-            return 1
-        if args.decode_node and not decode_candidates:
-            logger.error(f"No placement on decode node {args.decode_node!r}.")
-            return 1
-        if args.prefill_node and args.decode_node:
-            prefill_p = prefill_candidates[0]
-            decode_p = decode_candidates[0]
-        else:
-            pair = _pick_two_distinct_placements(placements)
-            if pair is None:
-                logger.error(
-                    "Need at least two distinct-node MLX placements for the same model."
-                )
-                return 1
-            prefill_p, decode_p = pair
-            if args.prefill_node:
-                prefill_p = prefill_candidates[0]
-            if args.decode_node:
-                decode_p = decode_candidates[0]
-    else:
-        prefill_node_id = (
-            _node_id_by_friendly(id_to_friendly, args.prefill_node)
-            if args.prefill_node
-            else None
-        )
-        decode_node_id = (
-            _node_id_by_friendly(id_to_friendly, args.decode_node)
-            if args.decode_node
-            else None
-        )
-        if args.prefill_node and prefill_node_id is None:
-            logger.error(f"Unknown node {args.prefill_node!r}.")
-            return 1
-        if args.decode_node and decode_node_id is None:
-            logger.error(f"Unknown node {args.decode_node!r}.")
-            return 1
-        prefill_placements = settle_and_fetch_placements(
-            client,
-            prefill_full_id,
-            prefill_args,
-            settle_timeout=args.settle_timeout,
-            node_id=prefill_node_id,
-        )
-        decode_placements = settle_and_fetch_placements(
-            client,
-            decode_full_id,
-            decode_args,
-            settle_timeout=args.settle_timeout,
-            node_id=decode_node_id,
-        )
-        if not prefill_placements:
-            logger.error(
-                f"No placement found for prefill model {prefill_full_id}"
-                f"{f' on node {args.prefill_node!r}' if args.prefill_node else ''}."
-            )
-            return 1
-        if not decode_placements:
-            logger.error(
-                f"No placement found for decode model {decode_full_id}"
-                f"{f' on node {args.decode_node!r}' if args.decode_node else ''}."
-            )
-            return 1
-        prefill_p = prefill_placements[0]
-        decode_p = decode_placements[0]
-
-    prefill_node_names = _placement_node_friendly_names(prefill_p, id_to_friendly)
-    decode_node_names = _placement_node_friendly_names(decode_p, id_to_friendly)
-    _ = unwrap_instance
-
-    prefill_instance = prefill_p["instance"]
-    decode_instance = decode_p["instance"]
-    prefill_id = instance_id_from_instance(prefill_instance)
-    decode_id = instance_id_from_instance(decode_instance)
-    prefill_meta = str(prefill_p.get("instance_meta", ""))
-    decode_meta = str(decode_p.get("instance_meta", ""))
-    prefill_nodes = nodes_used_in_instance(prefill_instance)
-    decode_nodes = nodes_used_in_instance(decode_instance)
-
-    logger.info("=" * 80)
-    logger.info(
-        f"PREFILL: {prefill_meta} / nodes={prefill_nodes} ({','.join(prefill_node_names)}) "
-        f"/ {prefill_short_id} ({prefill_full_id}) / instance_id={prefill_id}"
-    )
-    logger.info(
-        f"DECODE:  {decode_meta} / nodes={decode_nodes} ({','.join(decode_node_names)}) "
-        f"/ {decode_short_id} ({decode_full_id}) / instance_id={decode_id}"
-    )
-
-    if args.dry_run:
-        return 0
-
-    settle_deadline = (
-        time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
-    )
-
-    logger.info("Planning phase: prefill...")
-    run_planning_phase(
-        client,
-        prefill_full_id,
-        prefill_p,
-        args.danger_delete_downloads,
-        args.timeout,
-        settle_deadline,
-    )
-    logger.info("Planning phase: decode...")
-    run_planning_phase(
-        client,
-        decode_full_id,
-        decode_p,
-        args.danger_delete_downloads,
-        args.timeout,
-        settle_deadline,
-    )
-
-    if use_combinations:
-        pp_tg_pairs = list(itertools.product(pp_list, tg_list))
-    else:
-        pp_tg_pairs = list(zip(pp_list, tg_list, strict=True))
-
-    common_meta = {
-        "decode_model_short_id": decode_short_id,
-        "decode_model_id": decode_full_id,
-        "prefill_model_short_id": prefill_short_id,
-        "prefill_model_id": prefill_full_id,
-        "prefill_instance_id": prefill_id,
-        "prefill_instance_meta": prefill_meta,
-        "prefill_nodes": prefill_nodes,
-        "decode_instance_id": decode_id,
-        "decode_instance_meta": decode_meta,
-        "decode_nodes": decode_nodes,
-    }
-
-    all_rows: list[dict[str, Any]] = []
-    disagg_rows: list[dict[str, Any]] = []
-    decode_alone_rows: list[dict[str, Any]] = []
-    prefill_alone_rows: list[dict[str, Any]] = []
-    link_id = ""
-    prefill_alive = False
-    decode_alive = False
-    try:
-        logger.info("Creating prefill instance...")
-        client.request_json("POST", "/instance", body={"instance": prefill_instance})
-        wait_for_instance_ready(client, prefill_id)
-        prefill_alive = True
-        logger.info("Prefill instance ready")
-
-        if args.compare_baseline:
-            time.sleep(2)
-            prefill_alone_rows = _run_phase(
-                client=client,
-                label="prefill_alone",
-                pp_tg_pairs=pp_tg_pairs,
-                model_id=prefill_full_id,
-                prompt_sizer=prefill_prompt_sizer,
-                warmup=args.warmup,
-                repeat=args.repeat,
-                common_meta=common_meta,
-            )
-            all_rows.extend(prefill_alone_rows)
-
-        logger.info("Creating decode instance...")
-        client.request_json("POST", "/instance", body={"instance": decode_instance})
-        wait_for_instance_ready(client, decode_id)
-        decode_alive = True
-        logger.info("Decode instance ready")
-
-        logger.info("Linking instances (prefill → decode)...")
-        _create_instance_link(client, prefill_id, decode_id)
-        time.sleep(1)
-        links = _list_instance_links(client)
-        if not links:
-            logger.error("Link did not appear in state.")
-            return 1
-        link_id = str(links[-1].get("linkId") or links[-1].get("link_id") or "")
-        logger.info(f"Link created: {link_id}")
-        time.sleep(2)
-
-        disagg_rows = _run_phase(
-            client=client,
-            label="disaggregated",
-            pp_tg_pairs=pp_tg_pairs,
-            model_id=decode_full_id,
-            prompt_sizer=decode_prompt_sizer,
-            warmup=args.warmup,
-            repeat=args.repeat,
-            common_meta=common_meta,
-        )
-        all_rows.extend(disagg_rows)
-
-        if args.compare_baseline:
-            logger.info("Removing link and prefill instance to isolate decode_alone.")
-            with contextlib.suppress(ExoHttpError):
-                if link_id:
-                    _delete_instance_link(client, link_id)
-                    link_id = ""
-            with contextlib.suppress(ExoHttpError):
-                client.request_json("DELETE", f"/instance/{prefill_id}")
-            wait_for_instance_gone(client, prefill_id)
-            prefill_alive = False
-            time.sleep(2)
-
-            decode_alone_rows = _run_phase(
-                client=client,
-                label="decode_alone",
-                pp_tg_pairs=pp_tg_pairs,
-                model_id=decode_full_id,
-                prompt_sizer=decode_prompt_sizer,
-                warmup=args.warmup,
-                repeat=args.repeat,
-                common_meta=common_meta,
-            )
-            all_rows.extend(decode_alone_rows)
-
-            _print_diff(disagg_rows, decode_alone_rows, prefill_alone_rows)
-    finally:
-        with contextlib.suppress(ExoHttpError):
-            if link_id:
-                _delete_instance_link(client, link_id)
-        if decode_alive:
-            with contextlib.suppress(ExoHttpError):
-                client.request_json("DELETE", f"/instance/{decode_id}")
-            wait_for_instance_gone(client, decode_id)
-        if prefill_alive:
-            with contextlib.suppress(ExoHttpError):
-                client.request_json("DELETE", f"/instance/{prefill_id}")
-            wait_for_instance_gone(client, prefill_id)
-        logger.debug("Deleted both instances")
-
-    if args.stdout:
-        json.dump(all_rows, sys.stdout, indent=2, ensure_ascii=False)
-    elif args.json_out:
-        with open(args.json_out, "w", encoding="utf-8") as f:
-            json.dump(all_rows, f, indent=2, ensure_ascii=False)
-        logger.debug(f"\nWrote results JSON: {args.json_out}")
-
-    return 0
-
-
-if __name__ == "__main__":
-    sys.exit(main())

dashboard/src/lib/components/HeaderNav.svelte

@@ -1,8 +1,5 @@
 <script lang="ts">
   import { browser } from "$app/environment";
-  import { featureFlags } from "$lib/stores/app.svelte";
-
-  const showAdvanced = $derived(featureFlags()["disaggregation"] === true);
 
   interface Props {
     showHome?: boolean;
@@ -300,28 +297,5 @@
       </svg>
       <span class="hidden sm:inline">Integrations</span>
     </a>
-    {#if showAdvanced}
-      <a
-        href="/#/advanced"
-        class="text-xs md:text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-1.5 md:gap-2 cursor-pointer"
-        title="Advanced cluster settings"
-      >
-        <svg
-          class="w-4 h-4"
-          viewBox="0 0 24 24"
-          fill="none"
-          stroke="currentColor"
-          stroke-width="2"
-          stroke-linecap="round"
-          stroke-linejoin="round"
-        >
-          <circle cx="12" cy="12" r="3" />
-          <path
-            d="M19.4 15a1.65 1.65 0 0 0 .33 1.82l.06.06a2 2 0 0 1 0 2.83 2 2 0 0 1-2.83 0l-.06-.06a1.65 1.65 0 0 0-1.82-.33 1.65 1.65 0 0 0-1 1.51V21a2 2 0 0 1-4 0v-.09A1.65 1.65 0 0 0 9 19.4a1.65 1.65 0 0 0-1.82.33l-.06.06a2 2 0 0 1-2.83 0 2 2 0 0 1 0-2.83l.06-.06a1.65 1.65 0 0 0 .33-1.82 1.65 1.65 0 0 0-1.51-1H3a2 2 0 0 1 0-4h.09A1.65 1.65 0 0 0 4.6 9a1.65 1.65 0 0 0-.33-1.82l-.06-.06a2 2 0 0 1 0-2.83 2 2 0 0 1 2.83 0l.06.06a1.65 1.65 0 0 0 1.82.33H9a1.65 1.65 0 0 0 1-1.51V3a2 2 0 0 1 4 0v.09a1.65 1.65 0 0 0 1 1.51 1.65 1.65 0 0 0 1.82-.33l.06-.06a2 2 0 0 1 2.83 0 2 2 0 0 1 0 2.83l-.06.06a1.65 1.65 0 0 0-.33 1.82V9a1.65 1.65 0 0 0 1.51 1H21a2 2 0 0 1 0 4h-.09a1.65 1.65 0 0 0-1.51 1z"
-          />
-        </svg>
-        <span class="hidden sm:inline">Advanced</span>
-      </a>
-    {/if}
   </nav>
 </header>

dashboard/src/lib/components/PrefillDecodeDisaggregation.svelte

@@ -1,565 +0,0 @@
-<script lang="ts">
-  import { onMount, onDestroy } from "svelte";
-  import FamilyLogos from "$lib/components/FamilyLogos.svelte";
-  import {
-    instances,
-    instanceLinks,
-    nodeIdentities,
-    refreshState,
-    createInstanceLink,
-    updateInstanceLink,
-    deleteInstanceLink,
-    type Instance,
-  } from "$lib/stores/app.svelte";
-  import { deriveBaseModel, deriveFamily } from "$lib/utils/model_family";
-
-  type InstanceWrapper = {
-    MlxRingInstance?: Instance;
-    MlxJacclInstance?: Instance;
-    VllmInstance?: Instance;
-  };
-
-  let interval: ReturnType<typeof setInterval> | null = null;
-
-  onMount(() => {
-    refreshState();
-    interval = setInterval(refreshState, 3000);
-  });
-  onDestroy(() => {
-    if (interval) clearInterval(interval);
-  });
-
-  type InstanceRow = {
-    id: string;
-    modelId: string;
-    family: string;
-    baseModel: string;
-    nodeNames: string[];
-    nodeCount: number;
-  };
-
-  const instanceRows = $derived.by<InstanceRow[]>(() => {
-    const rows: InstanceRow[] = [];
-    const ids = nodeIdentities();
-    for (const [id, raw] of Object.entries(instances())) {
-      const wrapper = raw as InstanceWrapper;
-      const inst =
-        wrapper.MlxRingInstance ??
-        wrapper.MlxJacclInstance ??
-        wrapper.VllmInstance;
-      const modelId = inst?.shardAssignments?.modelId ?? "";
-      const nodeToRunner = inst?.shardAssignments?.nodeToRunner ?? {};
-      const nodeIds = Object.keys(nodeToRunner);
-      const nodeNames = nodeIds
-        .map((nodeId) => ids[nodeId]?.friendlyName ?? nodeId.slice(0, 6))
-        .filter((name) => !!name);
-      rows.push({
-        id,
-        modelId,
-        family: deriveFamily(modelId),
-        baseModel: deriveBaseModel(modelId),
-        nodeNames,
-        nodeCount: nodeIds.length,
-      });
-    }
-    rows.sort((a, b) => a.modelId.localeCompare(b.modelId));
-    return rows;
-  });
-
-  const instanceById = $derived(
-    Object.fromEntries(instanceRows.map((r) => [r.id, r])),
-  );
-
-  type LinkRow = {
-    linkId: string;
-    prefill: string[];
-    decode: string[];
-    families: string[];
-    multiNode: boolean;
-  };
-
-  const linkRows = $derived.by<LinkRow[]>(() => {
-    const rows: LinkRow[] = [];
-    for (const [, link] of Object.entries(instanceLinks())) {
-      const fams = new Set<string>();
-      let multiNode = false;
-      for (const id of [...link.prefillInstances, ...link.decodeInstances]) {
-        const r = instanceById[id];
-        if (r && r.baseModel) fams.add(r.baseModel.toLowerCase());
-        if (r && r.nodeCount > 1) multiNode = true;
-      }
-      rows.push({
-        linkId: link.linkId,
-        prefill: link.prefillInstances,
-        decode: link.decodeInstances,
-        families: Array.from(fams),
-        multiNode,
-      });
-    }
-    return rows;
-  });
-
-  let editingLinkId = $state<string | null>(null);
-  let editingPrefill = $state<Set<string>>(new Set());
-  let editingDecode = $state<Set<string>>(new Set());
-  let saving = $state(false);
-  let errorMessage = $state<string | null>(null);
-
-  function startCreate() {
-    editingLinkId = "new";
-    editingPrefill = new Set();
-    editingDecode = new Set();
-    errorMessage = null;
-  }
-
-  function startEdit(row: LinkRow) {
-    editingLinkId = row.linkId;
-    editingPrefill = new Set(row.prefill);
-    editingDecode = new Set(row.decode);
-    errorMessage = null;
-  }
-
-  function cancelEdit() {
-    editingLinkId = null;
-    editingPrefill = new Set();
-    editingDecode = new Set();
-    errorMessage = null;
-  }
-
-  type Role = "prefill" | "decode" | "none";
-
-  function roleOf(id: string): Role {
-    if (editingPrefill.has(id)) return "prefill";
-    if (editingDecode.has(id)) return "decode";
-    return "none";
-  }
-
-  function setRole(id: string, role: Role) {
-    const p = new Set(editingPrefill);
-    const d = new Set(editingDecode);
-    p.delete(id);
-    d.delete(id);
-    if (role === "prefill") p.add(id);
-    if (role === "decode") d.add(id);
-    editingPrefill = p;
-    editingDecode = d;
-  }
-
-  const editingFamilies = $derived.by<string[]>(() => {
-    const fams = new Set<string>();
-    for (const id of [...editingPrefill, ...editingDecode]) {
-      const r = instanceById[id];
-      if (r && r.baseModel) fams.add(r.baseModel.toLowerCase());
-    }
-    return Array.from(fams);
-  });
-
-  const editingMultiNode = $derived.by<string[]>(() => {
-    const names: string[] = [];
-    for (const id of [...editingPrefill, ...editingDecode]) {
-      const r = instanceById[id];
-      if (r && r.nodeCount > 1) {
-        names.push(r.baseModel || r.modelId);
-      }
-    }
-    return names;
-  });
-
-  const editingMismatch = $derived(editingFamilies.length > 1);
-  const canSave = $derived(
-    editingLinkId !== null &&
-      editingPrefill.size > 0 &&
-      editingDecode.size > 0 &&
-      !saving,
-  );
-
-  async function save() {
-    if (editingLinkId === null) return;
-    saving = true;
-    errorMessage = null;
-    try {
-      const prefill = Array.from(editingPrefill);
-      const decode = Array.from(editingDecode);
-      if (editingLinkId === "new") {
-        await createInstanceLink(prefill, decode);
-      } else {
-        await updateInstanceLink(editingLinkId, prefill, decode);
-      }
-      cancelEdit();
-      await refreshState();
-    } catch (err) {
-      errorMessage = err instanceof Error ? err.message : String(err);
-    } finally {
-      saving = false;
-    }
-  }
-
-  async function remove(linkId: string) {
-    if (!confirm("Remove this routing?")) return;
-    try {
-      await deleteInstanceLink(linkId);
-      if (editingLinkId === linkId) cancelEdit();
-      await refreshState();
-    } catch (err) {
-      errorMessage = err instanceof Error ? err.message : String(err);
-    }
-  }
-</script>
-
-<div class="font-mono text-foreground">
-  <div class="mb-6 space-y-4">
-    <details open class="group [&_summary::-webkit-details-marker]:hidden">
-      <summary
-        class="cursor-pointer list-none text-exo-yellow text-xs font-mono tracking-widest uppercase flex items-center gap-2 hover:opacity-80 transition-opacity"
-      >
-        <span
-          class="inline-block transition-transform group-open:rotate-90 text-exo-light-gray"
-          >▶</span
-        >
-        Prefill vs Decode
-      </summary>
-      <div class="mt-2 text-white/80 text-sm leading-relaxed">
-        Prefill is the compute-heavy pass that consumes the entire prompt and
-        builds a KV cache. Decode is the memory-bandwidth-bound loop that emits
-        tokens sequentially from that cache. The two phases have very different
-        bottlenecks, so running them on different hardware can be substantially
-        faster than doing both on one node.
-      </div>
-    </details>
-    <details class="group [&_summary::-webkit-details-marker]:hidden">
-      <summary
-        class="cursor-pointer list-none text-exo-yellow text-xs font-mono tracking-widest uppercase flex items-center gap-2 hover:opacity-80 transition-opacity"
-      >
-        <span
-          class="inline-block transition-transform group-open:rotate-90 text-exo-light-gray"
-          >▶</span
-        >
-        Linking Instances
-      </summary>
-      <div class="mt-2 text-white/80 text-sm leading-relaxed space-y-2">
-        <p>
-          A linked route here tells the cluster: when a request is sent to a
-          model in that cluster, the decode node (or the least active one if
-          there are multiple) will handle it. If it decides it must do a lot of
-          prefill not already cached in the prefix cache, it routes the request
-          to the prefill node over TCP IP. The prefill node streams the KV cache
-          back to the decode node which picks up from there.
-        </p>
-        <p>
-          Linked instances must be running the same model family — KV layouts
-          differ across architectures. More on the <a
-            class="text-exo-yellow underline underline-offset-2 hover:text-exo-yellow-darker transition-colors"
-            href="https://blog.exolabs.net/nvidia-dgx-spark/"
-            target="_blank"
-            rel="noreferrer noopener">blog</a
-          >.
-        </p>
-      </div>
-    </details>
-  </div>
-
-  {#if errorMessage}
-    <div
-      class="mb-4 px-4 py-3 bg-red-500/10 border border-red-500/40 text-red-300 text-sm"
-    >
-      {errorMessage}
-    </div>
-  {/if}
-
-  <section class="mt-12">
-    <h2
-      class="text-exo-yellow text-xs font-mono tracking-widest uppercase m-0 mb-3"
-    >
-      Existing routes
-    </h2>
-
-    {#if linkRows.length === 0}
-      {#if editingLinkId === null}
-        <div class="flex items-center justify-between">
-          <p class="text-exo-light-gray italic text-sm m-0">
-            No routes yet. Create one to enable remote prefill.
-          </p>
-          <button
-            class="px-3 py-1.5 text-xs font-mono tracking-wider uppercase bg-exo-yellow/15 border border-exo-yellow/50 text-exo-yellow hover:bg-exo-yellow/25 hover:border-exo-yellow/80 transition-colors"
-            onclick={startCreate}
-          >
-            + New route
-          </button>
-        </div>
-      {/if}
-    {:else}
-      {#if editingLinkId === null}
-        <div class="flex justify-end mb-3">
-          <button
-            class="px-3 py-1.5 text-xs font-mono tracking-wider uppercase bg-exo-yellow/15 border border-exo-yellow/50 text-exo-yellow hover:bg-exo-yellow/25 hover:border-exo-yellow/80 transition-colors"
-            onclick={startCreate}
-          >
-            + New route
-          </button>
-        </div>
-      {/if}
-      <div
-        class="bg-exo-dark-gray/60 border border-exo-medium-gray/40 flex flex-col"
-      >
-        {#each linkRows as row (row.linkId)}
-          {#if editingLinkId !== row.linkId}
-            <article
-              class="p-4 border-b border-exo-light-gray/25 last:border-b-0"
-            >
-              {#if row.multiNode}
-                <div
-                  class="mb-3 px-3 py-2 bg-red-500/10 border border-red-500/40 text-red-300 text-xs tracking-wide"
-                >
-                  ⚠ Multi-node instance detected. Remote prefill currently only
-                  works on single-node (rank-0) instances. This route will not
-                  function until that's supported.
-                </div>
-              {/if}
-              {#if row.families.length > 1}
-                <div
-                  class="mb-3 px-3 py-2 bg-amber-500/10 border border-amber-500/40 text-amber-300 text-xs tracking-wide"
-                >
-                  ⚠ Mixed model families: {row.families.join(", ")}
-                </div>
-              {/if}
-              <div
-                class="grid grid-cols-[1fr_auto_1fr_auto] items-center gap-x-3 gap-y-2"
-              >
-                <span
-                  class="inline-block justify-self-start text-[10px] font-mono tracking-widest uppercase px-2 py-0.5 bg-exo-yellow/15 border border-exo-yellow/40 text-exo-yellow"
-                  >Prefill</span
-                >
-                <span></span>
-                <span
-                  class="inline-block justify-self-start text-[10px] font-mono tracking-widest uppercase px-2 py-0.5 bg-exo-medium-gray/40 border border-exo-medium-gray/60 text-foreground"
-                  >Decode</span
-                >
-                <span></span>
-                <div class="min-w-0">
-                  <ul class="list-none p-0 m-0 flex flex-col gap-2">
-                    {#each row.prefill as id (id)}
-                      {@const r = instanceById[id]}
-                      {#if r}
-                        <li
-                          class="flex items-center gap-2 px-2.5 py-2 bg-exo-medium-gray/20 border border-exo-medium-gray/40"
-                        >
-                          <FamilyLogos family={r.family} />
-                          <div class="min-w-0 flex-1">
-                            <div
-                              class="text-exo-yellow text-xs font-mono truncate"
-                            >
-                              {r.baseModel || r.modelId}
-                            </div>
-                            <div
-                              class="text-exo-light-gray text-[11px] truncate"
-                            >
-                              {r.nodeNames.join(", ") || "?"}{r.nodeCount > 1
-                                ? ` (${r.nodeCount} nodes)`
-                                : ""}
-                            </div>
-                            <div
-                              class="text-exo-light-gray/40 text-[10px] font-mono truncate"
-                              title={r.id}
-                            >
-                              {r.id.slice(0, 8)}
-                            </div>
-                          </div>
-                        </li>
-                      {/if}
-                    {/each}
-                  </ul>
-                </div>
-                <div class="text-exo-yellow/60 text-xl px-2" aria-hidden="true">
-                  →
-                </div>
-                <div class="min-w-0">
-                  <ul class="list-none p-0 m-0 flex flex-col gap-2">
-                    {#each row.decode as id (id)}
-                      {@const r = instanceById[id]}
-                      {#if r}
-                        <li
-                          class="flex items-center gap-2 px-2.5 py-2 bg-exo-medium-gray/20 border border-exo-medium-gray/40"
-                        >
-                          <FamilyLogos family={r.family} />
-                          <div class="min-w-0 flex-1">
-                            <div
-                              class="text-exo-yellow text-xs font-mono truncate"
-                            >
-                              {r.baseModel || r.modelId}
-                            </div>
-                            <div
-                              class="text-exo-light-gray text-[11px] truncate"
-                            >
-                              {r.nodeNames.join(", ") || "?"}{r.nodeCount > 1
-                                ? ` (${r.nodeCount} nodes)`
-                                : ""}
-                            </div>
-                            <div
-                              class="text-exo-light-gray/40 text-[10px] font-mono truncate"
-                              title={r.id}
-                            >
-                              {r.id.slice(0, 8)}
-                            </div>
-                          </div>
-                        </li>
-                      {/if}
-                    {/each}
-                  </ul>
-                </div>
-                <div class="flex gap-2 pl-3">
-                  <button
-                    class="px-2 py-0.5 text-[11px] font-mono tracking-wider uppercase bg-exo-medium-gray/30 border border-exo-medium-gray/60 rounded text-foreground hover:border-exo-yellow/60 hover:text-exo-yellow disabled:opacity-40 disabled:cursor-not-allowed transition-colors"
-                    onclick={() => startEdit(row)}
-                    disabled={editingLinkId !== null}
-                  >
-                    Edit
-                  </button>
-                  <button
-                    class="px-2 py-0.5 text-[11px] font-mono tracking-wider uppercase bg-red-500/15 border border-red-500/40 rounded text-red-300 hover:bg-red-500/25 transition-colors"
-                    onclick={() => remove(row.linkId)}
-                  >
-                    Remove
-                  </button>
-                </div>
-              </div>
-            </article>
-          {/if}
-        {/each}
-      </div>
-    {/if}
-  </section>
-
-  {#if editingLinkId !== null && instanceRows.length === 0}
-    <section
-      class="mt-6 bg-exo-dark-gray/60 border border-exo-yellow/30 px-4 py-2.5 flex items-center justify-between gap-3"
-    >
-      <span class="text-exo-light-gray italic text-sm font-mono"
-        >No instances available.</span
-      >
-      <button
-        class="px-3 py-1 text-xs font-mono tracking-wider uppercase bg-exo-medium-gray/30 border border-exo-medium-gray/60 rounded text-foreground hover:border-exo-yellow/60 transition-colors"
-        onclick={cancelEdit}
-      >
-        Cancel
-      </button>
-    </section>
-  {:else if editingLinkId !== null}
-    <section class="mt-6 bg-exo-dark-gray/60 border border-exo-yellow/30 p-5">
-      <h2
-        class="text-exo-yellow text-xs font-mono tracking-widest uppercase m-0 mb-3"
-      >
-        {editingLinkId === "new" ? "New route" : "Edit route"}
-      </h2>
-
-      {#if editingMismatch}
-        <div
-          class="mb-3 px-3 py-2 bg-amber-500/10 border border-amber-500/40 text-amber-300 text-xs tracking-wide"
-        >
-          ⚠ Selected instances span multiple model families: <strong
-            >{editingFamilies.join(", ")}</strong
-          >. Linking across families produces a corrupt KV cache.
-        </div>
-      {/if}
-
-      {#if editingMultiNode.length > 0}
-        <div
-          class="mb-3 px-3 py-2 bg-red-500/10 border border-red-500/40 text-red-300 text-xs tracking-wide"
-        >
-          ⚠ Multi-node instance(s) selected: <strong
-            >{editingMultiNode.join(", ")}</strong
-          >. Remote prefill currently only works on single-node instances. This
-          route will not function until multi-node support lands.
-        </div>
-      {/if}
-
-      <p class="text-exo-light-gray text-xs mb-4">
-        Pick a role for each instance:
-        <span class="text-exo-yellow">Prefill</span>
-        serves KV cache,
-        <span class="text-foreground">Decode</span> consumes it.
-      </p>
-      <div
-        class="grid gap-2.5"
-        style="grid-template-columns: repeat(auto-fill, minmax(360px, 1fr));"
-      >
-        {#each instanceRows as row (row.id)}
-          {@const role = roleOf(row.id)}
-          <div
-            class="border p-3 flex flex-col gap-2.5 transition-colors {role ===
-            'prefill'
-              ? 'border-exo-yellow/60 bg-exo-dark-gray/60'
-              : role === 'decode'
-                ? 'border-exo-light-gray/60 bg-exo-dark-gray/60'
-                : 'border-exo-medium-gray/40 bg-exo-dark-gray/40'}"
-          >
-            <div class="flex items-center gap-2">
-              <FamilyLogos family={row.family} />
-              <div class="min-w-0 flex-1">
-                <div class="text-exo-yellow text-xs font-mono truncate">
-                  {row.baseModel || row.modelId}
-                </div>
-                <div class="text-exo-light-gray text-[11px] truncate">
-                  {row.nodeNames.join(", ") || "?"}{row.nodeCount > 1
-                    ? ` (${row.nodeCount} nodes)`
-                    : ""}
-                </div>
-                <div
-                  class="text-exo-light-gray/40 text-[10px] font-mono truncate"
-                  title={row.id}
-                >
-                  {row.id.slice(0, 8)}
-                </div>
-              </div>
-              {#if row.nodeCount > 1}
-                <span
-                  class="text-[9px] font-mono tracking-widest uppercase px-1.5 py-0.5 bg-red-500/15 border border-red-500/40 text-red-300"
-                  title="Multi-node instances are not supported by remote prefill yet."
-                  >Unsupported</span
-                >
-              {/if}
-            </div>
-            <div
-              class="flex rounded-md overflow-hidden border border-exo-light-gray/40 divide-x divide-exo-light-gray/40"
-            >
-              <button
-                class="flex-1 px-2 py-1 text-[11px] font-mono tracking-wider uppercase transition-colors {role ===
-                'prefill'
-                  ? 'bg-exo-yellow/20 text-exo-yellow'
-                  : 'bg-transparent text-white/80 hover:text-exo-yellow'}"
-                onclick={() =>
-                  setRole(row.id, role === "prefill" ? "none" : "prefill")}
-                >Prefill</button
-              >
-              <button
-                class="flex-1 px-2 py-1 text-[11px] font-mono tracking-wider uppercase transition-colors {role ===
-                'decode'
-                  ? 'bg-exo-medium-gray/50 text-foreground'
-                  : 'bg-transparent text-white/80 hover:text-foreground'}"
-                onclick={() =>
-                  setRole(row.id, role === "decode" ? "none" : "decode")}
-                >Decode</button
-              >
-            </div>
-          </div>
-        {/each}
-      </div>
-
-      <div class="flex gap-2 mt-5 justify-end">
-        <button
-          class="px-3 py-1.5 text-xs font-mono tracking-wider uppercase bg-exo-yellow/15 border border-exo-yellow/50 text-exo-yellow hover:bg-exo-yellow/25 hover:border-exo-yellow/80 disabled:opacity-40 disabled:cursor-not-allowed transition-colors"
-          onclick={save}
-          disabled={!canSave}
-        >
-          {saving ? "Saving..." : "Save route"}
-        </button>
-        <button
-          class="px-3 py-1.5 text-xs font-mono tracking-wider uppercase bg-exo-medium-gray/30 border border-exo-medium-gray/60 text-foreground hover:border-exo-yellow/60 disabled:opacity-40 disabled:cursor-not-allowed transition-colors"
-          onclick={cancelEdit}
-          disabled={saving}
-        >
-          Cancel
-        </button>
-      </div>
-    </section>
-  {/if}
-</div>

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

@@ -74,12 +74,6 @@ export interface Instance {
   };
 }
 
-export interface RawInstanceLink {
-  linkId: string;
-  prefillInstances: string[];
-  decodeInstances: string[];
-}
-
 // Granular node state types from the new state structure
 interface RawNodeIdentity {
   modelId?: string;
@@ -229,7 +223,6 @@ interface RawStateResponse {
     }
   >;
   runners?: Record<string, unknown>;
-  instanceLinks?: Record<string, RawInstanceLink>;
   downloads?: Record<string, unknown[]>;
   // New granular node state fields
   nodeIdentities?: Record<string, RawNodeIdentity>;
@@ -548,8 +541,6 @@ class AppStore {
   topologyData = $state<TopologyData | null>(null);
   instances = $state<Record<string, unknown>>({});
   runners = $state<Record<string, unknown>>({});
-  instanceLinks = $state<Record<string, RawInstanceLink>>({});
-  featureFlags = $state<Record<string, boolean>>({});
   downloads = $state<Record<string, unknown[]>>({});
   nodeDisk = $state<
     Record<
@@ -1283,7 +1274,6 @@ class AppStore {
 
   startPolling() {
     this.fetchState();
-    this.fetchFeatureFlags();
     this.fetchInterval = setInterval(() => this.fetchState(), 1000);
   }
 
@@ -1295,16 +1285,6 @@ class AppStore {
     this.stopPreviewsPolling();
   }
 
-  async fetchFeatureFlags() {
-    try {
-      const response = await fetch("/v1/feature-flags");
-      if (!response.ok) return;
-      this.featureFlags = await response.json();
-    } catch {
-      // Silently ignore — defaults to all-disabled.
-    }
-  }
-
   async fetchState() {
     try {
       const response = await fetch("/state");
@@ -1330,11 +1310,6 @@ class AppStore {
       if (data.runners) {
         this.runners = data.runners;
       }
-      if (data.instanceLinks) {
-        this.instanceLinks = data.instanceLinks;
-      } else {
-        this.instanceLinks = {};
-      }
       if (data.downloads) {
         this.downloads = data.downloads;
       }
@@ -1695,15 +1670,7 @@ class AppStore {
               }
             }
           }
-          const out: {
-            role: string;
-            content: string;
-            reasoning_content?: string;
-          } = { role: m.role, content: msgContent };
-          if (m.role === "assistant" && m.thinking) {
-            out.reasoning_content = m.thinking;
-          }
-          return out;
+          return { role: m.role, content: msgContent };
         }),
       ];
 
@@ -1910,15 +1877,7 @@ class AppStore {
       const apiMessages = [
         systemPrompt,
         ...targetConversation.messages.slice(0, -1).map((m) => {
-          const out: {
-            role: string;
-            content: string;
-            reasoning_content?: string;
-          } = { role: m.role, content: m.content };
-          if (m.role === "assistant" && m.thinking) {
-            out.reasoning_content = m.thinking;
-          }
-          return out;
+          return { role: m.role, content: m.content };
         }),
       ];
 
@@ -2449,15 +2408,10 @@ class AppStore {
               contentParts.push({ type: "text", text: textContent });
             }
 
-            const out: {
-              role: string;
-              content: typeof contentParts;
-              reasoning_content?: string;
-            } = { role: m.role, content: contentParts };
-            if (m.role === "assistant" && m.thinking) {
-              out.reasoning_content = m.thinking;
-            }
-            return out;
+            return {
+              role: m.role,
+              content: contentParts,
+            };
           }
 
           // Text-only message (original path)
@@ -2475,15 +2429,10 @@ class AppStore {
             }
           }
 
-          const out: {
-            role: string;
-            content: string;
-            reasoning_content?: string;
-          } = { role: m.role, content: msgContent };
-          if (m.role === "assistant" && m.thinking) {
-            out.reasoning_content = m.thinking;
-          }
-          return out;
+          return {
+            role: m.role,
+            content: msgContent,
+          };
         }),
       ];
 
@@ -3332,60 +3281,6 @@ class AppStore {
     }
   }
 
-  async createInstanceLink(
-    prefillInstances: string[],
-    decodeInstances: string[],
-  ): Promise<void> {
-    const response = await fetch("/v1/instance-links", {
-      method: "POST",
-      headers: { "Content-Type": "application/json" },
-      body: JSON.stringify({
-        prefill_instances: prefillInstances,
-        decode_instances: decodeInstances,
-      }),
-    });
-    if (!response.ok) {
-      throw new Error(
-        `Failed to create instance link: ${response.status} ${await response.text()}`,
-      );
-    }
-  }
-
-  async updateInstanceLink(
-    linkId: string,
-    prefillInstances: string[],
-    decodeInstances: string[],
-  ): Promise<void> {
-    const response = await fetch(
-      `/v1/instance-links/${encodeURIComponent(linkId)}`,
-      {
-        method: "PUT",
-        headers: { "Content-Type": "application/json" },
-        body: JSON.stringify({
-          prefill_instances: prefillInstances,
-          decode_instances: decodeInstances,
-        }),
-      },
-    );
-    if (!response.ok) {
-      throw new Error(
-        `Failed to update instance link: ${response.status} ${await response.text()}`,
-      );
-    }
-  }
-
-  async deleteInstanceLink(linkId: string): Promise<void> {
-    const response = await fetch(
-      `/v1/instance-links/${encodeURIComponent(linkId)}`,
-      { method: "DELETE" },
-    );
-    if (!response.ok) {
-      throw new Error(
-        `Failed to delete instance link: ${response.status} ${await response.text()}`,
-      );
-    }
-  }
-
   /**
    * Delete a downloaded model from a specific node
    */
@@ -3484,19 +3379,6 @@ export const prefillProgress = () => appStore.prefillProgress;
 export const topologyData = () => appStore.topologyData;
 export const instances = () => appStore.instances;
 export const runners = () => appStore.runners;
-export const instanceLinks = () => appStore.instanceLinks;
-export const featureFlags = () => appStore.featureFlags;
-export const createInstanceLink = (
-  prefillInstances: string[],
-  decodeInstances: string[],
-) => appStore.createInstanceLink(prefillInstances, decodeInstances);
-export const updateInstanceLink = (
-  linkId: string,
-  prefillInstances: string[],
-  decodeInstances: string[],
-) => appStore.updateInstanceLink(linkId, prefillInstances, decodeInstances);
-export const deleteInstanceLink = (linkId: string) =>
-  appStore.deleteInstanceLink(linkId);
 export const downloads = () => appStore.downloads;
 export const nodeDisk = () => appStore.nodeDisk;
 export const placementPreviews = () => appStore.placementPreviews;

dashboard/src/lib/utils/model_family.ts

@@ -1,44 +0,0 @@
-// Mirrors src/exo/shared/models/model_cards.py:derive_base_model
-const QUANT_SUFFIXES = new RegExp(
-  "[-_ ](?:MLX|MXFP[0-9]+|NVFP[0-9]+|GPTQ|AWQ|GGUF|fp16|bf16|fp8|int[0-9]+|[0-9]+(?:\\.[0-9]+)?bit|Q[0-9]+(?:_[A-Z0-9]+)?|gs[0-9]+)" +
-    "(?:[-_ ](?:MLX|Q[0-9]+|Int[0-9]+|[A-Z0-9]+|gs[0-9]+))*$",
-  "i",
-);
-
-function normalize(s: string): string {
-  return s
-    .replaceAll("-", " ")
-    .replaceAll("_", " ")
-    .replaceAll("  ", " ")
-    .trim();
-}
-
-export function deriveBaseModel(modelId: string): string {
-  const short = modelId.includes("/")
-    ? (modelId.split("/").pop() ?? modelId)
-    : modelId;
-  const stripped = short.replace(QUANT_SUFFIXES, "");
-  return normalize(stripped);
-}
-
-export function baseModelsCompatible(a: string, b: string): boolean {
-  return deriveBaseModel(a).toLowerCase() === deriveBaseModel(b).toLowerCase();
-}
-
-// Mirrors src/exo/shared/models/model_cards.py:derive_family
-export function deriveFamily(modelId: string): string {
-  const short = modelId.includes("/")
-    ? (modelId.split("/").pop() ?? modelId)
-    : modelId;
-  const stripped = short
-    .replace(QUANT_SUFFIXES, "")
-    .toLowerCase()
-    .replaceAll("_", "-");
-  const parts = stripped.split(/[-.]/);
-  const familyParts: string[] = [];
-  for (const p of parts) {
-    if (/^\d+$/.test(p) || /^\d+[bm]?$/i.test(p)) break;
-    familyParts.push(p);
-  }
-  return familyParts.length > 0 ? familyParts.join("-") : stripped;
-}

dashboard/src/routes/+page.svelte

@@ -3435,7 +3435,6 @@
             >
               <li>Connect nodes with TB5 cables</li>
               <li>Boot to Recovery (hold power 10s → Options)</li>
-              <li>Open Terminal from the Utilities menu</li>
               <li>
                 Run
                 <code class="text-yellow-300 bg-yellow-400/10 px-1 rounded"
@@ -4823,7 +4822,6 @@
                 >
                   <li>Connect nodes with TB5 cables</li>
                   <li>Boot to Recovery (hold power 10s → Options)</li>
-                  <li>Open Terminal from the Utilities menu</li>
                   <li>
                     Run
                     <code class="text-yellow-300 bg-yellow-400/10 px-1 rounded"
@@ -4970,7 +4968,6 @@
                   >
                     <li>Connect nodes with TB5 cables</li>
                     <li>Boot to Recovery (hold power 10s → Options)</li>
-                    <li>Open Terminal from the Utilities menu</li>
                     <li>
                       Run
                       <code

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

@@ -1,81 +0,0 @@
-<script lang="ts">
-  import { browser } from "$app/environment";
-  import HeaderNav from "$lib/components/HeaderNav.svelte";
-  import PrefillDecodeDisaggregation from "$lib/components/PrefillDecodeDisaggregation.svelte";
-  import { featureFlags, refreshState } from "$lib/stores/app.svelte";
-  import { onMount } from "svelte";
-
-  type TabId = "prefill-decode";
-
-  const tabs: { id: TabId; label: string }[] = [
-    { id: "prefill-decode", label: "Prefill / Decode" },
-  ];
-
-  let activeTab = $state<TabId>(tabs[0].id);
-  let flagsLoaded = $state(false);
-
-  onMount(() => {
-    refreshState().finally(() => {
-      flagsLoaded = true;
-    });
-  });
-
-  const flags = $derived(featureFlags());
-  const enabled = $derived(flags["disaggregation"] === true);
-
-  $effect(() => {
-    if (browser && flagsLoaded && !enabled) {
-      // No advanced features enabled — bounce home.
-      window.location.hash = "/";
-    }
-  });
-</script>
-
-<div class="min-h-screen bg-exo-dark-gray flex flex-col">
-  <HeaderNav />
-
-  <main class="flex-1 max-w-[1100px] mx-auto w-full px-4 md:px-6 py-8">
-    {#if !flagsLoaded}
-      <div class="text-exo-light-gray/60 text-sm">Loading…</div>
-    {:else if !enabled}
-      <div class="text-exo-light-gray/60 text-sm">
-        No advanced features enabled. Set <code
-          class="text-exo-yellow font-mono">ENABLE_DISAGGREGATION=true</code
-        > on the cluster to access prefill/decode disaggregation.
-      </div>
-    {:else}
-      <div class="mb-4">
-        <h1
-          class="text-white text-xl md:text-2xl font-semibold tracking-wide mb-2"
-        >
-          Advanced
-        </h1>
-        <p class="text-exo-light-gray/60 text-sm">
-          Cluster-level configuration. Most users don't need anything here.
-        </p>
-      </div>
-
-      <div
-        class="flex flex-wrap gap-2 mb-6 border-b border-exo-light-gray/10 pb-3"
-      >
-        {#each tabs as tab (tab.id)}
-          <button
-            onclick={() => (activeTab = tab.id)}
-            class="px-3 py-1.5 text-xs rounded-md transition-all cursor-pointer
-              {activeTab === tab.id
-              ? 'bg-exo-yellow/15 text-exo-yellow border border-exo-yellow/30'
-              : 'text-exo-light-gray/60 hover:text-white/80 border border-transparent hover:border-exo-light-gray/20'}"
-          >
-            {tab.label}
-          </button>
-        {/each}
-      </div>
-
-      <div class="space-y-4">
-        {#if activeTab === "prefill-decode"}
-          <PrefillDecodeDisaggregation />
-        {/if}
-      </div>
-    {/if}
-  </main>
-</div>

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

@@ -14,7 +14,6 @@
 
   let modelCapabilities = $state<Record<string, string[]>>({});
   let modelContextLengths = $state<Record<string, number>>({});
-  let modelReasoningDialects = $state<Record<string, string>>({});
 
   const runningModels = $derived.by(() => {
     const models: string[] = [];
@@ -133,7 +132,6 @@
     for (const modelId of runningModels) {
       const caps = modelCapabilities[modelId] || [];
       const ctxLen = modelContextLengths[modelId] || 0;
-      const dialect = modelReasoningDialects[modelId];
       const entry: Record<string, unknown> = { name: modelId };
       if (ctxLen > 0) {
         entry.limit = { context: ctxLen, output: Math.min(ctxLen, 16384) };
@@ -141,27 +139,6 @@
       if (caps.includes("vision")) {
         entry.modalities = { input: ["text", "image"], output: ["text"] };
       }
-      // Reasoning round-trip: opencode's `interleaved` field tells the
-      // openai-compatible adapter to send the assistant's prior
-      // reasoning_content back in subsequent turns. Emit it for dialects
-      // whose chat templates use prior reasoning:
-      //   - `tool_conditional` (DeepSeek V3.2 / V4): wrapper preserves all
-      //     reasoning when tools are present.
-      //   - `post_last_user` (Qwen3-Thinking, GLM 4.5+, MiniMax M2.x):
-      //     Jinja template reads reasoning_content for assistant turns since
-      //     the last user message — exactly the tool-chain window.
-      //   - `channel` (gpt-oss / Harmony): the model's Jinja template reads
-      //     `message.thinking` rather than `message.reasoning_content`, but
-      //     the server bridges `reasoning_content` → `thinking` before
-      //     rendering, so the round-trip works through the standard field.
-      // `suffix` (Kimi): reasoning lives in content; no separate field path.
-      if (
-        dialect === "tool_conditional" ||
-        dialect === "post_last_user" ||
-        dialect === "channel"
-      ) {
-        entry.interleaved = { field: "reasoning_content" };
-      }
       models[modelId] = entry;
     }
     if (Object.keys(models).length === 0) {
@@ -373,25 +350,16 @@
     try {
       const resp = await fetch("/v1/models");
       const data = (await resp.json()) as {
-        data: {
-          id: string;
-          capabilities: string[];
-          context_length: number;
-          reasoning_dialect?: string;
-        }[];
+        data: { id: string; capabilities: string[]; context_length: number }[];
       };
       const caps: Record<string, string[]> = {};
       const ctxs: Record<string, number> = {};
-      const dialects: Record<string, string> = {};
       for (const model of data.data) {
         caps[model.id] = model.capabilities || [];
         if (model.context_length > 0) ctxs[model.id] = model.context_length;
-        if (model.reasoning_dialect)
-          dialects[model.id] = model.reasoning_dialect;
       }
       modelCapabilities = caps;
       modelContextLengths = ctxs;
-      modelReasoningDialects = dialects;
     } catch {
       /* ignore */
     }

flake.lock

@@ -47,11 +47,11 @@
         "rust-analyzer-src": "rust-analyzer-src"
       },
       "locked": {
-        "lastModified": 1777708550,
-        "narHash": "sha256-Qif3UXT0l5OQq8H9pRWt4/ia4gF48MWK2oHKL8uVx8U=",
+        "lastModified": 1775807984,
+        "narHash": "sha256-Redoe3D9zGN5I9QPHWL9vfMVQBehY1fKsMiRXQ83X3w=",
         "owner": "nix-community",
         "repo": "fenix",
-        "rev": "74c1591efaff494756b8d35ebe357c6c2bbdca96",
+        "rev": "fcf90c0c4d368b2ca917a7afa6d08e98a397e5fd",
         "type": "github"
       },
       "original": {
@@ -218,11 +218,11 @@
     "rust-analyzer-src": {
       "flake": false,
       "locked": {
-        "lastModified": 1777639980,
-        "narHash": "sha256-6d7Hdurvbjc5uwJuc0YiK7rZBGj6Gs3uzfBFcTs+xCc=",
+        "lastModified": 1775745684,
+        "narHash": "sha256-8MbfLwd60FNa8dRFkjE+G3TT/x21G3Rsplm1bMBQUtU=",
         "owner": "rust-lang",
         "repo": "rust-analyzer",
-        "rev": "64cdaeb06f69b6b769a492edd88b022ae88e8ca2",
+        "rev": "64ddb549bc9a70d011328746fa46a8883f937b6b",
         "type": "github"
       },
       "original": {

flake.nix

@@ -146,7 +146,7 @@
                 config.treefmt.build.wrapper
 
                 # PYTHON
-                #self'.packages.editableVenv
+                self'.packages.editableVenv
                 uv
 
                 # RUST

pyproject.toml

@@ -15,11 +15,11 @@ dependencies = [
   "huggingface-hub>=1.8.0",
   "psutil>=7.0.0",
   "loguru>=0.7.3",
-  "exo-pyo3-bindings",                         # rust bindings
+  "exo-pyo3-bindings",                       # rust bindings
   "anyio==4.11.0",
   "mlx==0.31.2; sys_platform == 'darwin'",
   "mlx-lm; sys_platform=='darwin'",
-  "tiktoken>=0.12.0",                          # required for kimi k2 tokenizer
+  "tiktoken>=0.12.0",                        # required for kimi k2 tokenizer
   "hypercorn>=0.18.0",
   "openai-harmony>=0.0.8",
   "httpx>=0.28.1",
@@ -28,8 +28,8 @@ dependencies = [
   "python-multipart>=0.0.21",
   "msgspec>=0.19.0",
   "zstandard>=0.23.0",
-  "mlx-vlm>=0.3.11; sys_platform == 'darwin'",
-  "transformers>=5.6.2",
+  "mlx-vlm>=0.3.11",
+  "transformers>=5.0.0,<5.4.0",
 ]
 
 [project.scripts]
@@ -52,21 +52,18 @@ cpu = [
   "mlx==0.31.1; sys_platform == 'linux'",
   "mlx-cpu==0.31.1; sys_platform == 'linux'",
   "mlx-lm; sys_platform == 'linux'",
-  "mlx-vlm>=0.3.11; sys_platform== 'linux'",
   "torch>=2.10.0; sys_platform == 'linux'",
 ]
 cuda12 = [
   "mlx==0.31.1; sys_platform == 'linux'",
   "mlx-cuda-12==0.31.1; sys_platform == 'linux'",
   "mlx-lm; sys_platform == 'linux'",
-  "mlx-vlm>=0.3.11; sys_platform== 'linux'",
   "torch>=2.10.0; sys_platform == 'linux'",
 ]
 cuda13 = [
   "mlx==0.31.1; sys_platform == 'linux'",
   "mlx-cuda-13==0.31.1; sys_platform == 'linux'",
   "mlx-lm; sys_platform == 'linux'",
-  "mlx-vlm>=0.3.11; sys_platform== 'linux'",
   "torch>=2.10.0; sys_platform == 'linux'",
 ]
 
@@ -80,7 +77,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/deepseek-v4" }
+mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/fix-arrayscache-leak" }
 torch = [
   { index = "pytorch-cu130", marker = "sys_platform == 'linux' and extra == 'cuda13' and extra != 'cpu' and extra != 'cuda12'" },
   { index = "pytorch-cu120", marker = "sys_platform == 'linux' and extra == 'cuda12' and extra != 'cpu' and extra != 'cuda13'" },
@@ -157,7 +154,7 @@ required-version = ">=0.8.6"
 prerelease = "allow"
 environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
 conflicts = [[{ extra = "cuda12" }, { extra = "cuda13" }, { extra = "cpu" }]]
-constraint-dependencies = ["transformers>=5.6.2"]
+constraint-dependencies = ["transformers>=5.0.0,<5.4.0"]
 override-dependencies = [
   "mlx==0.31.1; sys_platform=='linux'",
   "mlx; sys_platform=='darwin'",

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

@@ -8,7 +8,6 @@ family = "deepseek"
 quantization = "4bit"
 base_model = "DeepSeek V3.1"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "post_last_user"
 
 context_length = 131072
 

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

@@ -8,7 +8,6 @@ family = "deepseek"
 quantization = "8bit"
 base_model = "DeepSeek V3.1"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "post_last_user"
 
 context_length = 131072
 

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

@@ -8,7 +8,6 @@ family = "deepseek"
 quantization = "4bit"
 base_model = "DeepSeek V3.2"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "tool_conditional"
 
 context_length = 131072
 

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

@@ -8,7 +8,6 @@ family = "deepseek"
 quantization = "8bit"
 base_model = "DeepSeek V3.2"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "tool_conditional"
 
 context_length = 131072
 

resources/inference_model_cards/mlx-community--DeepSeek-V4-Flash.toml

@@ -1,21 +0,0 @@
-model_id = "mlx-community/DeepSeek-V4-Flash"
-n_layers = 43
-hidden_size = 4096
-num_key_value_heads = 1
-supports_tensor = true
-tasks = ["TextGeneration"]
-family = "deepseek"
-quantization = "8bit"
-base_model = "DeepSeek V4 Flash"
-capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "tool_conditional"
-
-context_length = 1048576
-
-[storage_size]
-in_bytes = 155095760030
-
-# Source: https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash
-[sampling_defaults]
-temperature = 1.0
-top_p = 1.0

resources/inference_model_cards/mlx-community--DeepSeek-V4-Pro.toml

@@ -1,21 +0,0 @@
-model_id = "mlx-community/DeepSeek-V4-Pro"
-n_layers = 61
-hidden_size = 7168
-num_key_value_heads = 1
-supports_tensor = true
-tasks = ["TextGeneration"]
-family = "deepseek"
-quantization = "8bit"
-base_model = "DeepSeek V4 Pro"
-capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "tool_conditional"
-
-context_length = 1048576
-
-[storage_size]
-in_bytes = 849681803879
-
-# Source: https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro
-[sampling_defaults]
-temperature = 1.0
-top_p = 1.0

resources/inference_model_cards/mlx-community--GLM-4.5-Air-8bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "8bit"
 base_model = "GLM 4.5 Air"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "post_last_user"
+
 context_length = 131072
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.5-Air-bf16.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "bf16"
 base_model = "GLM 4.5 Air"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "post_last_user"
+
 context_length = 131072
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.7-4bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "4bit"
 base_model = "GLM 4.7"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "post_last_user"
+
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.7-6bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "6bit"
 base_model = "GLM 4.7"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "post_last_user"
+
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.7-8bit-gs32.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "8bit"
 base_model = "GLM 4.7"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "post_last_user"
+
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.7-Flash-4bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "4bit"
 base_model = "GLM 4.7 Flash"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "post_last_user"
+
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.7-Flash-5bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "5bit"
 base_model = "GLM 4.7 Flash"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "post_last_user"
+
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.7-Flash-6bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "6bit"
 base_model = "GLM 4.7 Flash"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "post_last_user"
+
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.7-Flash-8bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "8bit"
 base_model = "GLM 4.7 Flash"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "post_last_user"
+
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-5-8bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "8bit"
 base_model = "GLM-5"
 capabilities = ["text", "thinking"]
-reasoning_dialect = "post_last_user"
+
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-5-MXFP4-Q8.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "MXFP4-Q8"
 base_model = "GLM-5"
 capabilities = ["text", "thinking"]
-reasoning_dialect = "post_last_user"
+
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-5-bf16.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "bf16"
 base_model = "GLM-5"
 capabilities = ["text", "thinking"]
-reasoning_dialect = "post_last_user"
+
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-5.1-DQ4plus-q8.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "8bit"
 base_model = "GLM-5.1"
 capabilities = ["text", "thinking"]
-reasoning_dialect = "post_last_user"
+
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-5.1-MXFP4-Q8.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "MXFP4-Q8"
 base_model = "GLM-5.1"
 capabilities = ["text", "thinking"]
-reasoning_dialect = "post_last_user"
+
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-5.1.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "bf16"
 base_model = "GLM-5.1"
 capabilities = ["text", "thinking"]
-reasoning_dialect = "post_last_user"
+
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--Kimi-K2-Thinking.toml

@@ -8,7 +8,7 @@ family = "kimi"
 quantization = ""
 base_model = "Kimi K2"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "suffix"
+
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Kimi-K2.5.toml

@@ -8,7 +8,7 @@ family = "kimi"
 quantization = ""
 base_model = "Kimi K2.5"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "suffix"
+
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Kimi-K2.6-mlx-DQ3_K_M-q8.toml

@@ -8,7 +8,7 @@ family = "kimi"
 quantization = "3bit"
 base_model = "Kimi K2.6"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "suffix"
+
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.1-3bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "3bit"
 base_model = "MiniMax M2.1"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "post_last_user"
+
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.1-8bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "8bit"
 base_model = "MiniMax M2.1"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "post_last_user"
+
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.5-4bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "4bit"
 base_model = "MiniMax M2.5"
 capabilities = ["text", "thinking"]
-reasoning_dialect = "post_last_user"
+
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.5-6bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "6bit"
 base_model = "MiniMax M2.5"
 capabilities = ["text", "thinking"]
-reasoning_dialect = "post_last_user"
+
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.5-8bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "8bit"
 base_model = "MiniMax M2.5"
 capabilities = ["text", "thinking"]
-reasoning_dialect = "post_last_user"
+
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.7-4bit-mxfp4.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "4bit-mxfp4"
 base_model = "MiniMax M2.7"
 capabilities = ["text", "thinking"]
-reasoning_dialect = "post_last_user"
+
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.7-4bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "4bit"
 base_model = "MiniMax M2.7"
 capabilities = ["text", "thinking"]
-reasoning_dialect = "post_last_user"
+
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.7-5bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "5bit"
 base_model = "MiniMax M2.7"
 capabilities = ["text", "thinking"]
-reasoning_dialect = "post_last_user"
+
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.7-6bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "6bit"
 base_model = "MiniMax M2.7"
 capabilities = ["text", "thinking"]
-reasoning_dialect = "post_last_user"
+
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.7-8bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "8bit"
 base_model = "MiniMax M2.7"
 capabilities = ["text", "thinking"]
-reasoning_dialect = "post_last_user"
+
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.7.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "bf16"
 base_model = "MiniMax M2.7"
 capabilities = ["text", "thinking"]
-reasoning_dialect = "post_last_user"
+
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3-Next-80B-A3B-Thinking-4bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3 Next 80B"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3-Next-80B-A3B-Thinking-8bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3 Next 80B"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

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

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3.5 122B A10B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

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

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "6bit"
 base_model = "Qwen3.5 122B A10B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

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

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3.5 122B A10B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

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

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "bf16"
 base_model = "Qwen3.5 122B A10B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

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

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3.5 27B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

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

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3.5 27B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

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

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3.5 2B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

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

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3.5 35B A3B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

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

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3.5 35B A3B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

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

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3.5 397B A17B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

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

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "6bit"
 base_model = "Qwen3.5 397B A17B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

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

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3.5 397B A17B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

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

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3.5 9B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

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

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3.5 9B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

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

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3.6 27B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

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

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3.6 27B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.6-27B-bf16.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "bf16"
 base_model = "Qwen3.6 27B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

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

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3.6 35B A3B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

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

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "5bit"
 base_model = "Qwen3.6 35B A3B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

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

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3.6 35B A3B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.6-35B-A3B-bf16.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "bf16"
 base_model = "Qwen3.6 35B A3B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "post_last_user"
+
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--gpt-oss-120b-MXFP4-Q8.toml

@@ -8,7 +8,7 @@ family = "gpt-oss"
 quantization = "MXFP4-Q8"
 base_model = "GPT-OSS 120B"
 capabilities = ["text", "thinking"]
-reasoning_dialect = "channel"
+
 context_length = 131072
 
 [storage_size]

resources/inference_model_cards/mlx-community--gpt-oss-20b-MXFP4-Q8.toml

@@ -8,7 +8,7 @@ family = "gpt-oss"
 quantization = "MXFP4-Q8"
 base_model = "GPT-OSS 20B"
 capabilities = ["text", "thinking"]
-reasoning_dialect = "channel"
+
 context_length = 131072
 
 [storage_size]

rust/exo_pyo3_bindings/Cargo.toml

@@ -22,8 +22,7 @@ doc = false
 workspace = true
 
 [dependencies]
-networking.workspace = true
-extend.workspace = true
+networking = { workspace = true }
 
 # interop
 pyo3 = { version = "0.27.2", features = [
@@ -47,16 +46,21 @@ pyo3-async-runtimes = { version = "0.27.0", features = [
 ] }
 pyo3-log = "0.13.2"
 
+# macro dependencies
+extend = { workspace = true }
+delegate = { workspace = true }
+
 # async runtime
 tokio = { workspace = true, features = ["full", "tracing"] }
 futures-lite = { workspace = true }
-pin-project = "1.1.10"
+
+# utility dependencies
+util = { workspace = true }
 
 # Tracing
-log.workspace = true
-env_logger.workspace = true
+log = { workspace = true }
+env_logger = "0.11"
 
 # Networking
-zenoh.workspace = true
-zerompk.workspace = true
-rand = "0.10.1"
+libp2p = { workspace = true, features = ["full"] }
+pin-project = "1.1.10"

rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi

@@ -4,6 +4,12 @@
 import builtins
 import typing
 
+@typing.final
+class AllQueuesFullError(builtins.Exception):
+    def __new__(cls, *args: typing.Any) -> AllQueuesFullError: ...
+    def __repr__(self) -> builtins.str: ...
+    def __str__(self) -> builtins.str: ...
+
 @typing.final
 class Keypair:
     r"""
@@ -28,6 +34,12 @@ class Keypair:
         Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
         """
 
+@typing.final
+class MessageTooLargeError(builtins.Exception):
+    def __new__(cls, *args: typing.Any) -> MessageTooLargeError: ...
+    def __repr__(self) -> builtins.str: ...
+    def __str__(self) -> builtins.str: ...
+
 @typing.final
 class NetworkingHandle:
     def __new__(cls, identity: Keypair, bootstrap_peers: typing.Sequence[builtins.str], listen_port: builtins.int) -> NetworkingHandle: ...
@@ -51,22 +63,32 @@ class NetworkingHandle:
         """
     async def recv(self) -> PyFromSwarm: ...
 
+@typing.final
+class NoPeersSubscribedToTopicError(builtins.Exception):
+    def __new__(cls, *args: typing.Any) -> NoPeersSubscribedToTopicError: ...
+    def __repr__(self) -> builtins.str: ...
+    def __str__(self) -> builtins.str: ...
+
 class PyFromSwarm:
     @typing.final
     class Connection(PyFromSwarm):
-        __match_args__ = ("connected",)
+        __match_args__ = ("peer_id", "connected",)
+        @property
+        def peer_id(self) -> builtins.str: ...
         @property
         def connected(self) -> builtins.bool: ...
-        def __new__(cls, connected: builtins.bool) -> PyFromSwarm.Connection: ...
+        def __new__(cls, peer_id: builtins.str, connected: builtins.bool) -> PyFromSwarm.Connection: ...
     
     @typing.final
     class Message(PyFromSwarm):
-        __match_args__ = ("topic", "data",)
+        __match_args__ = ("origin", "topic", "data",)
+        @property
+        def origin(self) -> builtins.str: ...
         @property
         def topic(self) -> builtins.str: ...
         @property
         def data(self) -> bytes: ...
-        def __new__(cls, topic: builtins.str, data: bytes) -> PyFromSwarm.Message: ...
+        def __new__(cls, origin: builtins.str, topic: builtins.str, data: bytes) -> PyFromSwarm.Message: ...
     
     ...
 

rust/exo_pyo3_bindings/src/ident.rs

@@ -1,4 +1,5 @@
 use crate::ext::ResultExt as _;
+use libp2p::identity::Keypair;
 use pyo3::types::{PyBytes, PyBytesMethods as _};
 use pyo3::{Bound, PyResult, Python, pyclass, pymethods};
 use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
@@ -7,7 +8,7 @@ use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
 #[gen_stub_pyclass]
 #[pyclass(name = "Keypair", frozen)]
 #[repr(transparent)]
-pub struct PyKeypair(pub u128);
+pub struct PyKeypair(pub Keypair);
 
 #[gen_stub_pymethods]
 #[pymethods]
@@ -16,29 +17,31 @@ impl PyKeypair {
     /// Generate a new Ed25519 keypair.
     #[staticmethod]
     fn generate() -> Self {
-        Self(rand::random())
+        Self(Keypair::generate_ed25519())
     }
 
     /// Construct an Ed25519 keypair from secret key bytes
     #[staticmethod]
     fn from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
-        let bytes = Vec::from(bytes.as_bytes());
-        Ok(Self(u128::from_le_bytes(
-            bytes
-                .try_into()
-                .map_err(|_| "passed too many bytes to from_bytes")
-                .pyerr()?,
-        )))
+        let mut bytes = Vec::from(bytes.as_bytes());
+        Ok(Self(Keypair::ed25519_from_bytes(&mut bytes).pyerr()?))
     }
 
     /// Get the secret key bytes underlying the keypair
     fn to_bytes<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
-        let bytes = self.0.to_le_bytes();
+        let bytes = self
+            .0
+            .clone()
+            .try_into_ed25519()
+            .pyerr()?
+            .secret()
+            .as_ref()
+            .to_vec();
         Ok(PyBytes::new(py, &bytes))
     }
 
     /// Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
     fn to_node_id(&self) -> String {
-        format!("{:x}", self.0)
+        self.0.public().to_peer_id().to_base58()
     }
 }

rust/exo_pyo3_bindings/src/networking.rs

@@ -5,8 +5,12 @@ use crate::r#const::MPSC_CHANNEL_SIZE;
 use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
 use crate::ext::{ResultExt as _, TokioMpscSenderExt as _};
 use crate::ident::PyKeypair;
+use crate::networking::exception::{
+    PyAllQueuesFullError, PyMessageTooLargeError, PyNoPeersSubscribedToTopicError,
+};
 use crate::pyclass;
 use futures_lite::{Stream, StreamExt as _};
+use libp2p::gossipsub::PublishError;
 use networking::swarm::{FromSwarm, ToSwarm, create_swarm};
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::{PyModule, PyModuleMethods as _};
@@ -17,6 +21,114 @@ use pyo3_stub_gen::derive::{
 };
 use tokio::sync::{Mutex, mpsc, oneshot};
 
+mod exception {
+    use pyo3::types::PyTuple;
+    use pyo3::{exceptions::PyException, prelude::*};
+    use pyo3_stub_gen::derive::*;
+
+    #[gen_stub_pyclass]
+    #[pyclass(frozen, extends=PyException, name="NoPeersSubscribedToTopicError")]
+    pub struct PyNoPeersSubscribedToTopicError {}
+
+    impl PyNoPeersSubscribedToTopicError {
+        const MSG: &'static str = "\
+        No peers are currently subscribed to receive messages on this topic. \
+        Wait for peers to subscribe or check your network connectivity.";
+
+        ///   Creates a new  [ `PyErr` ]  of this type.
+        ///
+        ///   [`PyErr`] :  https://docs.rs/pyo3/latest/pyo3/struct.PyErr.html   "PyErr in pyo3"
+        pub(crate) fn new_err() -> PyErr {
+            PyErr::new::<Self, _>(()) // TODO: check if this needs to be replaced???
+        }
+    }
+
+    #[gen_stub_pymethods]
+    #[pymethods]
+    impl PyNoPeersSubscribedToTopicError {
+        #[new]
+        #[pyo3(signature = (*args))]
+        #[allow(unused_variables)]
+        pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self {
+            Self {}
+        }
+
+        fn __repr__(&self) -> String {
+            format!("PeerId(\"{}\")", Self::MSG)
+        }
+
+        fn __str__(&self) -> String {
+            Self::MSG.to_string()
+        }
+    }
+
+    #[gen_stub_pyclass]
+    #[pyclass(frozen, extends=PyException, name="AllQueuesFullError")]
+    pub struct PyAllQueuesFullError {}
+
+    impl PyAllQueuesFullError {
+        const MSG: &'static str =
+            "All libp2p peers are unresponsive, resend the message or reconnect.";
+
+        ///   Creates a new  [ `PyErr` ]  of this type.
+        ///
+        ///   [`PyErr`] :  https://docs.rs/pyo3/latest/pyo3/struct.PyErr.html   "PyErr in pyo3"
+        pub(crate) fn new_err() -> PyErr {
+            PyErr::new::<Self, _>(()) // TODO: check if this needs to be replaced???
+        }
+    }
+
+    #[gen_stub_pymethods]
+    #[pymethods]
+    impl PyAllQueuesFullError {
+        #[new]
+        #[pyo3(signature = (*args))]
+        #[allow(unused_variables)]
+        pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self {
+            Self {}
+        }
+
+        fn __repr__(&self) -> String {
+            format!("PeerId(\"{}\")", Self::MSG)
+        }
+
+        fn __str__(&self) -> String {
+            Self::MSG.to_string()
+        }
+    }
+
+    #[gen_stub_pyclass]
+    #[pyclass(frozen, extends=PyException, name="MessageTooLargeError")]
+    pub struct PyMessageTooLargeError {}
+
+    impl PyMessageTooLargeError {
+        const MSG: &'static str = "Gossipsub message exceeds max_transmit_size. Reduce prompt length or increase the limit.";
+
+        pub(crate) fn new_err() -> PyErr {
+            PyErr::new::<Self, _>(())
+        }
+    }
+
+    #[gen_stub_pymethods]
+    #[pymethods]
+    impl PyMessageTooLargeError {
+        #[new]
+        #[pyo3(signature = (*args))]
+        #[allow(unused_variables)]
+        pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self {
+            Self {}
+        }
+
+        fn __repr__(&self) -> String {
+            format!("MessageTooLargeError(\"{}\")", Self::MSG)
+        }
+
+        fn __str__(&self) -> String {
+            Self::MSG.to_string()
+        }
+    }
+}
+
 #[gen_stub_pyclass]
 #[pyclass(name = "NetworkingHandle")]
 struct PyNetworkingHandle {
@@ -28,15 +140,29 @@ struct PyNetworkingHandle {
 #[gen_stub_pyclass_complex_enum]
 #[pyclass]
 enum PyFromSwarm {
-    Connection { connected: bool },
-    Message { topic: String, data: Py<PyBytes> },
+    Connection {
+        peer_id: String,
+        connected: bool,
+    },
+    Message {
+        origin: String,
+        topic: String,
+        data: Py<PyBytes>,
+    },
 }
 impl From<FromSwarm> for PyFromSwarm {
     fn from(value: FromSwarm) -> Self {
         match value {
-            FromSwarm::Discovered {} => Self::Connection { connected: true },
-            FromSwarm::Expired {} => Self::Connection { connected: false },
-            FromSwarm::Message { topic, data } => Self::Message {
+            FromSwarm::Discovered { peer_id } => Self::Connection {
+                peer_id: peer_id.to_base58(),
+                connected: true,
+            },
+            FromSwarm::Expired { peer_id } => Self::Connection {
+                peer_id: peer_id.to_base58(),
+                connected: false,
+            },
+            FromSwarm::Message { from, topic, data } => Self::Message {
+                origin: from.to_base58(),
                 topic: topic,
                 data: data.pybytes(),
             },
@@ -69,8 +195,8 @@ 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)
-            .map(|it| it.into_stream())
-            .pyerr()?;
+            .pyerr()?
+            .into_stream();
 
         Ok(Self {
             swarm: Arc::new(Mutex::new(swarm)),
@@ -159,7 +285,14 @@ impl PyNetworkingHandle {
             .allow_threads_py() // allow-threads-aware async call
             .await
             .map_err(|_| PyErr::receiver_channel_closed())?
-            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+            .map_err(|e| match e {
+                PublishError::AllQueuesFull(_) => PyAllQueuesFullError::new_err(),
+                PublishError::MessageTooLarge => PyMessageTooLargeError::new_err(),
+                PublishError::NoPeersSubscribedToTopic => {
+                    PyNoPeersSubscribedToTopicError::new_err()
+                }
+                e => PyRuntimeError::new_err(e.to_string()),
+            })?;
         Ok(())
     }
 }
@@ -174,6 +307,10 @@ pyo3_stub_gen::inventory::submit! {
 }
 
 pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<exception::PyNoPeersSubscribedToTopicError>()?;
+    m.add_class::<exception::PyAllQueuesFullError>()?;
+    m.add_class::<exception::PyMessageTooLargeError>()?;
+
     m.add_class::<PyNetworkingHandle>()?;
     m.add_class::<PyFromSwarm>()?;
 

rust/exo_pyo3_bindings/tests/dummy.rs

@@ -0,0 +1,54 @@
+#[cfg(test)]
+mod tests {
+    use core::mem::drop;
+    use core::option::Option::Some;
+    use core::time::Duration;
+    use tokio;
+    use tokio::sync::mpsc;
+
+    #[tokio::test]
+    async fn test_drop_channel() {
+        struct Ping;
+
+        let (tx, mut rx) = mpsc::channel::<Ping>(10);
+
+        let _ = tokio::spawn(async move {
+            println!("TASK: entered");
+
+            loop {
+                tokio::select! {
+                    result = rx.recv() => {
+                        match result {
+                            Some(_) => {
+                                println!("TASK: pinged");
+                            }
+                            None => {
+                                println!("TASK: closing channel");
+                                break;
+                            }
+                        }
+                    }
+                    _ = tokio::time::sleep(Duration::from_secs_f32(0.1)) => {
+                        println!("TASK: heartbeat");
+                    }
+                }
+            }
+
+            println!("TASK: exited");
+        });
+
+        let tx2 = tx.clone();
+
+        tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
+
+        tx.send(Ping).await.expect("Should not fail");
+        drop(tx);
+
+        tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
+
+        tx2.send(Ping).await.expect("Should not fail");
+        drop(tx2);
+
+        tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
+    }
+}

rust/exo_pyo3_bindings/tests/test_python.py

@@ -4,6 +4,7 @@ import pytest
 from exo_pyo3_bindings import (
     Keypair,
     NetworkingHandle,
+    NoPeersSubscribedToTopicError,
     PyFromSwarm,
 )
 
@@ -12,15 +13,17 @@ from exo_pyo3_bindings import (
 async def test_sleep_on_multiple_items() -> None:
     print("PYTHON: starting handle")
     h = NetworkingHandle(Keypair.generate(), [], 0)
-    print("PYTHON: handle started")
 
     rt = asyncio.create_task(_await_recv(h))
 
     # sleep for 4 ticks
-    for i in range(10):
+    for i in range(4):
         await asyncio.sleep(1)
 
-        await h.gossipsub_publish("topic", b"somehting or other")
+        try:
+            await h.gossipsub_publish("topic", b"somehting or other")
+        except NoPeersSubscribedToTopicError as e:
+            print("caught it", e)
 
 
 async def _await_recv(h: NetworkingHandle):
@@ -31,6 +34,3 @@ async def _await_recv(h: NetworkingHandle):
                 print(f"PYTHON: connection update: {c}")
             case PyFromSwarm.Message() as m:
                 print(f"PYTHON: message: {m}")
-
-if __name__ == "__main__":
-    asyncio.run(test_sleep_on_multiple_items())

rust/networking/Cargo.toml

@@ -1,22 +1,42 @@
 [package]
 name = "networking"
-version.workspace = true
-edition.workspace = true
+version = { workspace = true }
+edition = { workspace = true }
+publish = false
 
-[dependencies]
-async-stream = "0.3.6"
-futures-lite.workspace = true
-netwatcher = { version = "0.6.0", features = ["tokio"] }
-parking_lot = "0.12.5"
-tokio = { workspace = true, features = ["full"] }
-zenoh = { version = "=1.9.0", features = ["internal", "plugins", "unstable"] }
-zenoh-plugin-storage-manager = { version = "=1.9.0", default-features = false }
-zenoh-plugin-trait = "=1.9.0"
-zerompk = { version = "0.4.2", features = ["derive"] }
-rand = "0.10.1"
-tracing = "0.1.44"
-log.workspace = true
+[lib]
+doctest = false
+name = "networking"
+path = "src/lib.rs"
 
 [lints]
 workspace = true
 
+[dependencies]
+# datastructures
+either = { workspace = true }
+
+# macro dependencies
+extend = { workspace = true }
+delegate = { workspace = true }
+
+# async
+async-stream = { workspace = true }
+futures-lite = { workspace = true }
+futures-timer = { workspace = true }
+tokio = { workspace = true, features = ["full"] }
+
+# utility dependencies
+util = { workspace = true }
+tracing-subscriber = { version = "0.3.19", features = [
+  "default",
+  "env-filter",
+] }
+keccak-const = { workspace = true }
+
+# tracing/logging
+log = { workspace = true }
+
+# networking
+libp2p = { workspace = true, features = ["full"] }
+pin-project = "1.1.10"

rust/networking/examples/chatroom.rs

@@ -0,0 +1,86 @@
+use futures_lite::StreamExt;
+use libp2p::identity;
+use networking::swarm;
+use networking::swarm::{FromSwarm, ToSwarm};
+use tokio::sync::{mpsc, oneshot};
+use tokio::{io, io::AsyncBufReadExt as _};
+use tracing_subscriber::EnvFilter;
+use tracing_subscriber::filter::LevelFilter;
+
+#[tokio::main]
+async fn main() {
+    let _ = tracing_subscriber::fmt()
+        .with_env_filter(EnvFilter::from_default_env().add_directive(LevelFilter::INFO.into()))
+        .try_init();
+
+    let (to_swarm, from_client) = mpsc::channel(20);
+
+    // Configure swarm
+    let mut swarm = swarm::create_swarm(
+        identity::Keypair::generate_ed25519(),
+        from_client,
+        vec![],
+        0,
+    )
+    .expect("Swarm creation failed")
+    .into_stream();
+
+    // Create a Gossipsub topic & subscribe
+    let (tx, rx) = oneshot::channel();
+    _ = to_swarm
+        .send(ToSwarm::Subscribe {
+            topic: "test-net".to_string(),
+            result_sender: tx,
+        })
+        .await
+        .expect("should send");
+
+    // Read full lines from stdin
+    let mut stdin = io::BufReader::new(io::stdin()).lines();
+    println!("Enter messages via STDIN and they will be sent to connected peers using Gossipsub");
+
+    tokio::task::spawn(async move {
+        rx.await
+            .expect("tx not dropped")
+            .expect("subscribe shouldn't fail");
+        loop {
+            if let Ok(Some(line)) = stdin.next_line().await {
+                let (tx, rx) = oneshot::channel();
+                if let Err(e) = to_swarm
+                    .send(swarm::ToSwarm::Publish {
+                        topic: "test-net".to_string(),
+                        data: line.as_bytes().to_vec(),
+                        result_sender: tx,
+                    })
+                    .await
+                {
+                    println!("Send error: {e:?}");
+                    return;
+                };
+                match rx.await {
+                    Ok(Err(e)) => println!("Publish error: {e:?}"),
+                    Err(e) => println!("Publish error: {e:?}"),
+                    Ok(_) => {}
+                }
+            }
+        }
+    });
+
+    // Kick it off
+    loop {
+        // on gossipsub outgoing
+        match swarm.next().await {
+            // on gossipsub incoming
+            Some(FromSwarm::Discovered { peer_id }) => {
+                println!("\n\nconnected to {peer_id}\n\n")
+            }
+            Some(FromSwarm::Expired { peer_id }) => {
+                println!("\n\ndisconnected from {peer_id}\n\n")
+            }
+            Some(FromSwarm::Message { from, topic, data }) => {
+                println!("{topic}/{from}:\n{}", String::from_utf8_lossy(&data))
+            }
+            None => {}
+        }
+    }
+}

rust/networking/examples/put_string.rs

@@ -1,22 +0,0 @@
-use networking;
-use tracing::info;
-use zenoh::Result;
-
-#[tokio::main]
-async fn main() -> Result<()> {
-    zenoh::init_log_from_env_or("info");
-    info!("Opening session...");
-    let cfg = networking::cfg(rand::random(), 0)?;
-    let session = networking::open(cfg).await?;
-    let _tok = session
-        .liveliness()
-        .declare_token(format!("nodes/{}/live", session.zid()))
-        .await?;
-    let key_expr = "storage/mem1/name";
-    let payload = "me";
-
-    info!("Putting Data ('{key_expr}': '{payload}')...");
-    session.put(key_expr, payload).await?;
-    tokio::signal::ctrl_c().await?;
-    Ok(())
-}

rust/networking/examples/serve_storage.rs

@@ -1,23 +0,0 @@
-use networking;
-use tracing::info;
-use zenoh::Result;
-
-#[tokio::main]
-async fn main() -> Result<()> {
-    zenoh::init_log_from_env_or("info");
-    info!("Opening session...");
-    let cfg = networking::cfg(rand::random(), 0)?;
-    let session = networking::open(cfg).await?;
-    let _tok = session
-        .liveliness()
-        .declare_token(format!("nodes/{}/live", session.zid()))
-        .await?;
-    let _sub = session
-        .liveliness()
-        .declare_subscriber("nodes/*/live")
-        .history(true)
-        .callback(|tok| println!("{tok:?}"))
-        .await?;
-    tokio::signal::ctrl_c().await?;
-    Ok(())
-}

rust/networking/src/RESEARCH_NOTES.txt

@@ -0,0 +1,44 @@
+https://github.com/ml-explore/mlx/commit/3fe98bacc7640d857acf3539f1d21b47a32e5609
+^raw sockets distributed -> `<net/ndrv.h>` -> https://newosxbook.com/code/xnu-3247.1.106/bsd/net/ndrv.h.auto.html
+--> header file for a networking component found in the macOS kernel (XNU) that defines structures for network device driver registration, specifically the ndrv_demux_desc and ndrv_protocol_desc structures used for demultiplexing protocol data at the network interface level. It specifies how to describe protocol data, such as an Ethernet type or a SNAP header, and how to associate these descriptions with a specific protocol family to receive matching packets.
+--> Used to bind an NDRV socket so that packets that match given protocol demux descriptions can be received.
+--> An NDRV socket is a special kind of socket in the Darwin/macOS operating system's XNU kernel, used for low-level network packet manipulation and binding to specific protocols for packet processing. It allows user-space applications or drivers to directly write Layer 2 (L2) network packets or interact with the network stack at a lower level, often by binding to protocol descriptors like the ndrv_protocol_desc. This type of socket is used for functions such as capturing and injecting packets, especially in network infrastructure software like routers or for kernel-level network monitoring and security tools.
+--> also called PF_NDRV sockets --> https://newosxbook.com/bonus/vol1ch16.html
+----> they are conceptually similar to https://scapy.disruptivelabs.in/networking/socket-interface PF_RAW or PF_PACKET
+
+https://stackoverflow.com/questions/17169298/af-packet-on-osx
+^AF_PACKET duplicates the packets as soon as it receives them from the physical layer (for incoming packets) or just before sending them out to the physical layer (for outgoing packets). -> this is on Linux only
+^it doesn't exist on OS X so you can use /dev/bpfX (Berkeley Packet Filter) for sniffing
+
+https://www.unix.com/man_page/mojave/4/ip/
+^OS X manpages for IP
+
+https://developer.apple.com/documentation/kernel/implementing_drivers_system_extensions_and_kexts
+^driver kit, system extensions & kexts for macOS
+
+----
+
+To set up a Linux system to use a Thunderbolt connection as a network device, connect the two computers with a Thunderbolt cable, load the thunderbolt-net kernel module (usually automatic but modprobe is an option for manual loading), and then the operating system will create virtual Ethernet interfaces (e.g., thunderbolt0) for networking. You can then use standard tools like ifconfig or your desktop environment's network manager to configure these new interfaces for a link-local network.
+--> https://gist.github.com/geosp/80fbd39e617b7d1d9421683df4ea224a
+----> here is a guide on how to set up thunderbolt-ethernet on linux
+----> I may be able to steal the thunderbolt-net code ideas to implement a kernel module for MacOS
+
+https://chatgpt.com/s/t_68af8e41a8548191993281a014f846a7
+^GPT discussion about making socket interface
+
+https://chatgpt.com/s/t_68afb798a85c8191973c02a0fa7a48a3 --> link-local address,,??
+https://chatgpt.com/s/t_68afb02987e08191b2b0044d3667ece2
+^GPT discussion about accessing TB on MacOS low level interactions
+
+--------------------------------
+
+https://www.intel.com/content/www/us/en/support/articles/000098893/software.html
+^Thunderbolt Share & Thunderbolt Networking Mode => intel's equivalent of thunderbolt bridge
+
+
+---------------------------------
+
+https://www.zerotier.com/blog/how-zerotier-eliminated-kernel-extensions-on-macos/
+-->fake ethernet devices on MacOS -> omg??? we can detect thunderbolt bridge, then bind to it, then re-expose it as fake ethernet??
+-->ps: https://chatgpt.com/s/t_68afb2b25fb881919526763fb5d7359c, AF/PF_NDRV are one and the same!!!
+-->https://github.com/zerotier/ZeroTierOne/blob/dev/osdep/MacEthernetTapAgent.c

rust/networking/src/discovery.rs

@@ -0,0 +1,390 @@
+use crate::ext::MultiaddrExt;
+use delegate::delegate;
+use either::Either;
+use futures_lite::FutureExt;
+use futures_timer::Delay;
+use libp2p::core::transport::PortUse;
+use libp2p::core::{ConnectedPoint, Endpoint};
+use libp2p::swarm::behaviour::ConnectionEstablished;
+use libp2p::swarm::dial_opts::DialOpts;
+use libp2p::swarm::{
+    CloseConnection, ConnectionClosed, ConnectionDenied, ConnectionHandler,
+    ConnectionHandlerSelect, ConnectionId, FromSwarm, NetworkBehaviour, THandler, THandlerInEvent,
+    THandlerOutEvent, ToSwarm, dummy,
+};
+use libp2p::{Multiaddr, PeerId, identity, mdns};
+use std::collections::{BTreeSet, HashMap};
+use std::convert::Infallible;
+use std::io;
+use std::net::IpAddr;
+use std::task::{Context, Poll};
+use std::time::Duration;
+use util::wakerdeque::WakerDeque;
+
+const RETRY_CONNECT_INTERVAL: Duration = Duration::from_secs(5);
+
+mod managed {
+    use libp2p::swarm::NetworkBehaviour;
+    use libp2p::{identity, mdns, ping};
+    use std::io;
+    use std::time::Duration;
+
+    const MDNS_RECORD_TTL: Duration = Duration::from_secs(2_500);
+    const MDNS_QUERY_INTERVAL: Duration = Duration::from_secs(1_500);
+    const PING_TIMEOUT: Duration = Duration::from_millis(2_500);
+    const PING_INTERVAL: Duration = Duration::from_millis(2_500);
+
+    #[derive(NetworkBehaviour)]
+    pub struct Behaviour {
+        mdns: mdns::tokio::Behaviour,
+        ping: ping::Behaviour,
+    }
+
+    impl Behaviour {
+        pub fn new(keypair: &identity::Keypair) -> io::Result<Self> {
+            Ok(Self {
+                mdns: mdns_behaviour(keypair)?,
+                ping: ping_behaviour(),
+            })
+        }
+    }
+
+    fn mdns_behaviour(keypair: &identity::Keypair) -> io::Result<mdns::tokio::Behaviour> {
+        use mdns::{Config, tokio};
+
+        // mDNS config => enable IPv6
+        let mdns_config = Config {
+            ttl: MDNS_RECORD_TTL,
+            query_interval: MDNS_QUERY_INTERVAL,
+
+            // enable_ipv6: true, // TODO: for some reason, TCP+mDNS don't work well with ipv6?? figure out how to make work
+            ..Default::default()
+        };
+
+        let mdns_behaviour = tokio::Behaviour::new(mdns_config, keypair.public().to_peer_id());
+        Ok(mdns_behaviour?)
+    }
+
+    fn ping_behaviour() -> ping::Behaviour {
+        ping::Behaviour::new(
+            ping::Config::new()
+                .with_timeout(PING_TIMEOUT)
+                .with_interval(PING_INTERVAL),
+        )
+    }
+}
+
+/// Events for when a listening connection is truly established and truly closed.
+#[derive(Debug, Clone)]
+pub enum Event {
+    ConnectionEstablished {
+        peer_id: PeerId,
+        connection_id: ConnectionId,
+        remote_ip: IpAddr,
+        remote_tcp_port: u16,
+    },
+    ConnectionClosed {
+        peer_id: PeerId,
+        connection_id: ConnectionId,
+        remote_ip: IpAddr,
+        remote_tcp_port: u16,
+    },
+}
+
+/// Discovery behavior that wraps mDNS to produce truly discovered durable peer-connections.
+///
+/// The behaviour operates as such:
+///  1) All true (listening) connections/disconnections are tracked, emitting corresponding events
+///     to the swarm.
+///  1) mDNS discovered/expired peers are tracked; discovered but not connected peers are dialed
+///     immediately, and expired but connected peers are disconnected from immediately.
+///  2) Every fixed interval: discovered but not connected peers are dialed, and expired but
+///     connected peers are disconnected from.
+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
+
+    // pending events to emmit => waker-backed Deque to control polling
+    pending_events: WakerDeque<ToSwarm<Event, Infallible>>,
+}
+
+impl Behaviour {
+    pub fn new(keypair: &identity::Keypair, bootstrap_peers: Vec<Multiaddr>) -> 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(),
+        })
+    }
+
+    fn dial(&mut self, peer_id: PeerId, addr: Multiaddr) {
+        self.pending_events.push_back(ToSwarm::Dial {
+            opts: DialOpts::peer_id(peer_id).addresses(vec![addr]).build(),
+        })
+    }
+
+    fn close_connection(&mut self, peer_id: PeerId, connection: ConnectionId) {
+        // push front to make this IMMEDIATE
+        self.pending_events.push_front(ToSwarm::CloseConnection {
+            peer_id,
+            connection: CloseConnection::One(connection),
+        })
+    }
+
+    fn handle_mdns_discovered(&mut self, peers: Vec<(PeerId, Multiaddr)>) {
+        for (p, ma) in peers {
+            self.dial(p, ma.clone()); // always connect
+
+            // get peer's multi-addresses or insert if missing
+            let Some(mas) = self.mdns_discovered.get_mut(&p) else {
+                self.mdns_discovered.insert(p, BTreeSet::from([ma]));
+                continue;
+            };
+
+            // multiaddress should never already be present - else something has gone wrong
+            let is_new_addr = mas.insert(ma);
+            assert!(is_new_addr, "cannot discover a discovered peer");
+        }
+    }
+
+    fn handle_mdns_expired(&mut self, peers: Vec<(PeerId, Multiaddr)>) {
+        for (p, ma) in peers {
+            // at this point, we *must* have the peer
+            let mas = self
+                .mdns_discovered
+                .get_mut(&p)
+                .expect("nonexistent peer cannot expire");
+
+            // at this point, we *must* have the multiaddress
+            let was_present = mas.remove(&ma);
+            assert!(was_present, "nonexistent multiaddress cannot expire");
+
+            // if empty, remove the peer-id entirely
+            if mas.is_empty() {
+                self.mdns_discovered.remove(&p);
+            }
+        }
+    }
+
+    fn on_connection_established(
+        &mut self,
+        peer_id: PeerId,
+        connection_id: ConnectionId,
+        remote_ip: IpAddr,
+        remote_tcp_port: u16,
+    ) {
+        // send out connected event
+        self.pending_events
+            .push_back(ToSwarm::GenerateEvent(Event::ConnectionEstablished {
+                peer_id,
+                connection_id,
+                remote_ip,
+                remote_tcp_port,
+            }));
+    }
+
+    fn on_connection_closed(
+        &mut self,
+        peer_id: PeerId,
+        connection_id: ConnectionId,
+        remote_ip: IpAddr,
+        remote_tcp_port: u16,
+    ) {
+        // send out disconnected event
+        self.pending_events
+            .push_back(ToSwarm::GenerateEvent(Event::ConnectionClosed {
+                peer_id,
+                connection_id,
+                remote_ip,
+                remote_tcp_port,
+            }));
+    }
+}
+
+impl NetworkBehaviour for Behaviour {
+    type ConnectionHandler =
+        ConnectionHandlerSelect<dummy::ConnectionHandler, THandler<managed::Behaviour>>;
+    type ToSwarm = Event;
+
+    // simply delegate to underlying mDNS behaviour
+
+    delegate! {
+        to self.managed {
+            fn handle_pending_inbound_connection(&mut self, connection_id: ConnectionId, local_addr: &Multiaddr, remote_addr: &Multiaddr) -> Result<(), ConnectionDenied>;
+            fn handle_pending_outbound_connection(&mut self, connection_id: ConnectionId, maybe_peer: Option<PeerId>, addresses: &[Multiaddr], effective_role: Endpoint) -> Result<Vec<Multiaddr>, ConnectionDenied>;
+        }
+    }
+
+    fn handle_established_inbound_connection(
+        &mut self,
+        connection_id: ConnectionId,
+        peer: PeerId,
+        local_addr: &Multiaddr,
+        remote_addr: &Multiaddr,
+    ) -> Result<THandler<Self>, ConnectionDenied> {
+        Ok(ConnectionHandler::select(
+            dummy::ConnectionHandler,
+            self.managed.handle_established_inbound_connection(
+                connection_id,
+                peer,
+                local_addr,
+                remote_addr,
+            )?,
+        ))
+    }
+
+    #[allow(clippy::needless_question_mark)]
+    fn handle_established_outbound_connection(
+        &mut self,
+        connection_id: ConnectionId,
+        peer: PeerId,
+        addr: &Multiaddr,
+        role_override: Endpoint,
+        port_use: PortUse,
+    ) -> Result<THandler<Self>, ConnectionDenied> {
+        Ok(ConnectionHandler::select(
+            dummy::ConnectionHandler,
+            self.managed.handle_established_outbound_connection(
+                connection_id,
+                peer,
+                addr,
+                role_override,
+                port_use,
+            )?,
+        ))
+    }
+
+    fn on_connection_handler_event(
+        &mut self,
+        peer_id: PeerId,
+        connection_id: ConnectionId,
+        event: THandlerOutEvent<Self>,
+    ) {
+        match event {
+            Either::Left(ev) => libp2p::core::util::unreachable(ev),
+            Either::Right(ev) => {
+                self.managed
+                    .on_connection_handler_event(peer_id, connection_id, ev)
+            }
+        }
+    }
+
+    // hook into these methods to drive behavior
+
+    fn on_swarm_event(&mut self, event: FromSwarm) {
+        self.managed.on_swarm_event(event); // let mDNS handle swarm events
+
+        // handle swarm events to update internal state:
+        match event {
+            FromSwarm::ConnectionEstablished(ConnectionEstablished {
+                peer_id,
+                connection_id,
+                endpoint,
+                ..
+            }) => {
+                let remote_address = match endpoint {
+                    ConnectedPoint::Dialer { address, .. } => address,
+                    ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr,
+                };
+
+                if let Some((ip, port)) = remote_address.try_to_tcp_addr() {
+                    // handle connection established event which is filtered correctly
+                    self.on_connection_established(peer_id, connection_id, ip, port)
+                }
+            }
+            FromSwarm::ConnectionClosed(ConnectionClosed {
+                peer_id,
+                connection_id,
+                endpoint,
+                ..
+            }) => {
+                let remote_address = match endpoint {
+                    ConnectedPoint::Dialer { address, .. } => address,
+                    ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr,
+                };
+
+                if let Some((ip, port)) = remote_address.try_to_tcp_addr() {
+                    // handle connection closed event which is filtered correctly
+                    self.on_connection_closed(peer_id, connection_id, ip, port)
+                }
+            }
+
+            // since we are running TCP/IP transport layer, we are assuming that
+            // no address changes can occur, hence encountering one is a fatal error
+            FromSwarm::AddressChange(a) => {
+                unreachable!("unhandlable: address change encountered: {:?}", a)
+            }
+            _ => {}
+        }
+    }
+
+    fn poll(&mut self, cx: &mut Context) -> Poll<ToSwarm<Self::ToSwarm, THandlerInEvent<Self>>> {
+        // delegate to managed behaviors for any behaviors they need to perform
+        match self.managed.poll(cx) {
+            Poll::Ready(ToSwarm::GenerateEvent(e)) => {
+                match e {
+                    // handle discovered and expired events from mDNS
+                    managed::BehaviourEvent::Mdns(e) => match e.clone() {
+                        mdns::Event::Discovered(peers) => {
+                            self.handle_mdns_discovered(peers);
+                        }
+                        mdns::Event::Expired(peers) => {
+                            self.handle_mdns_expired(peers);
+                        }
+                    },
+
+                    // handle ping events => if error then disconnect
+                    managed::BehaviourEvent::Ping(e) => {
+                        if let Err(_) = e.result {
+                            self.close_connection(e.peer, e.connection.clone())
+                        }
+                    }
+                }
+
+                // since we just consumed an event, we should immediately wake just in case
+                // there are more events to come where that came from
+                cx.waker().wake_by_ref();
+            }
+
+            // forward any other mDNS event to the swarm or its connection handler(s)
+            Poll::Ready(e) => {
+                return Poll::Ready(
+                    e.map_out(|_| unreachable!("events returning to swarm already handled"))
+                        .map_in(Either::Right),
+                );
+            }
+
+            Poll::Pending => {}
+        }
+
+        // retry connecting to all mDNS peers periodically (fails safely if already connected)
+        if self.retry_delay.poll(cx).is_ready() {
+            for (p, mas) in self.mdns_discovered.clone() {
+                for ma in mas {
+                    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
+        }
+
+        // send out any pending events from our own service
+        if let Some(e) = self.pending_events.pop_front(cx) {
+            return Poll::Ready(e.map_in(Either::Left));
+        }
+
+        // wait for pending events
+        Poll::Pending
+    }
+}

rust/networking/src/lib.rs

@@ -1,116 +1,44 @@
-use std::{
-    env,
-    ops::{Deref, DerefMut},
-    panic,
-};
-
-use netwatcher::WatchHandle;
-use tokio::{sync::mpsc, task::JoinHandle};
-use zenoh::{Result, Session as ZSession, config::WhatAmI, internal::runtime::Runtime};
-use zenoh_plugin_storage_manager::StoragesPlugin;
-use zenoh_plugin_trait::PluginsManager;
-
-pub use zenoh::{Config, config::ZenohId};
-
+//! TODO: crate documentation
+//!
+//! this is here as a placeholder documentation
+//!
+//!
+pub mod discovery;
 pub mod swarm;
 
-pub fn cfg(identity: u128, listen_port: u16) -> Result<zenoh::Config> {
-    let namespace = env::var("EXO_ZENOH_NAMESPACE").unwrap_or_else(|_| "exo".to_string());
-    let mut cfg = zenoh::Config::default();
-    // todo: cleanup
-    cfg.insert_json5("id", &format!("\"{identity:x}\""))?;
-    cfg.insert_json5("mode", "\"peer\"")?;
-    cfg.insert_json5("listen/endpoints", &format!("[\"tcp/[::]:{listen_port}\"]"))?;
-    cfg.insert_json5("scouting/multicast/enabled", "true")?;
-    cfg.insert_json5("scouting/multicast/autoconnect", "[]")?;
-    cfg.insert_json5("scouting/gossip/multihop", "true")?;
-    cfg.insert_json5("namespace", &format!("{namespace:?}"))?;
-    cfg.insert_json5("transport/link/tx/batch_size", "9216")?;
-    cfg.insert_json5("timestamping/enabled", "true")?;
-    cfg.insert_json5("plugins/storage_manager/__required__", "true")?;
-    cfg.insert_json5(
-        "plugins/storage_manager/storages/mem1",
-        r#"{
-            key_expr: "storage/mem1/**",
-            strip_prefix: "storage/mem1",
-            volume: "memory",
-            replication: {
-                interval: 2,
-            }
-        }"#,
-    )?;
-    Ok(cfg)
+/// Namespace for all the type/trait aliases used by this crate.
+pub(crate) mod alias {
+    use std::error::Error;
+
+    pub type AnyError = Box<dyn Error + Send + Sync + 'static>;
+    pub type AnyResult<T> = Result<T, AnyError>;
 }
 
-pub async fn open(cfg: zenoh::Config) -> Result<Session> {
-    let mut plugins = PluginsManager::static_plugins_only();
-    plugins.declare_static_plugin::<StoragesPlugin, _>("storage_manager", true);
-    let mut runtime = zenoh::internal::runtime::RuntimeBuilder::new(cfg)
-        .plugins_manager(plugins)
-        .build()
-        .await?;
-    let session = zenoh::session::init(runtime.clone().into()).await?;
-    runtime.start().await?;
-    let _watch_all_handle = watch_all(runtime).await?;
-    Ok(Session {
-        session,
-        _watch_all_handle,
-    })
-}
-async fn watch_all(runtime: Runtime) -> Result<WatchAllHandle> {
-    log::info!("spawning scout");
-    let mut cfg = Config::default();
-    cfg.insert_json5("scouting/multicast/ttl", "3")?;
-    cfg.insert_json5("scouting/multicast/interface", "\"auto\"")?;
-    let mut scout = zenoh::scout(WhatAmI::Peer, cfg.clone()).await?;
-    let (send, mut recv) = mpsc::unbounded_channel();
-    let _sync = netwatcher::watch_interfaces_with_callback(move |u| _ = send.send(u))?;
-    let _async = tokio::task::spawn(async move {
-        loop {
-            tokio::select! {
-                u = recv.recv() => {
-                    if u.is_none() {
-                        return Ok(());
-                    }
-                    log::info!("reloading scout");
-                    scout = zenoh::scout(WhatAmI::Peer, cfg.clone()).await?;
+/// Namespace for crate-wide extension traits/methods
+pub(crate) mod ext {
+    use extend::ext;
+    use libp2p::Multiaddr;
+    use libp2p::multiaddr::Protocol;
+    use std::net::IpAddr;
+
+    #[ext(pub, name = MultiaddrExt)]
+    impl Multiaddr {
+        /// If the multiaddress corresponds to a TCP address, extracts it
+        fn try_to_tcp_addr(&self) -> Option<(IpAddr, u16)> {
+            let mut ps = self.into_iter();
+            let ip = if let Some(p) = ps.next() {
+                match p {
+                    Protocol::Ip4(ip) => IpAddr::V4(ip),
+                    Protocol::Ip6(ip) => IpAddr::V6(ip),
+                    _ => return None,
                 }
-                hello = scout.recv_async() => {
-                    if let Ok(hello) = hello {
-                        // todo: auth
-                        runtime
-                            .connect_peer(&hello.zid().into(), hello.locators())
-                            .await;
-                    }
-                }
-            }
+            } else {
+                return None;
+            };
+            let Some(Protocol::Tcp(port)) = ps.next() else {
+                return None;
+            };
+            Some((ip, port))
         }
-    });
-    Ok(WatchAllHandle { _sync, _async })
-}
-
-pub struct Session {
-    pub session: ZSession,
-    _watch_all_handle: WatchAllHandle,
-}
-impl Deref for Session {
-    type Target = ZSession;
-    fn deref(&self) -> &Self::Target {
-        &self.session
     }
 }
-impl DerefMut for Session {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.session
-    }
-}
-impl Drop for WatchAllHandle {
-    fn drop(&mut self) {
-        self._async.abort();
-    }
-}
-
-pub struct WatchAllHandle {
-    _sync: WatchHandle,
-    _async: JoinHandle<Result<()>>,
-}

rust/networking/src/swarm.rs

@@ -1,22 +1,24 @@
-//! Compat shim for the old libp2p code
-
-use std::collections::HashMap;
 use std::pin::Pin;
 
-use futures_lite::Stream;
-use tokio::sync::mpsc;
-use tokio::sync::oneshot;
-use tracing::info;
-use zenoh::Result;
-use zenoh::Session;
-use zenoh::handlers::FifoChannelHandler;
-use zenoh::liveliness::LivelinessToken;
-use zenoh::pubsub::Subscriber;
-use zenoh::sample::Sample;
-use zenoh::sample::SampleKind;
-use zerompk::{FromMessagePack, ToMessagePack};
+use crate::swarm::transport::tcp_transport;
+use crate::{alias, discovery};
+pub use behaviour::{Behaviour, BehaviourEvent};
+use futures_lite::{Stream, StreamExt};
+use libp2p::{PeerId, SwarmBuilder, gossipsub, identity, swarm::SwarmEvent};
+use tokio::sync::{mpsc, oneshot};
 
-#[derive(Debug)]
+/// The current version of the network: this prevents devices running different versions of the
+/// software from interacting with each other.
+///
+/// TODO: right now this is a hardcoded constant; figure out what the versioning semantics should
+///       even be, and how to inject the right version into this config/initialization. E.g. should
+///       this be passed in as a parameter? What about rapidly changing versions in debug builds?
+///       this is all VERY very hard to figure out and needs to be mulled over as a team.
+pub const NETWORK_VERSION: &[u8] = b"v0.0.1";
+pub const OVERRIDE_VERSION_ENV_VAR: &str = "EXO_LIBP2P_NAMESPACE";
+
+// Uses oneshot senders to emulate function calling apis while avoiding requiring unique ownership
+// of the Swarm.
 pub enum ToSwarm {
     Unsubscribe {
         topic: String,
@@ -24,66 +26,52 @@ pub enum ToSwarm {
     },
     Subscribe {
         topic: String,
-        result_sender: oneshot::Sender<Result<bool>>,
+        result_sender: oneshot::Sender<Result<bool, gossipsub::SubscriptionError>>,
     },
     Publish {
         topic: String,
         data: Vec<u8>,
-        result_sender: oneshot::Sender<Result<()>>,
+        result_sender: oneshot::Sender<Result<gossipsub::MessageId, gossipsub::PublishError>>,
     },
 }
-#[derive(Debug, ToMessagePack, FromMessagePack)]
 pub enum FromSwarm {
-    Message { topic: String, data: Vec<u8> },
-    Discovered {},
-    Expired {},
+    Message {
+        from: PeerId,
+        topic: String,
+        data: Vec<u8>,
+    },
+    Discovered {
+        peer_id: PeerId,
+    },
+    Expired {
+        peer_id: PeerId,
+    },
 }
 
-pub type Topics = HashMap<String, Subscriber<()>>;
 pub struct Swarm {
-    cfg: zenoh::Config,
+    swarm: libp2p::Swarm<Behaviour>,
     from_client: mpsc::Receiver<ToSwarm>,
 }
 
 impl Swarm {
     pub fn into_stream(self) -> Pin<Box<dyn Stream<Item = FromSwarm> + Send>> {
         let Swarm {
-            cfg,
+            mut swarm,
             mut from_client,
         } = self;
         let stream = async_stream::stream! {
-            let (mut to_topics, mut from_topics) = mpsc::channel(1024);
-            let mut topics = Topics::new();
-            let Ok(mut session) = crate::open(cfg).await else { return; };
-            let Ok((_token, discovery)) = register_liveness(&mut session).await else { return; };
             loop {
                 tokio::select! {
                     msg = from_client.recv() => {
                         let Some(msg) = msg else { break };
-                        on_message(&mut session, &mut topics, &mut to_topics, msg).await;
+                        on_message(&mut swarm, msg);
                     }
-                    event = from_topics.recv() => {
-                        if let Some(event) = event {
-                            yield event
+                    event = swarm.next() => {
+                        let Some(event) = event else { break };
+                        if let Some(item) = filter_swarm_event(event) {
+                            yield item;
                         }
                     }
-                    token = discovery.recv_async() => {
-                        if let Ok(token) = token {
-                            let key_expr = token.key_expr().as_str().to_owned();
-                            let nid = key_expr.strip_prefix("nodes/").and_then(|s| s.strip_suffix("/live"));
-                            yield match token.kind() {
-                                SampleKind::Put => {
-                                    info!("discovered: {nid:?}");
-                                    FromSwarm::Discovered {}
-                                }
-                                SampleKind::Delete => {
-                                    info!("expired: {nid:?}");
-                                    FromSwarm::Expired {}
-                                }
-                            }
-                        }
-
-                    }
                 }
             }
         };
@@ -91,96 +79,208 @@ impl Swarm {
     }
 }
 
-async fn register_liveness(
-    session: &mut Session,
-) -> Result<(LivelinessToken, Subscriber<FifoChannelHandler<Sample>>)> {
-    let token = session
-        .liveliness()
-        .declare_token(format!("nodes/{}/live", session.zid()))
-        .await?;
-    let sub = session
-        .liveliness()
-        .declare_subscriber("nodes/*/live")
-        .history(true)
-        .await?;
-    Ok((token, sub))
-}
-
-async fn on_message(
-    session: &mut Session,
-    topics: &mut Topics,
-    to_topics: &mut mpsc::Sender<FromSwarm>,
-    msg: ToSwarm,
-) {
-    match msg {
-        ToSwarm::Publish {
+fn on_message(swarm: &mut libp2p::Swarm<Behaviour>, message: ToSwarm) {
+    match message {
+        ToSwarm::Subscribe {
             topic,
-            data,
             result_sender,
         } => {
-            let res = session.put(format!("topics/{topic}"), data).await;
-            _ = result_sender.send(res);
+            let result = swarm
+                .behaviour_mut()
+                .gossipsub
+                .subscribe(&gossipsub::IdentTopic::new(topic));
+            _ = result_sender.send(result);
         }
         ToSwarm::Unsubscribe {
             topic,
             result_sender,
         } => {
-            let Some((_, subscriber)) = topics.remove_entry(&topic) else {
-                _ = result_sender.send(false);
-                return;
-            };
-            _ = subscriber.undeclare().await;
-            _ = result_sender.send(true);
+            let result = swarm
+                .behaviour_mut()
+                .gossipsub
+                .unsubscribe(&gossipsub::IdentTopic::new(topic));
+            _ = result_sender.send(result);
         }
-        ToSwarm::Subscribe {
+        ToSwarm::Publish {
             topic,
+            data,
             result_sender,
         } => {
-            assert!(topic.is_ascii());
-            if topics.contains_key(&topic) {
-                _ = result_sender.send(Ok(false));
-                return;
-            }
-            let subscriber = match session
-                .declare_subscriber(format!("topics/{topic}"))
-                .allowed_origin(zenoh::sample::Locality::Remote)
-                .callback({
-                    let sender = to_topics.clone();
-                    let topic = topic.clone();
-                    move |sample| {
-                        if sample.kind() != SampleKind::Put {
-                            return;
-                        }
-                        _ = sender.try_send(FromSwarm::Message {
-                            topic: topic.clone(),
-                            data: sample.payload().to_bytes().to_vec(),
-                        });
-                    }
-                })
-                .await
-            {
-                Ok(p) => p,
-                Err(e) => {
-                    _ = result_sender.send(Err(e));
-                    return;
-                }
-            };
-            assert!(topics.insert(topic, subscriber).is_none());
-            _ = result_sender.send(Ok(true));
+            let result = swarm
+                .behaviour_mut()
+                .gossipsub
+                .publish(gossipsub::IdentTopic::new(topic), data);
+            _ = result_sender.send(result);
         }
     }
 }
 
+fn filter_swarm_event(event: SwarmEvent<BehaviourEvent>) -> Option<FromSwarm> {
+    match event {
+        SwarmEvent::Behaviour(BehaviourEvent::Gossipsub(gossipsub::Event::Message {
+            message:
+                gossipsub::Message {
+                    source: Some(peer_id),
+                    topic,
+                    data,
+                    ..
+                },
+            ..
+        })) => Some(FromSwarm::Message {
+            from: peer_id,
+            topic: topic.into_string(),
+            data,
+        }),
+        SwarmEvent::Behaviour(BehaviourEvent::Discovery(
+            discovery::Event::ConnectionEstablished { peer_id, .. },
+        )) => Some(FromSwarm::Discovered { peer_id }),
+        SwarmEvent::Behaviour(BehaviourEvent::Discovery(discovery::Event::ConnectionClosed {
+            peer_id,
+            ..
+        })) => Some(FromSwarm::Expired { peer_id }),
+        _ => None,
+    }
+}
+
+/// Create and configure a swarm.
+///
+/// - `listen_port`: TCP port to listen on. `0` lets the OS assign one.
+/// - `bootstrap_peers`: multiaddrs to dial for environments without mDNS.
 pub fn create_swarm(
-    identity: u128,
+    keypair: identity::Keypair,
     from_client: mpsc::Receiver<ToSwarm>,
     bootstrap_peers: Vec<String>,
     listen_port: u16,
-) -> Result<Swarm> {
-    // todo: bootstrap
-    if !bootstrap_peers.is_empty() || listen_port != 0 {
-        todo!();
-    }
-    let cfg = crate::cfg(identity, listen_port)?;
-    Ok(Swarm { cfg, from_client })
+) -> 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))?
+        .build();
+
+    swarm.listen_on(format!("/ip4/0.0.0.0/tcp/{listen_port}").parse()?)?;
+    Ok(Swarm { swarm, from_client })
+}
+
+mod transport {
+    use crate::alias;
+    use crate::swarm::{NETWORK_VERSION, OVERRIDE_VERSION_ENV_VAR};
+    use futures_lite::{AsyncRead, AsyncWrite};
+    use keccak_const::Sha3_256;
+    use libp2p::core::muxing;
+    use libp2p::core::transport::Boxed;
+    use libp2p::pnet::{PnetError, PnetOutput};
+    use libp2p::{PeerId, Transport, identity, noise, pnet, yamux};
+    use std::{env, sync::LazyLock};
+
+    /// Key used for networking's private network; parametrized on the [`NETWORK_VERSION`].
+    /// See [`pnet_upgrade`] for more.
+    static PNET_PRESHARED_KEY: LazyLock<[u8; 32]> = LazyLock::new(|| {
+        let builder = Sha3_256::new().update(b"exo_discovery_network");
+
+        if let Ok(var) = env::var(OVERRIDE_VERSION_ENV_VAR) {
+            let bytes = var.into_bytes();
+            builder.update(&bytes)
+        } else {
+            builder.update(NETWORK_VERSION)
+        }
+        .finalize()
+    });
+
+    /// Make the Swarm run on a private network, as to not clash with public libp2p nodes and
+    /// also different-versioned instances of this same network.
+    /// This is implemented as an additional "upgrade" ontop of existing [`libp2p::Transport`] layers.
+    async fn pnet_upgrade<TSocket>(
+        socket: TSocket,
+        _: impl Sized,
+    ) -> Result<PnetOutput<TSocket>, PnetError>
+    where
+        TSocket: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+    {
+        use pnet::{PnetConfig, PreSharedKey};
+        PnetConfig::new(PreSharedKey::new(*PNET_PRESHARED_KEY))
+            .handshake(socket)
+            .await
+    }
+
+    /// TCP/IP transport layer configuration.
+    pub fn tcp_transport(
+        keypair: &identity::Keypair,
+    ) -> alias::AnyResult<Boxed<(PeerId, muxing::StreamMuxerBox)>> {
+        use libp2p::{
+            core::upgrade::Version,
+            tcp::{Config, tokio},
+        };
+
+        // `TCP_NODELAY` enabled => avoid latency
+        let tcp_config = Config::default().nodelay(true);
+
+        // V1 + lazy flushing => 0-RTT negotiation
+        let upgrade_version = Version::V1Lazy;
+
+        // Noise is faster than TLS + we don't care much for security
+        let noise_config = noise::Config::new(keypair)?;
+
+        // Use default Yamux config for multiplexing
+        let yamux_config = yamux::Config::default();
+
+        // Create new Tokio-driven TCP/IP transport layer
+        let base_transport = tokio::Transport::new(tcp_config)
+            .and_then(pnet_upgrade)
+            .upgrade(upgrade_version)
+            .authenticate(noise_config)
+            .multiplex(yamux_config);
+
+        // Return boxed transport (to flatten complex type)
+        Ok(base_transport.boxed())
+    }
+}
+
+mod behaviour {
+    use crate::{alias, discovery};
+    use libp2p::swarm::NetworkBehaviour;
+    use libp2p::{gossipsub, identity};
+
+    /// Behavior of the Swarm which composes all desired behaviors:
+    /// Right now its just [`discovery::Behaviour`] and [`gossipsub::Behaviour`].
+    #[derive(NetworkBehaviour)]
+    pub struct Behaviour {
+        pub discovery: discovery::Behaviour,
+        pub gossipsub: gossipsub::Behaviour,
+    }
+
+    impl Behaviour {
+        pub fn new(
+            keypair: &identity::Keypair,
+            bootstrap_peers: Vec<libp2p::Multiaddr>,
+        ) -> alias::AnyResult<Self> {
+            Ok(Self {
+                discovery: discovery::Behaviour::new(keypair, bootstrap_peers)?,
+                gossipsub: gossipsub_behaviour(keypair),
+            })
+        }
+    }
+
+    fn gossipsub_behaviour(keypair: &identity::Keypair) -> gossipsub::Behaviour {
+        use gossipsub::{ConfigBuilder, MessageAuthenticity, ValidationMode};
+
+        // build a gossipsub network behaviour
+        //  => signed message authenticity + strict validation mode means the message-ID is
+        //     automatically provided by gossipsub w/out needing to provide custom message-ID function
+        gossipsub::Behaviour::new(
+            MessageAuthenticity::Signed(keypair.clone()),
+            ConfigBuilder::default()
+                .max_transmit_size(8 * 1024 * 1024)
+                .validation_mode(ValidationMode::Strict)
+                .build()
+                .expect("the configuration should always be valid"),
+        )
+        .expect("creating gossipsub behavior should always work")
+    }
 }