Merge remote-tracking branch 'origin/main' into alexcheema/prefill-progress

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

.github/workflows/pipeline.yml

@@ -8,6 +8,33 @@ on:
       - main
 
 jobs:
+  typecheck:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          lfs: false
+
+      - uses: cachix/install-nix-action@v31
+        with:
+          nix_path: nixpkgs=channel:nixos-unstable
+
+      - uses: cachix/cachix-action@v14
+        name: Configure Cachix
+        with:
+          name: exo
+          authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
+
+      - name: Load nix develop environment
+        run: nix run github:nicknovitski/nix-develop/v1
+
+      - name: Sync dependencies
+        run: uv sync --all-packages
+
+      - name: Run type checker
+        run: uv run basedpyright --project pyproject.toml
+
   nix:
     name: Build and check (${{ matrix.system }})
     runs-on: ${{ matrix.runner }}

.mlx_typings/mlx_lm/models/glm_moe_dsa.pyi

@@ -1,46 +0,0 @@
-"""Type stubs for mlx_lm.models.glm_moe_dsa"""
-
-from dataclasses import dataclass
-from typing import Any, Dict, Optional
-
-from .base import BaseModelArgs
-from .deepseek_v32 import Model as DSV32Model
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    model_type: str
-    vocab_size: int
-    hidden_size: int
-    index_head_dim: int
-    index_n_heads: int
-    index_topk: 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: Optional[int]
-    routed_scaling_factor: float
-    kv_lora_rank: int
-    q_lora_rank: int
-    qk_rope_head_dim: int
-    v_head_dim: int
-    qk_nope_head_dim: int
-    topk_method: str
-    scoring_func: str
-    norm_topk_prob: bool
-    n_group: int
-    topk_group: int
-    num_experts_per_tok: int
-    moe_layer_freq: int
-    first_k_dense_replace: int
-    max_position_embeddings: int
-    rms_norm_eps: float
-    rope_parameters: Dict[str, Any]
-    attention_bias: bool
-    rope_scaling: Dict[str, Any] | None
-    rope_theta: float | None
-
-class Model(DSV32Model):
-    def __init__(self, config: ModelArgs) -> None: ...

Cargo.lock

@@ -141,6 +141,12 @@ version = "0.3.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb"
 
+[[package]]
+name = "arrayvec"
+version = "0.7.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50"
+
 [[package]]
 name = "asn1-rs"
 version = "0.7.1"
@@ -298,6 +304,19 @@ version = "1.8.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "55248b47b0caf0546f7988906588779981c43bb1bc9d0c44087278f80cdb44ba"
 
+[[package]]
+name = "bigdecimal"
+version = "0.4.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "560f42649de9fa436b73517378a147ec21f6c997a546581df4b4b31677828934"
+dependencies = [
+ "autocfg",
+ "libm",
+ "num-bigint",
+ "num-integer",
+ "num-traits",
+]
+
 [[package]]
 name = "bimap"
 version = "0.6.3"
@@ -497,6 +516,15 @@ version = "0.4.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2f421161cb492475f1661ddc9815a745a1c894592070661180fdec3d4872e9c3"
 
+[[package]]
+name = "convert_case"
+version = "0.10.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "633458d4ef8c78b72454de2d54fd6ab2e60f9e02be22f3c6104cdc8a4e0fceb9"
+dependencies = [
+ "unicode-segmentation",
+]
+
 [[package]]
 name = "core-foundation"
 version = "0.9.4"
@@ -718,6 +746,29 @@ dependencies = [
  "powerfmt",
 ]
 
+[[package]]
+name = "derive_more"
+version = "2.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "10b768e943bed7bf2cab53df09f4bc34bfd217cdb57d971e769874c9a6710618"
+dependencies = [
+ "derive_more-impl",
+]
+
+[[package]]
+name = "derive_more-impl"
+version = "2.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6d286bfdaf75e988b4a78e013ecd79c581e06399ab53fbacd2d916c2f904f30b"
+dependencies = [
+ "convert_case",
+ "proc-macro2",
+ "quote",
+ "rustc_version",
+ "syn 2.0.111",
+ "unicode-xid",
+]
+
 [[package]]
 name = "digest"
 version = "0.10.7"
@@ -888,17 +939,22 @@ name = "exo_pyo3_bindings"
 version = "0.0.1"
 dependencies = [
  "delegate",
+ "derive_more",
  "env_logger",
  "extend",
- "futures-lite",
+ "futures",
+ "impl-trait-for-tuples",
  "libp2p",
  "log",
  "networking",
+ "once_cell",
  "pin-project",
  "pyo3",
  "pyo3-async-runtimes",
  "pyo3-log",
  "pyo3-stub-gen",
+ "thiserror 2.0.17",
+ "thread_local",
  "tokio",
  "util",
 ]
@@ -914,12 +970,6 @@ dependencies = [
  "syn 2.0.111",
 ]
 
-[[package]]
-name = "fastrand"
-version = "2.3.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be"
-
 [[package]]
 name = "ff"
 version = "0.13.1"
@@ -1028,10 +1078,7 @@ version = "2.6.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f78e10609fe0e0b3f4157ffab1876319b5b0db102a2c60dc4626306dc46b44ad"
 dependencies = [
- "fastrand",
  "futures-core",
- "futures-io",
- "parking",
  "pin-project-lite",
 ]
 
@@ -1593,6 +1640,17 @@ dependencies = [
  "xmltree",
 ]
 
+[[package]]
+name = "impl-trait-for-tuples"
+version = "0.2.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a0eb5a3343abf848c0984fe4604b2b105da9539376e24fc0a3b0007411ae4fd9"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.111",
+]
+
 [[package]]
 name = "indexmap"
 version = "2.12.1"
@@ -1771,6 +1829,12 @@ version = "0.2.178"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "37c93d8daa9d8a012fd8ab92f088405fb202ea0b6ab73ee2482ae66af4f42091"
 
+[[package]]
+name = "libm"
+version = "0.2.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f9fbbcab51052fe104eb5e5d351cf728d30a5be1fe14d9be8a3b097481fb97de"
+
 [[package]]
 name = "libp2p"
 version = "0.56.0"
@@ -2760,13 +2824,16 @@ name = "networking"
 version = "0.0.1"
 dependencies = [
  "delegate",
+ "derive_more",
  "either",
  "extend",
- "futures-lite",
+ "futures",
  "futures-timer",
+ "impl-trait-for-tuples",
  "keccak-const",
  "libp2p",
  "log",
+ "thiserror 2.0.17",
  "tokio",
  "tracing-subscriber",
  "util",
@@ -2851,6 +2918,17 @@ dependencies = [
  "num-traits",
 ]
 
+[[package]]
+name = "num-rational"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824"
+dependencies = [
+ "num-bigint",
+ "num-integer",
+ "num-traits",
+]
+
 [[package]]
 name = "num-traits"
 version = "0.2.19"
@@ -3201,14 +3279,28 @@ version = "0.27.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ab53c047fcd1a1d2a8820fe84f05d6be69e9526be40cb03b73f86b6b03e6d87d"
 dependencies = [
+ "bigdecimal",
+ "either",
+ "hashbrown 0.16.1",
+ "indexmap",
  "indoc",
+ "inventory",
  "libc",
+ "lock_api",
  "memoffset",
+ "num-bigint",
+ "num-complex",
+ "num-rational",
+ "num-traits",
  "once_cell",
+ "ordered-float",
+ "parking_lot",
  "portable-atomic",
  "pyo3-build-config",
  "pyo3-ffi",
  "pyo3-macros",
+ "rust_decimal",
+ "smallvec",
  "unindent",
 ]
 
@@ -3649,6 +3741,16 @@ dependencies = [
  "tokio",
 ]
 
+[[package]]
+name = "rust_decimal"
+version = "1.39.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "35affe401787a9bd846712274d97654355d21b2a2c092a3139aabe31e9022282"
+dependencies = [
+ "arrayvec",
+ "num-traits",
+]
+
 [[package]]
 name = "rustc-hash"
 version = "1.1.0"
@@ -4513,12 +4615,24 @@ version = "1.0.22"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9312f7c4f6ff9069b165498234ce8be658059c6728633667c526e27dc2cf1df5"
 
+[[package]]
+name = "unicode-segmentation"
+version = "1.12.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493"
+
 [[package]]
 name = "unicode-width"
 version = "0.2.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254"
 
+[[package]]
+name = "unicode-xid"
+version = "0.2.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853"
+
 [[package]]
 name = "unicode_names2"
 version = "1.3.0"

Cargo.toml

@@ -26,20 +26,49 @@ opt-level = 3
 networking = { path = "rust/networking" }
 util = { path = "rust/util" }
 
+# Proc-macro authoring tools
+syn = "2.0"
+quote = "1.0"
+proc-macro2 = "1.0"
+darling = "0.20"
+
 # Macro dependecies
 extend = "1.2"
 delegate = "0.13"
+impl-trait-for-tuples = "0.2"
+clap = "4.5"
+derive_more = { version = "2.0.1", features = ["display"] }
+pin-project = "1"
 
 # Utility dependencies
+itertools = "0.14"
+thiserror = "2"
+internment = "0.8"
+recursion = "0.5"
+regex = "1.11"
+once_cell = "1.21"
+thread_local = "1.1"
+bon = "3.4"
+generativity = "1.1"
+anyhow = "1.0"
 keccak-const = "0.2"
 
+# Functional generics/lenses frameworks
+frunk_core = "0.4"
+frunk = "0.4"
+frunk_utils = "0.2"
+frunk-enum-core = "0.3"
+
 # Async dependencies
 tokio = "1.46"
-futures-lite = "2.6.1"
+futures = "0.3"
+futures-util = "0.3"
 futures-timer = "3.0"
 
 # Data structures
 either = "1.15"
+ordered-float = "5.0"
+ahash = "0.8"
 
 # Tracing/logging
 log = "0.4"

MISSED_THINGS.md

@@ -5,21 +5,21 @@
 [X] Fetching download status of all models on start
 [X] Deduplication of tasks in plan_step.
 [X] resolve_allow_patterns should just be wildcard now.
-[X] no mx_barrier in genreate.py mlx_generate at the end.
+[] no mx_barrier in genreate.py mlx_generate at the end.
 [] cache assertion not needed in auto_parallel.py PipelineLastLayer.
-[X] GPTOSS support dropped in auto_parallel.py.
-[X] sharding changed "all-to-sharded" became _all_to_sharded in auto_parallel.py.
-[X] same as above with "sharded-to-all" became _sharded_to_all in auto_parallel.py.
-[X] Dropped support for Ministral3Model, DeepseekV32Model, Glm4MoeModel, Qwen3NextModel, GptOssMode in auto_parallel.py.
+[] GPTOSS support dropped in auto_parallel.py.
+[] sharding changed "all-to-sharded" became _all_to_sharded in auto_parallel.py.
+[] same as above with "sharded-to-all" became _sharded_to_all in auto_parallel.py.
+[] Dropped support for Ministral3Model, DeepseekV32Model, Glm4MoeModel, Qwen3NextModel, GptOssMode in auto_parallel.py.
 [] Dropped prefill/decode code in auto_parallel.py and utils_mlx.py.
 [X] KV_CACHE_BITS should be None to disable quantized KV cache.
-[X] Dropped _set_nofile_limit in utils_mlx.py.
-[X] We have group optional in load_mlx_items in utils_mlx.py.
-[X] Dropped add_missing_chat_templates for GptOss in load_mlx_items in utils_mlx.py.
-[X] Dropped model.make_cache in make_kv_cache in utils_mlx.py.
+[] Dropped _set_nofile_limit in utils_mlx.py.
+[] We have group optional in load_mlx_items in utils_mlx.py.
+[] Dropped add_missing_chat_templates for GptOss in load_mlx_items in utils_mlx.py.
+[] Dropped model.make_cache in make_kv_cache in utils_mlx.py.
 [X] We put cache limit back in utils_mlx.py.
-[X] topology.py remove_node removes the connections after checking if node is is in self._node_id_to_rx_id_map. on beta_1 it checks after, so would remove stale connections I guess?
-[X] Missing Glm 4.7 model cards (this isn't ready yet but should be picked up, probably create an issue... the blocker is transforemrs version doesn't support the tokenizer for Glm 4.7. rc-1 does but we can't upgrade as it breaks other things.)
+[] topology.py remove_node removes the connections after checking if node is is in self._node_id_to_rx_id_map. on beta_1 it checks after, so would remove stale connections I guess?
+[] Missing Glm 4.7 model cards (this isn't ready yet but should be picked up, probably create an issue... the blocker is transforemrs version doesn't support the tokenizer for Glm 4.7. rc-1 does but we can't upgrade as it breaks other things.)
 [] try-except in _command_processor only excepts ValueError. This was silently failing leading to un-debuggable errors (we had a KeyError that was happening ). Changed this to catch Exception instead of ValueError. See exo-v2 89ae38405e0052e3c22405daf094b065878aa873 and fb99fea69b5a39017efc90c5dad0072e677455f0.
 [X] In placement.py, place_instance no longer looks at model_meta.supports_tensor and check if this tensor parallel number of nodes is supported by the model's tensor dimensions.
 [X] In placement.py, place_instanec, we no longer have the special case to exclude DeepSeek v3.1 pipeline parallel (it doesn't work).

README.md

@@ -72,23 +72,16 @@ There are two ways to run exo:
 
 ### Run from Source (macOS)
 
-If you have [Nix](https://nixos.org/) installed, you can skip most of the steps below and run exo directly (after accepting the Cachix cache):
-
-```bash
-nix run .#exo
-```
-
 **Prerequisites:**
-- [Xcode](https://developer.apple.com/xcode/) (provides the Metal ToolChain required for MLX compilation)
 - [brew](https://github.com/Homebrew/brew) (for simple package management on macOS)
-
+  
   ```bash
   /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"
   ```
 - [uv](https://github.com/astral-sh/uv) (for Python dependency management)
 - [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
 - [node](https://github.com/nodejs/node) (for building the dashboard)
-
+  
   ```bash
   brew install uv macmon node
   ```

app/EXO/EXO/ExoProcessController.swift

@@ -126,37 +126,11 @@ final class ExoProcessController: ObservableObject {
             return
         }
         process.terminationHandler = nil
-        status = .stopped
-
-        guard process.isRunning else {
-            self.process = nil
-            return
+        if process.isRunning {
+            process.terminate()
         }
-
-        let proc = process
         self.process = nil
-
-        Task.detached {
-            proc.interrupt()
-
-            for _ in 0..<50 {
-                if !proc.isRunning { return }
-                try? await Task.sleep(nanoseconds: 100_000_000)
-            }
-
-            if proc.isRunning {
-                proc.terminate()
-            }
-
-            for _ in 0..<30 {
-                if !proc.isRunning { return }
-                try? await Task.sleep(nanoseconds: 100_000_000)
-            }
-
-            if proc.isRunning {
-                kill(proc.processIdentifier, SIGKILL)
-            }
-        }
+        status = .stopped
     }
 
     func restart() {

bench/bench.toml

@@ -1,7 +0,0 @@
-# Canary benchmark manifest
-#
-# Lists the suite files to include. Each file defines benchmarks
-# with shared constraints, topology, and default args.
-include = [
-    "single-m3-ultra.toml",
-]

bench/exo_bench.py

@@ -288,151 +288,6 @@ def resolve_model_short_id(client: ExoClient, model_arg: str) -> tuple[str, str]
     raise ValueError(f"Model not found in /models: {model_arg}")
 
 
-def run_planning_phase(
-    client: ExoClient,
-    full_model_id: str,
-    preview: dict[str, Any],
-    danger_delete: bool,
-    timeout: float,
-    settle_deadline: float | None,
-) -> None:
-    """Check disk space and ensure model is downloaded before benchmarking."""
-    # Get model size from /models
-    models = client.request_json("GET", "/models") or {}
-    model_bytes = 0
-    for m in models.get("data", []):
-        if m.get("hugging_face_id") == full_model_id:
-            model_bytes = m.get("storage_size_megabytes", 0) * 1024 * 1024
-            break
-
-    if not model_bytes:
-        logger.warning(
-            f"Could not determine size for {full_model_id}, skipping disk check"
-        )
-        return
-
-    # Get nodes from preview
-    inner = unwrap_instance(preview["instance"])
-    node_ids = list(inner["shardAssignments"]["nodeToRunner"].keys())
-    runner_to_shard = inner["shardAssignments"]["runnerToShard"]
-
-    state = client.request_json("GET", "/state")
-    downloads = state.get("downloads", {})
-    node_disk = state.get("nodeDisk", {})
-
-    for node_id in node_ids:
-        node_downloads = downloads.get(node_id, [])
-
-        # Check if model already downloaded on this node
-        already_downloaded = any(
-            "DownloadCompleted" in p
-            and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
-                "modelId"
-            ]
-            == full_model_id
-            for p in node_downloads
-        )
-        if already_downloaded:
-            continue
-
-        # Wait for disk info if settle_deadline is set
-        disk_info = node_disk.get(node_id, {})
-        backoff = _SETTLE_INITIAL_BACKOFF_S
-        while not disk_info and settle_deadline and time.monotonic() < settle_deadline:
-            remaining = settle_deadline - time.monotonic()
-            logger.info(
-                f"Waiting for disk info on {node_id} ({remaining:.0f}s remaining)..."
-            )
-            time.sleep(min(backoff, remaining))
-            backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S)
-            state = client.request_json("GET", "/state")
-            node_disk = state.get("nodeDisk", {})
-            disk_info = node_disk.get(node_id, {})
-
-        if not disk_info:
-            logger.warning(f"No disk info for {node_id}, skipping space check")
-            continue
-
-        avail = disk_info.get("available", {}).get("inBytes", 0)
-        if avail >= model_bytes:
-            continue
-
-        if not danger_delete:
-            raise RuntimeError(
-                f"Insufficient disk on {node_id}: need {model_bytes // (1024**3)}GB, "
-                f"have {avail // (1024**3)}GB. Use --danger-delete-downloads to free space."
-            )
-
-        # Delete from smallest to largest
-        completed = [
-            (
-                unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
-                    "modelId"
-                ],
-                p["DownloadCompleted"]["totalBytes"]["inBytes"],
-            )
-            for p in node_downloads
-            if "DownloadCompleted" in p
-        ]
-        for del_model, size in sorted(completed, key=lambda x: x[1]):
-            logger.info(f"Deleting {del_model} from {node_id} ({size // (1024**2)}MB)")
-            client.request_json("DELETE", f"/download/{node_id}/{del_model}")
-            avail += size
-            if avail >= model_bytes:
-                break
-
-        if avail < model_bytes:
-            raise RuntimeError(f"Could not free enough space on {node_id}")
-
-    # Start downloads (idempotent)
-    for node_id in node_ids:
-        runner_id = inner["shardAssignments"]["nodeToRunner"][node_id]
-        shard = runner_to_shard[runner_id]
-        client.request_json(
-            "POST",
-            "/download/start",
-            body={
-                "targetNodeId": node_id,
-                "shardMetadata": shard,
-            },
-        )
-        logger.info(f"Started download on {node_id}")
-
-    # Wait for downloads
-    start = time.time()
-    while time.time() - start < timeout:
-        state = client.request_json("GET", "/state")
-        downloads = state.get("downloads", {})
-        all_done = True
-        for node_id in node_ids:
-            done = any(
-                "DownloadCompleted" in p
-                and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])[
-                    "modelCard"
-                ]["modelId"]
-                == full_model_id
-                for p in downloads.get(node_id, [])
-            )
-            failed = [
-                p["DownloadFailed"]["errorMessage"]
-                for p in downloads.get(node_id, [])
-                if "DownloadFailed" in p
-                and unwrap_instance(p["DownloadFailed"]["shardMetadata"])["modelCard"][
-                    "modelId"
-                ]
-                == full_model_id
-            ]
-            if failed:
-                raise RuntimeError(f"Download failed on {node_id}: {failed[0]}")
-            if not done:
-                all_done = False
-        if all_done:
-            return
-        time.sleep(1)
-
-    raise TimeoutError("Downloads did not complete in time")
-
-
 def placement_filter(instance_meta: str, wanted: str) -> bool:
     s = (instance_meta or "").lower()
     if wanted == "both":
@@ -680,11 +535,6 @@ def main() -> int:
         default=0,
         help="Max seconds to wait for the cluster to produce valid placements (0 = try once).",
     )
-    ap.add_argument(
-        "--danger-delete-downloads",
-        action="store_true",
-        help="Delete existing models from smallest to largest to make room for benchmark model.",
-    )
     args = ap.parse_args()
 
     pp_list = parse_int_list(args.pp)
@@ -719,16 +569,13 @@ def main() -> int:
         logger.error("[exo-bench] tokenizer usable but prompt sizing failed")
         raise
 
-    settle_deadline = (
-        time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
-    )
-
     selected = fetch_and_filter_placements(client, full_model_id, args)
 
-    if not selected and settle_deadline:
+    if not selected and args.settle_timeout > 0:
         backoff = _SETTLE_INITIAL_BACKOFF_S
-        while not selected and time.monotonic() < settle_deadline:
-            remaining = settle_deadline - time.monotonic()
+        deadline = time.monotonic() + args.settle_timeout
+        while not selected and time.monotonic() < deadline:
+            remaining = deadline - time.monotonic()
             logger.warning(
                 f"No valid placements yet (cluster may still be settling). "
                 f"Retrying in {backoff:.1f}s ({remaining:.0f}s remaining)..."
@@ -760,16 +607,6 @@ def main() -> int:
     if args.dry_run:
         return 0
 
-    logger.info("Planning phase: checking downloads...")
-    run_planning_phase(
-        client,
-        full_model_id,
-        selected[0],
-        args.danger_delete_downloads,
-        args.timeout,
-        settle_deadline,
-    )
-
     all_rows: list[dict[str, Any]] = []
 
     for preview in selected:

bench/single-m3-ultra.toml

@@ -1,189 +0,0 @@
-# Single-node M3 Ultra benchmarks
-#
-# Shared constraints applied to ALL benchmarks in this file.
-constraints = [
-    "All(MacOsBuild(=25D125))",
-    "Hosts(=1)",
-    "All(Chip(m3_ultra))",
-    "All(GpuCores(=80))",
-]
-
-[topology]
-type = "none"
-
-# Default args merged into each benchmark's args (benchmark-level args win).
-[defaults]
-pp = [512, 2048, 8192, 16384]
-tg = 128
-
-[[benchmark]]
-model = "mlx-community/Meta-Llama-3.1-70B-Instruct-4bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/gpt-oss-120b-MXFP4-Q8"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/GLM-4.7-Flash-8bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Qwen3-Coder-Next-6bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Qwen3-30B-A3B-8bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Qwen3-0.6B-4bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Qwen3-0.6B-8bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Llama-3.2-1B-Instruct-4bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Llama-3.2-3B-Instruct-4bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Llama-3.2-3B-Instruct-8bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Meta-Llama-3.1-8B-Instruct-4bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Meta-Llama-3.1-8B-Instruct-8bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Meta-Llama-3.1-8B-Instruct-bf16"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/gpt-oss-20b-MXFP4-Q8"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Qwen3-30B-A3B-4bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/GLM-4.7-Flash-4bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/GLM-4.7-Flash-5bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/GLM-4.7-Flash-6bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Llama-3.3-70B-Instruct-4bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Qwen3-Coder-Next-4bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Qwen3-Coder-Next-5bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Qwen3-Coder-Next-8bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Qwen3-Next-80B-A3B-Instruct-4bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Qwen3-Next-80B-A3B-Instruct-8bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Qwen3-Next-80B-A3B-Thinking-4bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Qwen3-Next-80B-A3B-Thinking-8bit"
-extra_constraints = ["All(Memory(>=96GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Llama-3.3-70B-Instruct-8bit"
-extra_constraints = ["All(Memory(>=256GiB))"]
-
-[[benchmark]]
-model = "mlx-community/llama-3.3-70b-instruct-fp16"
-extra_constraints = ["All(Memory(>=256GiB))"]
-
-[[benchmark]]
-model = "mlx-community/GLM-4.5-Air-8bit"
-extra_constraints = ["All(Memory(>=256GiB))"]
-
-[[benchmark]]
-model = "mlx-community/GLM-4.5-Air-bf16"
-extra_constraints = ["All(Memory(>=256GiB))"]
-
-[[benchmark]]
-model = "mlx-community/GLM-4.7-4bit"
-extra_constraints = ["All(Memory(>=256GiB))"]
-
-[[benchmark]]
-model = "mlx-community/MiniMax-M2.1-3bit"
-extra_constraints = ["All(Memory(>=256GiB))"]
-
-[[benchmark]]
-model = "mlx-community/MiniMax-M2.1-8bit"
-extra_constraints = ["All(Memory(>=256GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Qwen3-235B-A22B-Instruct-2507-4bit"
-extra_constraints = ["All(Memory(>=256GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Qwen3-Coder-Next-bf16"
-extra_constraints = ["All(Memory(>=256GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Step-3.5-Flash-4bit"
-extra_constraints = ["All(Memory(>=256GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Step-3.5-Flash-6bit"
-extra_constraints = ["All(Memory(>=256GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Step-3.5-Flash-8Bit"
-extra_constraints = ["All(Memory(>=256GiB))"]
-
-[[benchmark]]
-model = "mlx-community/DeepSeek-V3.1-4bit"
-extra_constraints = ["All(Memory(>=512GiB))"]
-
-[[benchmark]]
-model = "mlx-community/GLM-4.7-6bit"
-extra_constraints = ["All(Memory(>=512GiB))"]
-
-[[benchmark]]
-model = "mlx-community/GLM-4.7-8bit-gs32"
-extra_constraints = ["All(Memory(>=512GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Qwen3-235B-A22B-Instruct-2507-8bit"
-extra_constraints = ["All(Memory(>=512GiB))"]
-
-[[benchmark]]
-model = "mlx-community/Qwen3-Coder-480B-A35B-Instruct-4bit"
-extra_constraints = ["All(Memory(>=512GiB))"]

dashboard/src/lib/components/ChatForm.svelte

@@ -103,7 +103,7 @@
   const modelSupportsThinking = $derived(() => {
     if (!currentModel) return false;
     const caps = modelCapabilities[currentModel] || [];
-    return caps.includes("thinking_toggle") && caps.includes("text");
+    return caps.includes("thinking") && caps.includes("text");
   });
 
   const isEditOnlyWithoutImage = $derived(
@@ -265,7 +265,6 @@
 
   function handleSubmit() {
     if ((!message.trim() && uploadedFiles.length === 0) || loading) return;
-    if (isEditOnlyWithoutImage) return;
 
     const content = message.trim();
     const files = [...uploadedFiles];
@@ -290,11 +289,7 @@
       if (imageFile.preview) {
         editImage(content, imageFile.preview);
       }
-    } else if (
-      currentModel &&
-      modelSupportsTextToImage(currentModel) &&
-      content
-    ) {
+    } else if (isImageModel() && content) {
       // Use image generation for text-to-image models
       generateImage(content);
     } else {

dashboard/src/lib/components/ChatMessages.svelte

@@ -9,7 +9,6 @@
     regenerateFromToken,
     setEditingImage,
   } from "$lib/stores/app.svelte";
-  import type { Message } from "$lib/stores/app.svelte";
   import type { MessageAttachment } from "$lib/stores/app.svelte";
   import MarkdownContent from "./MarkdownContent.svelte";
   import TokenHeatmap from "./TokenHeatmap.svelte";
@@ -225,7 +224,6 @@
   }
 
   function handleDeleteClick(messageId: string) {
-    if (loading) return;
     deleteConfirmId = messageId;
   }
 
@@ -256,7 +254,7 @@
 </script>
 
 <div class="flex flex-col gap-4 sm:gap-6 {className}">
-  {#each messageList as message, i (message.id)}
+  {#each messageList as message (message.id)}
     <div
       class="group flex {message.role === 'user'
         ? 'justify-end'
@@ -318,11 +316,9 @@
           <!-- Delete confirmation -->
           <div class="bg-red-500/10 border border-red-500/30 rounded-lg p-3">
             <p class="text-xs text-red-400 mb-3">
-              {#if i === messageList.length - 1}
-                Delete this message?
-              {:else}
-                Delete this message and all messages after it?
-              {/if}
+              Delete this message{message.role === "user"
+                ? " and all responses after it"
+                : ""}?
             </p>
             <div class="flex gap-2 justify-end">
               <button
@@ -754,13 +750,8 @@
             <!-- Delete button -->
             <button
               onclick={() => handleDeleteClick(message.id)}
-              disabled={loading}
-              class="p-1.5 transition-colors rounded {loading
-                ? 'text-exo-light-gray/30 cursor-not-allowed'
-                : 'text-exo-light-gray hover:text-red-400 hover:bg-red-500/10 cursor-pointer'}"
-              title={loading
-                ? "Cannot delete while generating"
-                : "Delete message"}
+              class="p-1.5 text-exo-light-gray hover:text-red-400 transition-colors rounded hover:bg-red-500/10 cursor-pointer"
+              title="Delete message"
             >
               <svg
                 class="w-3.5 h-3.5"

dashboard/src/lib/components/ImageParamsPanel.svelte

@@ -59,14 +59,13 @@
   }
 
   const sizeOptions: ImageGenerationParams["size"][] = [
-    "auto",
     "512x512",
     "768x768",
     "1024x1024",
     "1024x768",
     "768x1024",
-    "1024x1536",
-    "1536x1024",
+    "1024x1365",
+    "1365x1024",
   ];
 
   const qualityOptions: ImageGenerationParams["quality"][] = [
@@ -177,90 +176,92 @@
 <div class="border-b border-exo-medium-gray/30 px-3 py-2">
   <!-- Basic params row -->
   <div class="flex items-center gap-3 flex-wrap">
-    <!-- Size -->
-    <div class="flex items-center gap-1.5">
-      <span class="text-xs text-exo-light-gray uppercase tracking-wider"
-        >SIZE:</span
-      >
-      <div class="relative">
-        <button
-          bind:this={sizeButtonRef}
-          type="button"
-          onclick={() => (isSizeDropdownOpen = !isSizeDropdownOpen)}
-          class="bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-2 pr-6 py-1 text-xs font-mono text-exo-yellow cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70 {isSizeDropdownOpen
-            ? 'border-exo-yellow/70'
-            : ''}"
+    <!-- Size (hidden in edit mode - output size comes from input image) -->
+    {#if !isEditMode}
+      <div class="flex items-center gap-1.5">
+        <span class="text-xs text-exo-light-gray uppercase tracking-wider"
+          >SIZE:</span
         >
-          {params.size.toUpperCase()}
-        </button>
-        <div
-          class="absolute right-1.5 top-1/2 -translate-y-1/2 pointer-events-none transition-transform duration-200 {isSizeDropdownOpen
-            ? 'rotate-180'
-            : ''}"
-        >
-          <svg
-            class="w-3 h-3 text-exo-yellow/60"
-            fill="none"
-            viewBox="0 0 24 24"
-            stroke="currentColor"
+        <div class="relative">
+          <button
+            bind:this={sizeButtonRef}
+            type="button"
+            onclick={() => (isSizeDropdownOpen = !isSizeDropdownOpen)}
+            class="bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-2 pr-6 py-1 text-xs font-mono text-exo-yellow cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70 {isSizeDropdownOpen
+              ? 'border-exo-yellow/70'
+              : ''}"
           >
-            <path
-              stroke-linecap="round"
-              stroke-linejoin="round"
-              stroke-width="2"
-              d="M19 9l-7 7-7-7"
-            />
-          </svg>
-        </div>
-      </div>
-
-      {#if isSizeDropdownOpen}
-        <!-- Backdrop to close dropdown -->
-        <button
-          type="button"
-          class="fixed inset-0 z-[9998] cursor-default"
-          onclick={() => (isSizeDropdownOpen = false)}
-          aria-label="Close dropdown"
-        ></button>
-
-        <!-- Dropdown Panel - fixed positioning to escape overflow:hidden -->
-        <div
-          class="fixed bg-exo-dark-gray border border-exo-yellow/30 rounded shadow-lg shadow-black/50 z-[9999] max-h-48 overflow-y-auto overflow-x-hidden min-w-max"
-          style="bottom: calc(100vh - {sizeDropdownPosition()
-            .top}px + 4px); left: {sizeDropdownPosition().left}px;"
-        >
-          <div class="py-1">
-            {#each sizeOptions as size}
-              <button
-                type="button"
-                onclick={() => selectSize(size)}
-                class="w-full px-3 py-1.5 text-left text-xs font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {params.size ===
-                size
-                  ? 'bg-transparent text-exo-yellow'
-                  : 'text-exo-light-gray hover:text-exo-yellow'}"
-              >
-                {#if params.size === size}
-                  <svg
-                    class="w-3 h-3 flex-shrink-0"
-                    fill="currentColor"
-                    viewBox="0 0 20 20"
-                  >
-                    <path
-                      fill-rule="evenodd"
-                      d="M16.707 5.293a1 1 0 010 1.414l-8 8a1 1 0 01-1.414 0l-4-4a1 1 0 011.414-1.414L8 12.586l7.293-7.293a1 1 0 011.414 0z"
-                      clip-rule="evenodd"
-                    />
-                  </svg>
-                {:else}
-                  <span class="w-3"></span>
-                {/if}
-                <span>{size.toUpperCase()}</span>
-              </button>
-            {/each}
+            {params.size}
+          </button>
+          <div
+            class="absolute right-1.5 top-1/2 -translate-y-1/2 pointer-events-none transition-transform duration-200 {isSizeDropdownOpen
+              ? 'rotate-180'
+              : ''}"
+          >
+            <svg
+              class="w-3 h-3 text-exo-yellow/60"
+              fill="none"
+              viewBox="0 0 24 24"
+              stroke="currentColor"
+            >
+              <path
+                stroke-linecap="round"
+                stroke-linejoin="round"
+                stroke-width="2"
+                d="M19 9l-7 7-7-7"
+              />
+            </svg>
           </div>
         </div>
-      {/if}
-    </div>
+
+        {#if isSizeDropdownOpen}
+          <!-- Backdrop to close dropdown -->
+          <button
+            type="button"
+            class="fixed inset-0 z-[9998] cursor-default"
+            onclick={() => (isSizeDropdownOpen = false)}
+            aria-label="Close dropdown"
+          ></button>
+
+          <!-- Dropdown Panel - fixed positioning to escape overflow:hidden -->
+          <div
+            class="fixed bg-exo-dark-gray border border-exo-yellow/30 rounded shadow-lg shadow-black/50 z-[9999] max-h-48 overflow-y-auto min-w-max"
+            style="bottom: calc(100vh - {sizeDropdownPosition()
+              .top}px + 4px); left: {sizeDropdownPosition().left}px;"
+          >
+            <div class="py-1">
+              {#each sizeOptions as size}
+                <button
+                  type="button"
+                  onclick={() => selectSize(size)}
+                  class="w-full px-3 py-1.5 text-left text-xs font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {params.size ===
+                  size
+                    ? 'bg-transparent text-exo-yellow'
+                    : 'text-exo-light-gray hover:text-exo-yellow'}"
+                >
+                  {#if params.size === size}
+                    <svg
+                      class="w-3 h-3 flex-shrink-0"
+                      fill="currentColor"
+                      viewBox="0 0 20 20"
+                    >
+                      <path
+                        fill-rule="evenodd"
+                        d="M16.707 5.293a1 1 0 010 1.414l-8 8a1 1 0 01-1.414 0l-4-4a1 1 0 011.414-1.414L8 12.586l7.293-7.293a1 1 0 011.414 0z"
+                        clip-rule="evenodd"
+                      />
+                    </svg>
+                  {:else}
+                    <span class="w-3"></span>
+                  {/if}
+                  <span>{size}</span>
+                </button>
+              {/each}
+            </div>
+          </div>
+        {/if}
+      </div>
+    {/if}
 
     <!-- Quality -->
     <div class="flex items-center gap-1.5">
@@ -310,7 +311,7 @@
 
         <!-- Dropdown Panel - fixed positioning to escape overflow:hidden -->
         <div
-          class="fixed bg-exo-dark-gray border border-exo-yellow/30 rounded shadow-lg shadow-black/50 z-[9999] max-h-48 overflow-y-auto overflow-x-hidden min-w-max"
+          class="fixed bg-exo-dark-gray border border-exo-yellow/30 rounded shadow-lg shadow-black/50 z-[9999] max-h-48 overflow-y-auto min-w-max"
           style="bottom: calc(100vh - {qualityDropdownPosition()
             .top}px + 4px); left: {qualityDropdownPosition().left}px;"
         >

dashboard/src/lib/components/PrefillProgressBar.svelte

@@ -0,0 +1,51 @@
+<script lang="ts">
+  import type { PrefillProgress } from "$lib/stores/app.svelte";
+
+  interface Props {
+    progress: PrefillProgress;
+    class?: string;
+  }
+
+  let { progress, class: className = "" }: Props = $props();
+
+  const percentage = $derived(
+    progress.total > 0
+      ? Math.round((progress.processed / progress.total) * 100)
+      : 0,
+  );
+
+  function formatTokenCount(count: number): string {
+    if (count >= 1000) {
+      return `${(count / 1000).toFixed(1)}k`;
+    }
+    return count.toString();
+  }
+</script>
+
+<div class="prefill-progress {className}">
+  <div
+    class="flex items-center justify-between text-xs text-exo-light-gray mb-1"
+  >
+    <span>Processing prompt</span>
+    <span class="font-mono">
+      {formatTokenCount(progress.processed)} / {formatTokenCount(
+        progress.total,
+      )} tokens
+    </span>
+  </div>
+  <div class="h-1.5 bg-exo-black/60 rounded-full overflow-hidden">
+    <div
+      class="h-full bg-exo-yellow rounded-full transition-all duration-150 ease-out"
+      style="width: {percentage}%"
+    ></div>
+  </div>
+  <div class="text-right text-xs text-exo-light-gray/70 mt-0.5 font-mono">
+    {percentage}%
+  </div>
+</div>
+
+<style>
+  .prefill-progress {
+    width: 100%;
+  }
+</style>

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

@@ -273,6 +273,11 @@ export interface TokenData {
   topLogprobs: TopLogprob[];
 }
 
+export interface PrefillProgress {
+  processed: number;
+  total: number;
+}
+
 export interface Message {
   id: string;
   role: "user" | "assistant" | "system";
@@ -306,14 +311,13 @@ const IMAGE_PARAMS_STORAGE_KEY = "exo-image-generation-params";
 export interface ImageGenerationParams {
   // Basic params
   size:
-    | "auto"
     | "512x512"
     | "768x768"
     | "1024x1024"
     | "1024x768"
     | "768x1024"
-    | "1024x1536"
-    | "1536x1024";
+    | "1024x1365"
+    | "1365x1024";
   quality: "low" | "medium" | "high";
   outputFormat: "png" | "jpeg";
   numImages: number;
@@ -337,7 +341,7 @@ export interface EditingImage {
 }
 
 const DEFAULT_IMAGE_PARAMS: ImageGenerationParams = {
-  size: "auto",
+  size: "1024x1024",
   quality: "medium",
   outputFormat: "png",
   numImages: 1,

flake.nix

@@ -74,6 +74,7 @@
       perSystem =
         { config, self', inputs', pkgs, lib, system, ... }:
         let
+          fenixToolchain = inputs'.fenix.packages.complete;
           # Use pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs)
           pkgsSwift = import inputs.nixpkgs-swift { inherit system; };
         in
@@ -114,7 +115,7 @@
           packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin (
             let
               uvLock = builtins.fromTOML (builtins.readFile ./uv.lock);
-              mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx" && p.source ? git) uvLock.package);
+              mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx") uvLock.package);
               uvLockMlxVersion = mlxPackage.version;
             in
             {

nix/mlx.nix

@@ -41,16 +41,16 @@ let
 
   mlx = stdenv.mkDerivation rec {
     pname = "mlx";
-    version = let v = "0.30.7.dev20260218+14841977"; in
+    version = let v = "0.30.6"; in
       assert v == uvLockMlxVersion || throw "MLX version mismatch: nix/mlx.nix has ${v} but uv.lock has ${uvLockMlxVersion}. Update both the version and hash in nix/mlx.nix.";
       v;
     pyproject = true;
 
     src = fetchFromGitHub {
-      owner = "rltakashige";
-      repo = "mlx-jaccl-fix-small-recv";
-      rev = "1484197707f35186ad3bd614357c7c47fdf86ebc";
-      hash = "sha256-FupCMoK/SF/ldfKuvMSAKECcOP8c+ANgkQlPZttDsLk=";
+      owner = "ml-explore";
+      repo = "mlx";
+      tag = "v${version}";
+      hash = "sha256-avD5EGhwgmPdXLAyQSqTO6AXk/W3ziH+f6AetjK3Sdo=";
     };
 
     patches = [

pyproject.toml

@@ -17,9 +17,9 @@ dependencies = [
     "loguru>=0.7.3",
     "exo_pyo3_bindings", # rust bindings
     "anyio==4.11.0",
-    "mlx; sys_platform == 'darwin'",
+    "mlx==0.30.6; sys_platform == 'darwin'",
     "mlx[cpu]==0.30.6; sys_platform == 'linux'",
-    "mlx-lm==0.30.7",
+    "mlx-lm==0.30.6",
     "tiktoken>=0.12.0", # required for kimi k2 tokenizer
     "hypercorn>=0.18.0",
     "openai-harmony>=0.0.8",
@@ -64,7 +64,6 @@ members = [
 
 [tool.uv.sources]
 exo_pyo3_bindings = { workspace = true }
-mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
 #mlx-lm = { git = "https://github.com/davidmcc73/mlx-lm", branch = "stable" }
 # Uncomment to use local mlx/mlx-lm development versions:
 # mlx = { path = "/Users/Shared/mlx", editable=true }
@@ -133,7 +132,7 @@ markers = [
 env = [
   "EXO_TESTS=1"
 ]
-addopts = "-m 'not slow' --ignore=tests/start_distributed_test.py"
+addopts = "-m 'not slow'"
 filterwarnings = [
     "ignore:builtin type Swig:DeprecationWarning",
 ]

python/parts.nix

@@ -14,9 +14,7 @@
 
       # Override overlay to inject Nix-built components
       exoOverlay = final: prev: {
-        # Replace workspace exo_pyo3_bindings with Nix-built wheel.
-        # Preserve passthru so mkVirtualEnv can resolve dependency groups.
-        # Copy .pyi stub + py.typed marker so basedpyright can find the types.
+        # Replace workspace exo_pyo3_bindings with Nix-built wheel
         exo-pyo3-bindings = pkgs.stdenv.mkDerivation {
           pname = "exo-pyo3-bindings";
           version = "0.1.0";
@@ -24,12 +22,6 @@
           # Install from pre-built wheel
           nativeBuildInputs = [ final.pyprojectWheelHook ];
           dontStrip = true;
-          passthru = prev.exo-pyo3-bindings.passthru or { };
-          postInstall = ''
-            local siteDir=$out/${final.python.sitePackages}/exo_pyo3_bindings
-            cp ${inputs.self}/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi $siteDir/
-            touch $siteDir/py.typed
-          '';
         };
       };
 
@@ -37,47 +29,17 @@
 
       # Overlay to provide build systems and custom packages
       buildSystemsOverlay = final: prev: {
+        # Use our pure Nix-built MLX with Metal support
+        mlx = self'.packages.mlx;
+
         # mlx-lm is a git dependency that needs setuptools
         mlx-lm = prev.mlx-lm.overrideAttrs (old: {
           nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
             final.setuptools
           ];
         });
-      } // lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin {
-        # Use our pure Nix-built MLX with Metal support (macOS only)
-        mlx = self'.packages.mlx;
       };
 
-      # Additional overlay for Linux-specific fixes (type checking env).
-      # Native wheels have shared lib dependencies we don't need at type-check time.
-      linuxOverlay = final: prev:
-        let
-          ignoreMissing = drv: drv.overrideAttrs { autoPatchelfIgnoreMissingDeps = [ "*" ]; };
-          nvidiaPackages = lib.filterAttrs (name: _: lib.hasPrefix "nvidia-" name) prev;
-        in
-        lib.optionalAttrs pkgs.stdenv.hostPlatform.isLinux (
-          (lib.mapAttrs (_: ignoreMissing) nvidiaPackages) // {
-            mlx = ignoreMissing prev.mlx;
-            mlx-cuda-13 = prev.mlx-cuda-13.overrideAttrs (old: {
-              buildInputs = (old.buildInputs or [ ]) ++ [
-                final.nvidia-cublas
-                final.nvidia-cuda-nvrtc
-                final.nvidia-cudnn-cu13
-                final.nvidia-nccl-cu13
-              ];
-              preFixup = ''
-                addAutoPatchelfSearchPath ${final.nvidia-cublas}
-                addAutoPatchelfSearchPath ${final.nvidia-cuda-nvrtc}
-                addAutoPatchelfSearchPath ${final.nvidia-cudnn-cu13}
-                addAutoPatchelfSearchPath ${final.nvidia-nccl-cu13}
-              '';
-              autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-            });
-            torch = ignoreMissing prev.torch;
-            triton = ignoreMissing prev.triton;
-          }
-        );
-
       pythonSet = (pkgs.callPackage inputs.pyproject-nix.build.packages {
         inherit python;
       }).overrideScope (
@@ -86,28 +48,16 @@
           overlay
           exoOverlay
           buildSystemsOverlay
-          linuxOverlay
         ]
       );
-      # mlx-cpu and mlx-cuda-13 both ship mlx/ site-packages files; keep first.
-      # mlx-cpu/mlx-cuda-13 and nvidia-cudnn-cu12/cu13 ship overlapping files.
-      venvCollisionPaths = lib.optionals pkgs.stdenv.hostPlatform.isLinux [
-        "lib/python3.13/site-packages/mlx*"
-        "lib/python3.13/site-packages/nvidia*"
-      ];
-
-      exoVenv = (pythonSet.mkVirtualEnv "exo-env" workspace.deps.default).overrideAttrs {
-        venvIgnoreCollisions = venvCollisionPaths;
-      };
+      exoVenv = pythonSet.mkVirtualEnv "exo-env" workspace.deps.default;
 
       # Virtual environment with dev dependencies for testing
-      testVenv = (pythonSet.mkVirtualEnv "exo-test-env" (
+      testVenv = pythonSet.mkVirtualEnv "exo-test-env" (
         workspace.deps.default // {
           exo = [ "dev" ]; # Include pytest, pytest-asyncio, pytest-env
         }
-      )).overrideAttrs {
-        venvIgnoreCollisions = venvCollisionPaths;
-      };
+      );
 
       mkPythonScript = name: path: pkgs.writeShellApplication {
         inherit name;
@@ -168,21 +118,6 @@
           ${pkgs.ruff}/bin/ruff check ${inputs.self}
           touch $out
         '';
-
-        # Hermetic basedpyright type checking
-        typecheck = pkgs.runCommand "typecheck"
-          {
-            nativeBuildInputs = [
-              testVenv
-              pkgs.basedpyright
-            ];
-          }
-          ''
-            cd ${inputs.self}
-            export HOME=$TMPDIR
-            basedpyright --pythonpath ${testVenv}/bin/python
-            touch $out
-          '';
       };
     };
 }

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

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "deepseek"
 quantization = "4bit"
 base_model = "DeepSeek V3.1"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 405874409472

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

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "deepseek"
 quantization = "8bit"
 base_model = "DeepSeek V3.1"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 765577920512

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

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "glm"
 quantization = "8bit"
 base_model = "GLM 4.5 Air"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 122406567936

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

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "glm"
 quantization = "bf16"
 base_model = "GLM 4.5 Air"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 229780750336

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

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "glm"
 quantization = "4bit"
 base_model = "GLM 4.7"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 198556925568

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

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "glm"
 quantization = "6bit"
 base_model = "GLM 4.7"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 286737579648

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

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "glm"
 quantization = "8bit"
 base_model = "GLM 4.7"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 396963397248

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

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "glm"
 quantization = "4bit"
 base_model = "GLM 4.7 Flash"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 19327352832

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

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "glm"
 quantization = "5bit"
 base_model = "GLM 4.7 Flash"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 22548578304

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

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "glm"
 quantization = "6bit"
 base_model = "GLM 4.7 Flash"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 26843545600

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

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "glm"
 quantization = "8bit"
 base_model = "GLM 4.7 Flash"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 34359738368

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

@@ -1,12 +0,0 @@
-model_id = "mlx-community/GLM-5-8bit-MXFP8"
-n_layers = 78
-hidden_size = 6144
-supports_tensor = true
-tasks = ["TextGeneration"]
-family = "glm"
-quantization = "8bit"
-base_model = "GLM-5"
-capabilities = ["text", "thinking"]
-
-[storage_size]
-in_bytes = 790517400864

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

@@ -1,12 +0,0 @@
-model_id = "mlx-community/GLM-5-MXFP4-Q8"
-n_layers = 78
-hidden_size = 6144
-supports_tensor = true
-tasks = ["TextGeneration"]
-family = "glm"
-quantization = "MXFP4-Q8"
-base_model = "GLM-5"
-capabilities = ["text", "thinking"]
-
-[storage_size]
-in_bytes = 405478939008

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

@@ -1,12 +0,0 @@
-model_id = "mlx-community/GLM-5"
-n_layers = 78
-hidden_size = 6144
-supports_tensor = true
-tasks = ["TextGeneration"]
-family = "glm"
-quantization = "bf16"
-base_model = "GLM-5"
-capabilities = ["text", "thinking"]
-
-[storage_size]
-in_bytes = 1487822475264

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

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "kimi"
 quantization = ""
 base_model = "Kimi K2"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 706522120192

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

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "kimi"
 quantization = ""
 base_model = "Kimi K2.5"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 662498705408

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

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "minimax"
 quantization = "3bit"
 base_model = "MiniMax M2.1"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 100086644736

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

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "minimax"
 quantization = "8bit"
 base_model = "MiniMax M2.1"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 242986745856

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

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3 0.6B"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 342884352

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

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3 0.6B"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 698351616

resources/inference_model_cards/mlx-community--Qwen3-235B-A22B-Instruct-2507-4bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3 235B"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 141733920768

resources/inference_model_cards/mlx-community--Qwen3-235B-A22B-Instruct-2507-8bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3 235B"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 268435456000

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

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3 30B"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 17612931072

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

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3 30B"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 33279705088

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

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3 Next 80B"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 47080074240

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

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3 Next 80B"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 88814387200

resources/inference_model_cards/mlx-community--Step-3.5-Flash-4bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "step"
 quantization = "4bit"
 base_model = "Step 3.5 Flash"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 114572190076

resources/inference_model_cards/mlx-community--Step-3.5-Flash-6bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "step"
 quantization = "6bit"
 base_model = "Step 3.5 Flash"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 159039627774

resources/inference_model_cards/mlx-community--Step-3.5-Flash-8Bit.toml

@@ -6,7 +6,7 @@ tasks = ["TextGeneration"]
 family = "step"
 quantization = "8bit"
 base_model = "Step 3.5 Flash"
-capabilities = ["text", "thinking", "thinking_toggle"]
+capabilities = ["text", "thinking"]
 
 [storage_size]
 in_bytes = 209082699847

rust/exo_pyo3_bindings/Cargo.toml

@@ -25,17 +25,17 @@ workspace = true
 networking = { workspace = true }
 
 # interop
-pyo3 = { version = "0.27.2", features = [
-    # "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13
-    # "nightly", # enables better-supported GIL integration
+pyo3 = { version = "0.27.1", features = [
+    # "abi3-py311", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.11
+    "nightly", # enables better-supported GIL integration
     "experimental-async", # async support in #[pyfunction] & #[pymethods]
     #"experimental-inspect", # inspection of generated binary => easier to automate type-hint generation
     #"py-clone", # adding Clone-ing of `Py<T>` without GIL (may cause panics - remove if panics happen)
-    # "multiple-pymethods", # allows multiple #[pymethods] sections per class
+    "multiple-pymethods", # allows multiple #[pymethods] sections per class
 
     # integrations with other libraries
-    # "arc_lock", "bigdecimal", "either", "hashbrown", "indexmap", "num-bigint", "num-complex", "num-rational",
-    # "ordered-float", "rust_decimal", "smallvec",
+    "arc_lock", "bigdecimal", "either", "hashbrown", "indexmap", "num-bigint", "num-complex", "num-rational",
+    "ordered-float", "rust_decimal", "smallvec",
     # "anyhow", "chrono", "chrono-local", "chrono-tz", "eyre", "jiff-02", "lock_api", "parking-lot", "time",  "serde",
 ] }
 pyo3-stub-gen = { version = "0.17.2" }
@@ -45,18 +45,33 @@ pyo3-log = "0.13.2"
 # macro dependencies
 extend = { workspace = true }
 delegate = { workspace = true }
+impl-trait-for-tuples = { workspace = true }
+derive_more = { workspace = true }
+pin-project = { workspace = true }
 
 # async runtime
 tokio = { workspace = true, features = ["full", "tracing"] }
-futures-lite = { workspace = true }
+futures = { workspace = true }
 
 # utility dependencies
+once_cell = "1.21.3"
+thread_local = "1.1.9"
 util = { workspace = true }
+thiserror = { workspace = true }
+#internment = { workspace = true }
+#recursion = { workspace = true }
+#generativity = { workspace = true }
+#itertools = { workspace = true }
+
 
 # Tracing
+#tracing = "0.1"
+#tracing-subscriber = "0.3"
+#console-subscriber = "0.1.5"
+#tracing-log = "0.2.0"
 log = { workspace = true }
 env_logger = "0.11"
 
+
 # Networking
 libp2p = { workspace = true, features = ["full"] }
-pin-project = "1.1.10"

rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi

@@ -19,7 +19,7 @@ class ConnectionUpdate:
         Whether this is a connection or disconnection event
         """
     @property
-    def peer_id(self) -> builtins.str:
+    def peer_id(self) -> PeerId:
         r"""
         Identity of the peer that we have connected to or disconnected from.
         """
@@ -40,22 +40,92 @@ class Keypair:
     Identity keypair of a node.
     """
     @staticmethod
-    def generate() -> Keypair:
+    def generate_ed25519() -> Keypair:
         r"""
         Generate a new Ed25519 keypair.
         """
     @staticmethod
-    def from_bytes(bytes: bytes) -> Keypair:
+    def generate_ecdsa() -> Keypair:
         r"""
-        Construct an Ed25519 keypair from secret key bytes
+        Generate a new ECDSA keypair.
+        """
+    @staticmethod
+    def generate_secp256k1() -> Keypair:
+        r"""
+        Generate a new Secp256k1 keypair.
+        """
+    @staticmethod
+    def from_protobuf_encoding(bytes: bytes) -> Keypair:
+        r"""
+        Decode a private key from a protobuf structure and parse it as a `Keypair`.
+        """
+    @staticmethod
+    def rsa_from_pkcs8(bytes: bytes) -> Keypair:
+        r"""
+        Decode an keypair from a DER-encoded secret key in PKCS#8 `PrivateKeyInfo`
+        format (i.e. unencrypted) as defined in [RFC5208].
+        
+        [RFC5208]: https://tools.ietf.org/html/rfc5208#section-5
+        """
+    @staticmethod
+    def secp256k1_from_der(bytes: bytes) -> Keypair:
+        r"""
+        Decode a keypair from a DER-encoded Secp256k1 secret key in an `ECPrivateKey`
+        structure as defined in [RFC5915].
+        
+        [RFC5915]: https://tools.ietf.org/html/rfc5915
+        """
+    @staticmethod
+    def ed25519_from_bytes(bytes: bytes) -> Keypair: ...
+    def to_protobuf_encoding(self) -> bytes:
+        r"""
+        Encode a private key as protobuf structure.
+        """
+    def to_peer_id(self) -> PeerId:
+        r"""
+        Convert the `Keypair` into the corresponding `PeerId`.
+        """
+
+@typing.final
+class Multiaddr:
+    r"""
+    Representation of a Multiaddr.
+    """
+    @staticmethod
+    def empty() -> Multiaddr:
+        r"""
+        Create a new, empty multiaddress.
+        """
+    @staticmethod
+    def with_capacity(n: builtins.int) -> Multiaddr:
+        r"""
+        Create a new, empty multiaddress with the given capacity.
+        """
+    @staticmethod
+    def from_bytes(bytes: bytes) -> Multiaddr:
+        r"""
+        Parse a `Multiaddr` value from its byte slice representation.
+        """
+    @staticmethod
+    def from_string(string: builtins.str) -> Multiaddr:
+        r"""
+        Parse a `Multiaddr` value from its string representation.
+        """
+    def len(self) -> builtins.int:
+        r"""
+        Return the length in bytes of this multiaddress.
+        """
+    def is_empty(self) -> builtins.bool:
+        r"""
+        Returns true if the length of this multiaddress is 0.
         """
     def to_bytes(self) -> bytes:
         r"""
-        Get the secret key bytes underlying the keypair
+        Return a copy of this [`Multiaddr`]'s byte representation.
         """
-    def to_node_id(self) -> builtins.str:
+    def to_string(self) -> builtins.str:
         r"""
-        Convert the `Keypair` into the corresponding `PeerId` string, which we use as our NodeId.
+        Convert a Multiaddr to a string.
         """
 
 @typing.final
@@ -110,6 +180,37 @@ class NoPeersSubscribedToTopicError(builtins.Exception):
     def __repr__(self) -> builtins.str: ...
     def __str__(self) -> builtins.str: ...
 
+@typing.final
+class PeerId:
+    r"""
+    Identifier of a peer of the network.
+    
+    The data is a `CIDv0` compatible multihash of the protobuf encoded public key of the peer
+    as specified in [specs/peer-ids](https://github.com/libp2p/specs/blob/master/peer-ids/peer-ids.md).
+    """
+    @staticmethod
+    def random() -> PeerId:
+        r"""
+        Generates a random peer ID from a cryptographically secure PRNG.
+        
+        This is useful for randomly walking on a DHT, or for testing purposes.
+        """
+    @staticmethod
+    def from_bytes(bytes: bytes) -> PeerId:
+        r"""
+        Parses a `PeerId` from bytes.
+        """
+    def to_bytes(self) -> bytes:
+        r"""
+        Returns a raw bytes representation of this `PeerId`.
+        """
+    def to_base58(self) -> builtins.str:
+        r"""
+        Returns a base-58 encoded string of this `PeerId`.
+        """
+    def __repr__(self) -> builtins.str: ...
+    def __str__(self) -> builtins.str: ...
+
 @typing.final
 class ConnectionUpdateType(enum.Enum):
     r"""

rust/exo_pyo3_bindings/src/allow_threading.rs

@@ -2,10 +2,11 @@
 //!
 
 use pin_project::pin_project;
+use pyo3::marker::Ungil;
 use pyo3::prelude::*;
 use std::{
     future::Future,
-    pin::Pin,
+    pin::{Pin, pin},
     task::{Context, Poll},
 };
 
@@ -25,13 +26,15 @@ where
 
 impl<F> Future for AllowThreads<F>
 where
-    F: Future + Send,
-    F::Output: Send,
+    F: Future + Ungil,
+    F::Output: Ungil,
 {
     type Output = F::Output;
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let waker = cx.waker();
-        Python::attach(|py| py.detach(|| self.project().0.poll(&mut Context::from_waker(waker))))
+        Python::with_gil(|py| {
+            py.allow_threads(|| self.project().0.poll(&mut Context::from_waker(waker)))
+        })
     }
 }

rust/exo_pyo3_bindings/src/examples/mod.rs

@@ -0,0 +1,240 @@
+//! This module exists to hold examples of some pyo3 patterns that may be too complex to
+//! re-create from scratch, but too inhomogenous to create an abstraction/wrapper around.
+//!
+//! Pattern examples include:
+//!  - Async task handles: with GC-integrated cleanup
+//!  - Sync/async callbacks from python: with propper eventloop handling
+//!
+//! Mutability pattern: https://pyo3.rs/v0.26.0/async-await.html#send--static-constraint
+//!  - Store mutable fields in tokio's `Mutex<T>`
+//!  - For async code: take `&self` and `.lock().await`
+//!  - For sync code: take `&mut self` and `.get_mut()`
+
+use crate::ext::{PyResultExt as _, ResultExt as _, TokioRuntimeExt as _};
+use futures::FutureExt as _;
+use futures::future::BoxFuture;
+use pyo3::exceptions::PyRuntimeError;
+use pyo3::prelude::{PyModule, PyModuleMethods as _};
+use pyo3::{
+    Bound, Py, PyAny, PyErr, PyResult, PyTraverseError, PyVisit, Python, pyclass, pymethods,
+};
+use std::time::Duration;
+use tokio::sync::mpsc;
+use tokio::sync::mpsc::error::TryRecvError;
+
+fn needs_tokio_runtime() {
+    tokio::runtime::Handle::current();
+}
+
+type SyncCallback = Box<dyn Fn() + Send + Sync>;
+type AsyncCallback = Box<dyn Fn() -> BoxFuture<'static, ()> + Send + Sync>;
+
+enum AsyncTaskMessage {
+    SyncCallback(SyncCallback),
+    AsyncCallback(AsyncCallback),
+}
+
+async fn async_task(
+    sender: mpsc::UnboundedSender<()>,
+    mut receiver: mpsc::UnboundedReceiver<AsyncTaskMessage>,
+) {
+    log::info!("RUST: async task started");
+
+    // task state
+    let mut interval = tokio::time::interval(Duration::from_secs(1));
+
+    let mut sync_cbs: Vec<SyncCallback> = vec![];
+    let mut async_cbs: Vec<AsyncCallback> = vec![];
+
+    loop {
+        tokio::select! {
+            // handle incoming messages from task-handle
+            message = receiver.recv() => {
+                // handle closed channel by exiting
+                let Some(message) = message else {
+                    log::info!("RUST: channel closed");
+                    break;
+                };
+
+                // dispatch incoming event
+                match message {
+                    AsyncTaskMessage::SyncCallback(cb) => {
+                        sync_cbs.push(cb);
+                    }
+                    AsyncTaskMessage::AsyncCallback(cb) => {
+                        async_cbs.push(cb);
+                    }
+                }
+            }
+
+            // handle all other events
+            _ = interval.tick() => {
+                log::info!("RUST: async task tick");
+
+                // call back all sync callbacks
+                for cb in &sync_cbs {
+                    cb();
+                }
+
+                // call back all async callbacks
+                for cb in &async_cbs {
+                    cb().await;
+                }
+
+                // send event on unbounded channel
+                sender.send(()).expect("handle receiver cannot be closed/dropped");
+            }
+        }
+    }
+
+    log::info!("RUST: async task stopped");
+}
+
+// #[gen_stub_pyclass]
+#[pyclass(name = "AsyncTaskHandle")]
+#[derive(Debug)]
+struct PyAsyncTaskHandle {
+    sender: Option<mpsc::UnboundedSender<AsyncTaskMessage>>,
+    receiver: mpsc::UnboundedReceiver<()>,
+}
+
+#[allow(clippy::expect_used)]
+impl PyAsyncTaskHandle {
+    const fn sender(&self) -> &mpsc::UnboundedSender<AsyncTaskMessage> {
+        self.sender
+            .as_ref()
+            .expect("The sender should only be None after de-initialization.")
+    }
+
+    const fn sender_mut(&mut self) -> &mpsc::UnboundedSender<AsyncTaskMessage> {
+        self.sender
+            .as_mut()
+            .expect("The sender should only be None after de-initialization.")
+    }
+
+    const fn new(
+        sender: mpsc::UnboundedSender<AsyncTaskMessage>,
+        receiver: mpsc::UnboundedReceiver<()>,
+    ) -> Self {
+        Self {
+            sender: Some(sender),
+            receiver,
+        }
+    }
+}
+
+// #[gen_stub_pymethods]
+#[pymethods]
+impl PyAsyncTaskHandle {
+    #[new]
+    fn py_new(py: Python<'_>) -> PyResult<Self> {
+        use pyo3_async_runtimes::tokio::get_runtime;
+
+        // create communication channel TOWARDS our task
+        let (h_sender, t_receiver) = mpsc::unbounded_channel::<AsyncTaskMessage>();
+
+        // create communication channel FROM our task
+        let (t_sender, h_receiver) = mpsc::unbounded_channel::<()>();
+
+        // perform necessary setup within tokio context - or it crashes
+        let () = get_runtime().block_on(async { needs_tokio_runtime() });
+
+        // spawn tokio task with this thread's task-locals - without this, async callbacks on the new threads will not work!!
+        _ = get_runtime().spawn_with_scope(py, async move {
+            async_task(t_sender, t_receiver).await;
+        });
+        Ok(Self::new(h_sender, h_receiver))
+    }
+
+    /// NOTE: exceptions in callbacks are silently ignored until end of execution
+    fn add_sync_callback(
+        &self,
+        // #[gen_stub(override_type(
+        //     type_repr="collections.abc.Callable[[], None]",
+        //     imports=("collections.abc")
+        // ))]
+        callback: Py<PyAny>,
+    ) -> PyResult<()> {
+        // blocking call to async method -> can do non-blocking if needed
+        self.sender()
+            .send(AsyncTaskMessage::SyncCallback(Box::new(move || {
+                _ = Python::with_gil(|py| callback.call0(py).write_unraisable_with(py));
+            })))
+            .pyerr()?;
+        Ok(())
+    }
+
+    /// NOTE: exceptions in callbacks are silently ignored until end of execution
+    fn add_async_callback(
+        &self,
+        // #[gen_stub(override_type(
+        //     type_repr="collections.abc.Callable[[], collections.abc.Awaitable[None]]",
+        //     imports=("collections.abc")
+        // ))]
+        callback: Py<PyAny>,
+    ) -> PyResult<()> {
+        // blocking call to async method -> can do non-blocking if needed
+        self.sender()
+            .send(AsyncTaskMessage::AsyncCallback(Box::new(move || {
+                let c = Python::with_gil(|py| callback.clone_ref(py));
+                async move {
+                    if let Some(f) = Python::with_gil(|py| {
+                        let coroutine = c.call0(py).write_unraisable_with(py)?;
+                        pyo3_async_runtimes::tokio::into_future(coroutine.into_bound(py))
+                            .write_unraisable_with(py)
+                    }) {
+                        _ = f.await.write_unraisable();
+                    }
+                }
+                .boxed()
+            })))
+            .pyerr()?;
+        Ok(())
+    }
+
+    async fn receive_unit(&mut self) -> PyResult<()> {
+        self.receiver
+            .recv()
+            .await
+            .ok_or(PyErr::new::<PyRuntimeError, _>(
+                "cannot receive unit on closed channel",
+            ))
+    }
+
+    fn drain_units(&mut self) -> PyResult<i32> {
+        let mut cnt = 0;
+        loop {
+            match self.receiver.try_recv() {
+                Err(TryRecvError::Disconnected) => {
+                    return Err(PyErr::new::<PyRuntimeError, _>(
+                        "cannot receive unit on closed channel",
+                    ));
+                }
+                Err(TryRecvError::Empty) => return Ok(cnt),
+                Ok(()) => {
+                    cnt += 1;
+                    continue;
+                }
+            }
+        }
+    }
+
+    // #[gen_stub(skip)]
+    const fn __traverse__(&self, _visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
+        Ok(()) // This is needed purely so `__clear__` can work
+    }
+
+    // #[gen_stub(skip)]
+    fn __clear__(&mut self) {
+        // TODO: may or may not need to await a "kill-signal" oneshot channel message,
+        //       to ensure that the networking task is done BEFORE exiting the clear function...
+        //       but this may require GIL?? and it may not be safe to call GIL here??
+        self.sender = None; // Using Option<T> as a trick to force `sender` channel to be dropped
+    }
+}
+
+pub fn examples_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<PyAsyncTaskHandle>()?;
+
+    Ok(())
+}

rust/exo_pyo3_bindings/src/ident.rs

@@ -1,47 +0,0 @@
-use crate::ext::ResultExt as _;
-use libp2p::identity::Keypair;
-use pyo3::types::{PyBytes, PyBytesMethods};
-use pyo3::{Bound, PyResult, Python, pyclass, pymethods};
-use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
-
-/// Identity keypair of a node.
-#[gen_stub_pyclass]
-#[pyclass(name = "Keypair", frozen)]
-#[repr(transparent)]
-pub struct PyKeypair(pub Keypair);
-
-#[gen_stub_pymethods]
-#[pymethods]
-#[allow(clippy::needless_pass_by_value)]
-impl PyKeypair {
-    /// Generate a new Ed25519 keypair.
-    #[staticmethod]
-    fn generate() -> Self {
-        Self(Keypair::generate_ed25519())
-    }
-
-    /// Construct an Ed25519 keypair from secret key bytes
-    #[staticmethod]
-    fn from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
-        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
-            .clone()
-            .try_into_ed25519()
-            .expect("we only use ed25519 keys")
-            .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 {
-        self.0.public().to_peer_id().to_base58()
-    }
-}

rust/exo_pyo3_bindings/src/lib.rs

@@ -4,14 +4,28 @@
 //!
 //!
 
-mod allow_threading;
-mod ident;
-mod networking;
+// enable Rust-unstable features for convenience
+#![feature(trait_alias)]
+#![feature(tuple_trait)]
+#![feature(unboxed_closures)]
+// #![feature(stmt_expr_attributes)]
+// #![feature(assert_matches)]
+// #![feature(async_fn_in_dyn_trait)]
+// #![feature(async_for_loop)]
+// #![feature(auto_traits)]
+// #![feature(negative_impls)]
+
+extern crate core;
+mod allow_threading;
+mod examples;
+pub(crate) mod networking;
+pub(crate) mod pylibp2p;
 
-use crate::ident::PyKeypair;
 use crate::networking::networking_submodule;
+use crate::pylibp2p::ident::ident_submodule;
+use crate::pylibp2p::multiaddr::multiaddr_submodule;
 use pyo3::prelude::PyModule;
-use pyo3::types::PyModuleMethods;
+use pyo3::prelude::*;
 use pyo3::{Bound, PyResult, pyclass, pymodule};
 use pyo3_stub_gen::define_stub_info_gatherer;
 
@@ -20,11 +34,24 @@ pub(crate) mod r#const {
     pub const MPSC_CHANNEL_SIZE: usize = 1024;
 }
 
+/// Namespace for all the type/trait aliases used by this crate.
+pub(crate) mod alias {
+    use std::error::Error;
+    use std::marker::Tuple;
+
+    pub trait SendFn<Args: Tuple + Send + 'static, Output> =
+        Fn<Args, Output = Output> + Send + 'static;
+
+    pub type AnyError = Box<dyn Error + Send + Sync + 'static>;
+    pub type AnyResult<T> = Result<T, AnyError>;
+}
+
 /// Namespace for crate-wide extension traits/methods
 pub(crate) mod ext {
     use crate::allow_threading::AllowThreads;
     use extend::ext;
     use pyo3::exceptions::{PyConnectionError, PyRuntimeError};
+    use pyo3::marker::Ungil;
     use pyo3::types::PyBytes;
     use pyo3::{Py, PyErr, PyResult, Python};
     use tokio::runtime::Runtime;
@@ -35,7 +62,7 @@ pub(crate) mod ext {
     #[ext(pub, name = ByteArrayExt)]
     impl [u8] {
         fn pybytes(&self) -> Py<PyBytes> {
-            Python::attach(|py| PyBytes::new(py, self).unbind())
+            Python::with_gil(|py| PyBytes::new(py, self).unbind())
         }
     }
 
@@ -71,7 +98,7 @@ pub(crate) mod ext {
     #[ext(pub, name = PyResultExt)]
     impl<T> PyResult<T> {
         fn write_unraisable(self) -> Option<T> {
-            Python::attach(|py| self.write_unraisable_with(py))
+            Python::with_gil(|py| self.write_unraisable_with(py))
         }
 
         fn write_unraisable_with(self, py: Python<'_>) -> Option<T> {
@@ -148,6 +175,24 @@ pub(crate) mod ext {
     }
 }
 
+pub(crate) mod private {
+    use std::marker::Sized;
+
+    /// Sealed traits support
+    pub trait Sealed {}
+    impl<T: ?Sized> Sealed for T {}
+}
+
+/// A wrapper around [`Py`] that implements [`Clone`] using [`Python::with_gil`].
+#[repr(transparent)]
+pub(crate) struct ClonePy<T>(pub Py<T>);
+
+impl<T> Clone for ClonePy<T> {
+    fn clone(&self) -> Self {
+        Python::with_gil(|py| Self(self.0.clone_ref(py)))
+    }
+}
+
 /// A Python module implemented in Rust. The name of this function must match
 /// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to
 /// import the module.
@@ -159,7 +204,8 @@ fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     // TODO: for now this is all NOT a submodule, but figure out how to make the submodule system
     //       work with maturin, where the types generate correctly, in the right folder, without
     //       too many importing issues...
-    m.add_class::<PyKeypair>()?;
+    ident_submodule(m)?;
+    multiaddr_submodule(m)?;
     networking_submodule(m)?;
 
     // top-level constructs

rust/exo_pyo3_bindings/src/networking.rs

@@ -8,12 +8,12 @@
 use crate::r#const::MPSC_CHANNEL_SIZE;
 use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
 use crate::ext::{ResultExt as _, TokioMpscReceiverExt as _, TokioMpscSenderExt as _};
-use crate::ident::PyKeypair;
 use crate::pyclass;
+use crate::pylibp2p::ident::{PyKeypair, PyPeerId};
 use libp2p::futures::StreamExt as _;
-use libp2p::gossipsub;
 use libp2p::gossipsub::{IdentTopic, Message, MessageId, PublishError};
 use libp2p::swarm::SwarmEvent;
+use libp2p::{gossipsub, mdns};
 use networking::discovery;
 use networking::swarm::create_swarm;
 use pyo3::prelude::{PyModule, PyModuleMethods as _};
@@ -25,7 +25,7 @@ use tokio::sync::{Mutex, mpsc, oneshot};
 
 mod exception {
     use pyo3::types::PyTuple;
-    use pyo3::{exceptions::PyException, prelude::*};
+    use pyo3::{PyErrArguments, exceptions::PyException, prelude::*};
     use pyo3_stub_gen::derive::*;
 
     #[gen_stub_pyclass]
@@ -119,7 +119,7 @@ struct PyConnectionUpdate {
 
     /// Identity of the peer that we have connected to or disconnected from.
     #[pyo3(get)]
-    peer_id: String,
+    peer_id: PyPeerId,
 
     /// Remote connection's IPv4 address.
     #[pyo3(get)]
@@ -155,6 +155,7 @@ async fn networking_task(
 ) {
     use SwarmEvent::*;
     use ToTask::*;
+    use mdns::Event::*;
     use networking::swarm::BehaviourEvent::*;
 
     log::info!("RUST: networking task started");
@@ -251,7 +252,7 @@ async fn networking_task(
                         // send connection event to channel (or exit if connection closed)
                         if let Err(e) = connection_update_tx.send(PyConnectionUpdate {
                             update_type: PyConnectionUpdateType::Connected,
-                            peer_id: peer_id.to_base58(),
+                            peer_id: PyPeerId(peer_id),
                             remote_ipv4,
                             remote_tcp_port,
                         }).await {
@@ -272,7 +273,7 @@ async fn networking_task(
                         // send disconnection event to channel (or exit if connection closed)
                         if let Err(e) = connection_update_tx.send(PyConnectionUpdate {
                             update_type: PyConnectionUpdateType::Disconnected,
-                            peer_id: peer_id.to_base58(),
+                            peer_id: PyPeerId(peer_id),
                             remote_ipv4,
                             remote_tcp_port,
                         }).await {
@@ -484,7 +485,7 @@ impl PyNetworkingHandle {
         let (tx, rx) = oneshot::channel();
 
         // send off request to subscribe
-        let data = Python::attach(|py| Vec::from(data.as_bytes(py)));
+        let data = Python::with_gil(|py| Vec::from(data.as_bytes(py)));
         self.to_task_tx()
             .send_py(ToTask::GossipsubPublish {
                 topic,

rust/exo_pyo3_bindings/src/pylibp2p/ident.rs

@@ -0,0 +1,159 @@
+use crate::ext::ResultExt as _;
+use libp2p::PeerId;
+use libp2p::identity::Keypair;
+use pyo3::prelude::{PyBytesMethods as _, PyModule, PyModuleMethods as _};
+use pyo3::types::PyBytes;
+use pyo3::{Bound, PyResult, Python, pyclass, pymethods};
+use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
+
+/// Identity keypair of a node.
+#[gen_stub_pyclass]
+#[pyclass(name = "Keypair", frozen)]
+#[repr(transparent)]
+pub struct PyKeypair(pub Keypair);
+
+#[gen_stub_pymethods]
+#[pymethods]
+#[allow(clippy::needless_pass_by_value)]
+impl PyKeypair {
+    /// Generate a new Ed25519 keypair.
+    #[staticmethod]
+    fn generate_ed25519() -> Self {
+        Self(Keypair::generate_ed25519())
+    }
+
+    /// Generate a new ECDSA keypair.
+    #[staticmethod]
+    fn generate_ecdsa() -> Self {
+        Self(Keypair::generate_ecdsa())
+    }
+
+    /// Generate a new Secp256k1 keypair.
+    #[staticmethod]
+    fn generate_secp256k1() -> Self {
+        Self(Keypair::generate_secp256k1())
+    }
+
+    /// Decode a private key from a protobuf structure and parse it as a `Keypair`.
+    #[staticmethod]
+    fn from_protobuf_encoding(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
+        let bytes = Vec::from(bytes.as_bytes());
+        Ok(Self(Keypair::from_protobuf_encoding(&bytes).pyerr()?))
+    }
+
+    /// Decode an keypair from a DER-encoded secret key in PKCS#8 `PrivateKeyInfo`
+    /// format (i.e. unencrypted) as defined in [RFC5208].
+    ///
+    /// [RFC5208]: https://tools.ietf.org/html/rfc5208#section-5
+    #[staticmethod]
+    fn rsa_from_pkcs8(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
+        let mut bytes = Vec::from(bytes.as_bytes());
+        Ok(Self(Keypair::rsa_from_pkcs8(&mut bytes).pyerr()?))
+    }
+
+    /// Decode a keypair from a DER-encoded Secp256k1 secret key in an `ECPrivateKey`
+    /// structure as defined in [RFC5915].
+    ///
+    /// [RFC5915]: https://tools.ietf.org/html/rfc5915
+    #[staticmethod]
+    fn secp256k1_from_der(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
+        let mut bytes = Vec::from(bytes.as_bytes());
+        Ok(Self(Keypair::secp256k1_from_der(&mut bytes).pyerr()?))
+    }
+
+    #[staticmethod]
+    fn ed25519_from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
+        let mut bytes = Vec::from(bytes.as_bytes());
+        Ok(Self(Keypair::ed25519_from_bytes(&mut bytes).pyerr()?))
+    }
+
+    /// Encode a private key as protobuf structure.
+    fn to_protobuf_encoding<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
+        let bytes = self.0.to_protobuf_encoding().pyerr()?;
+        Ok(PyBytes::new(py, &bytes))
+    }
+
+    /// Convert the `Keypair` into the corresponding `PeerId`.
+    fn to_peer_id(&self) -> PyPeerId {
+        PyPeerId(self.0.public().to_peer_id())
+    }
+
+    // /// Hidden constructor for pickling support. TODO: figure out how to do pickling...
+    // #[gen_stub(skip)]
+    // #[new]
+    // fn py_new(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
+    //     Self::from_protobuf_encoding(bytes)
+    // }
+    //
+    // #[gen_stub(skip)]
+    // fn __setstate__(&mut self, state: Bound<'_, PyBytes>) -> PyResult<()> {
+    //     *self = Self::from_protobuf_encoding(state)?;
+    //     Ok(())
+    // }
+    //
+    // #[gen_stub(skip)]
+    // fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
+    //     self.to_protobuf_encoding(py)
+    // }
+    //
+    // #[gen_stub(skip)]
+    // pub fn __getnewargs__<'py>(&self, py: Python<'py>) -> PyResult<(Bound<'py, PyBytes>,)> {
+    //     Ok((self.to_protobuf_encoding(py)?,))
+    // }
+}
+
+/// Identifier of a peer of the network.
+///
+/// The data is a `CIDv0` compatible multihash of the protobuf encoded public key of the peer
+/// as specified in [specs/peer-ids](https://github.com/libp2p/specs/blob/master/peer-ids/peer-ids.md).
+#[gen_stub_pyclass]
+#[pyclass(name = "PeerId", frozen)]
+#[derive(Debug, Clone)]
+#[repr(transparent)]
+pub struct PyPeerId(pub PeerId);
+
+#[gen_stub_pymethods]
+#[pymethods]
+#[allow(clippy::needless_pass_by_value)]
+impl PyPeerId {
+    /// Generates a random peer ID from a cryptographically secure PRNG.
+    ///
+    /// This is useful for randomly walking on a DHT, or for testing purposes.
+    #[staticmethod]
+    fn random() -> Self {
+        Self(PeerId::random())
+    }
+
+    /// Parses a `PeerId` from bytes.
+    #[staticmethod]
+    fn from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
+        let bytes = Vec::from(bytes.as_bytes());
+        Ok(Self(PeerId::from_bytes(&bytes).pyerr()?))
+    }
+
+    /// Returns a raw bytes representation of this `PeerId`.
+    fn to_bytes<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
+        let bytes = self.0.to_bytes();
+        PyBytes::new(py, &bytes)
+    }
+
+    /// Returns a base-58 encoded string of this `PeerId`.
+    fn to_base58(&self) -> String {
+        self.0.to_base58()
+    }
+
+    fn __repr__(&self) -> String {
+        format!("PeerId({})", self.to_base58())
+    }
+
+    fn __str__(&self) -> String {
+        self.to_base58()
+    }
+}
+
+pub fn ident_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<PyKeypair>()?;
+    m.add_class::<PyPeerId>()?;
+
+    Ok(())
+}

rust/exo_pyo3_bindings/src/pylibp2p/mod.rs

@@ -0,0 +1,8 @@
+//! A module for exposing Rust's libp2p datatypes over Pyo3
+//!
+//! TODO: right now we are coupled to libp2p's identity, but eventually we want to create our own
+//!       independent identity type of some kind or another. This may require handshaking.
+//!
+
+pub mod ident;
+pub mod multiaddr;

rust/exo_pyo3_bindings/src/pylibp2p/multiaddr.rs

@@ -0,0 +1,81 @@
+use crate::ext::ResultExt as _;
+use libp2p::Multiaddr;
+use pyo3::prelude::{PyBytesMethods as _, PyModule, PyModuleMethods as _};
+use pyo3::types::PyBytes;
+use pyo3::{Bound, PyResult, Python, pyclass, pymethods};
+use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
+use std::str::FromStr as _;
+
+/// Representation of a Multiaddr.
+#[gen_stub_pyclass]
+#[pyclass(name = "Multiaddr", frozen)]
+#[derive(Debug, Clone)]
+#[repr(transparent)]
+pub struct PyMultiaddr(pub Multiaddr);
+
+#[gen_stub_pymethods]
+#[pymethods]
+#[allow(clippy::needless_pass_by_value)]
+impl PyMultiaddr {
+    /// Create a new, empty multiaddress.
+    #[staticmethod]
+    fn empty() -> Self {
+        Self(Multiaddr::empty())
+    }
+
+    /// Create a new, empty multiaddress with the given capacity.
+    #[staticmethod]
+    fn with_capacity(n: usize) -> Self {
+        Self(Multiaddr::with_capacity(n))
+    }
+
+    /// Parse a `Multiaddr` value from its byte slice representation.
+    #[staticmethod]
+    fn from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
+        let bytes = Vec::from(bytes.as_bytes());
+        Ok(Self(Multiaddr::try_from(bytes).pyerr()?))
+    }
+
+    /// Parse a `Multiaddr` value from its string representation.
+    #[staticmethod]
+    fn from_string(string: String) -> PyResult<Self> {
+        Ok(Self(Multiaddr::from_str(&string).pyerr()?))
+    }
+
+    /// Return the length in bytes of this multiaddress.
+    fn len(&self) -> usize {
+        self.0.len()
+    }
+
+    /// Returns true if the length of this multiaddress is 0.
+    fn is_empty(&self) -> bool {
+        self.0.is_empty()
+    }
+
+    /// Return a copy of this [`Multiaddr`]'s byte representation.
+    fn to_bytes<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
+        let bytes = self.0.to_vec();
+        PyBytes::new(py, &bytes)
+    }
+
+    /// Convert a Multiaddr to a string.
+    fn to_string(&self) -> String {
+        self.0.to_string()
+    }
+
+    #[gen_stub(skip)]
+    fn __repr__(&self) -> String {
+        format!("Multiaddr({})", self.0)
+    }
+
+    #[gen_stub(skip)]
+    fn __str__(&self) -> String {
+        self.to_string()
+    }
+}
+
+pub fn multiaddr_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<PyMultiaddr>()?;
+
+    Ok(())
+}

rust/networking/Cargo.toml

@@ -19,14 +19,21 @@ either = { workspace = true }
 # macro dependencies
 extend = { workspace = true }
 delegate = { workspace = true }
+impl-trait-for-tuples = { workspace = true }
+derive_more = { workspace = true }
 
 # async
 tokio = { workspace = true, features = ["full"] }
-futures-lite = { workspace = true }
+futures = { workspace = true }
 futures-timer = { workspace = true }
 
 # utility dependencies
 util = { workspace = true }
+thiserror = { workspace = true }
+#internment = { workspace = true }
+#recursion = { workspace = true }
+#generativity = { workspace = true }
+#itertools = { workspace = true }
 tracing-subscriber = { version = "0.3.19", features = ["default", "env-filter"] }
 keccak-const = { workspace = true }
 
@@ -34,4 +41,4 @@ keccak-const = { workspace = true }
 log = { workspace = true }
 
 # networking
-libp2p = { workspace = true, features = ["full"] }
+libp2p = { workspace = true, features = ["full"] }

rust/networking/examples/chatroom.rs

@@ -1,4 +1,4 @@
-use futures_lite::StreamExt;
+use futures::stream::StreamExt as _;
 use libp2p::{gossipsub, identity, swarm::SwarmEvent};
 use networking::{discovery, swarm};
 use tokio::{io, io::AsyncBufReadExt as _, select};
@@ -38,19 +38,19 @@ async fn main() {
                     println!("Publish error: {e:?}");
                 }
             }
-            event = swarm.next() => match event {
+            event = swarm.select_next_some() => match event {
                 // on gossipsub incoming
-                Some(SwarmEvent::Behaviour(swarm::BehaviourEvent::Gossipsub(gossipsub::Event::Message {
+                SwarmEvent::Behaviour(swarm::BehaviourEvent::Gossipsub(gossipsub::Event::Message {
                     propagation_source: peer_id,
                     message_id: id,
                     message,
-                }))) => println!(
+                })) => println!(
                         "\n\nGot message: '{}' with id: {id} from peer: {peer_id}\n\n",
                         String::from_utf8_lossy(&message.data),
                     ),
 
                 // on discovery
-                Some(SwarmEvent::Behaviour(swarm::BehaviourEvent::Discovery(e)) )=> match e {
+                SwarmEvent::Behaviour(swarm::BehaviourEvent::Discovery(e)) => match e {
                     discovery::Event::ConnectionEstablished {
                         peer_id, connection_id, remote_ip, remote_tcp_port
                     } => {
@@ -64,7 +64,7 @@ async fn main() {
                 }
 
                 // ignore outgoing errors: those are normal
-                e@Some(SwarmEvent::OutgoingConnectionError { .. }) => { log::debug!("Outgoing connection error: {e:?}"); }
+                e@SwarmEvent::OutgoingConnectionError { .. } => { log::debug!("Outgoing connection error: {e:?}"); }
 
                 // otherwise log any other event
                 e => { log::info!("Other event {e:?}"); }

rust/networking/examples/chatroom_manual.rs

@@ -0,0 +1,127 @@
+// Copyright 2018 Parity Technologies (UK) Ltd.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a
+// copy of this software and associated documentation files (the "Software"),
+// to deal in the Software without restriction, including without limitation
+// the rights to use, copy, modify, merge, publish, distribute, sublicense,
+// and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+// DEALINGS IN THE SOFTWARE.
+
+use futures::stream::StreamExt;
+use libp2p::{
+    gossipsub, mdns, noise,
+    swarm::{NetworkBehaviour, SwarmEvent},
+    tcp, yamux,
+};
+use std::time::Duration;
+use std::{error::Error, hash::Hash};
+use tokio::{io, io::AsyncBufReadExt, select};
+use tracing_subscriber::EnvFilter;
+
+// We create a custom network behaviour that combines Gossipsub and Mdns.
+#[derive(NetworkBehaviour)]
+struct MyBehaviour {
+    gossipsub: gossipsub::Behaviour,
+    mdns: mdns::tokio::Behaviour,
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn Error>> {
+    let _ = tracing_subscriber::fmt()
+        .with_env_filter(EnvFilter::from_default_env())
+        .try_init();
+
+    let mut swarm = libp2p::SwarmBuilder::with_new_identity()
+        .with_tokio()
+        .with_tcp(
+            tcp::Config::default(),
+            noise::Config::new,
+            yamux::Config::default,
+        )?
+        .with_behaviour(|key| {
+            // Set a custom gossipsub configuration
+            let gossipsub_config = gossipsub::ConfigBuilder::default()
+                .heartbeat_interval(Duration::from_secs(10))
+                .validation_mode(gossipsub::ValidationMode::Strict) // This sets the kind of message validation. The default is Strict (enforce message signing)
+                .build()
+                .map_err(io::Error::other)?; // Temporary hack because `build` does not return a proper `std::error::Error`.
+
+            // build a gossipsub network behaviour
+            let gossipsub = gossipsub::Behaviour::new(
+                gossipsub::MessageAuthenticity::Signed(key.clone()),
+                gossipsub_config,
+            )?;
+
+            let mdns =
+                mdns::tokio::Behaviour::new(mdns::Config::default(), key.public().to_peer_id())?;
+            Ok(MyBehaviour { gossipsub, mdns })
+        })?
+        .build();
+
+    println!("Running swarm with identity {}", swarm.local_peer_id());
+
+    // Create a Gossipsub topic
+    let topic = gossipsub::IdentTopic::new("test-net");
+    // subscribes to our topic
+    swarm.behaviour_mut().gossipsub.subscribe(&topic)?;
+
+    // Read full lines from stdin
+    let mut stdin = io::BufReader::new(io::stdin()).lines();
+
+    // Listen on all interfaces and whatever port the OS assigns
+    swarm.listen_on("/ip4/0.0.0.0/tcp/0".parse()?)?;
+
+    println!("Enter messages via STDIN and they will be sent to connected peers using Gossipsub");
+
+    // Kick it off
+    loop {
+        select! {
+            Ok(Some(line)) = stdin.next_line() => {
+                if let Err(e) = swarm
+                    .behaviour_mut().gossipsub
+                    .publish(topic.clone(), line.as_bytes()) {
+                    println!("Publish error: {e:?}");
+                }
+            }
+            event = swarm.select_next_some() => match event {
+                SwarmEvent::Behaviour(MyBehaviourEvent::Mdns(mdns::Event::Discovered(list))) => {
+                    for (peer_id, multiaddr) in list {
+                        println!("mDNS discovered a new peer: {peer_id} on {multiaddr}");
+                        swarm.behaviour_mut().gossipsub.add_explicit_peer(&peer_id);
+                    }
+                },
+                SwarmEvent::Behaviour(MyBehaviourEvent::Mdns(mdns::Event::Expired(list))) => {
+                    for (peer_id, multiaddr) in list {
+                        println!("mDNS discover peer has expired: {peer_id} on {multiaddr}");
+                        swarm.behaviour_mut().gossipsub.remove_explicit_peer(&peer_id);
+                    }
+                },
+                SwarmEvent::Behaviour(MyBehaviourEvent::Gossipsub(gossipsub::Event::Message {
+                    propagation_source: peer_id,
+                    message_id: id,
+                    message,
+                })) => println!(
+                        "Got message: '{}' with id: {id} from peer: {peer_id}",
+                        String::from_utf8_lossy(&message.data),
+                    ),
+                SwarmEvent::NewListenAddr { address, .. } => {
+                    println!("Local node is listening on {address}");
+                }
+                e => {
+                    println!("Other swarm event: {:?}", e);
+                }
+            }
+        }
+    }
+}

rust/networking/src/discovery.rs

@@ -1,7 +1,8 @@
 use crate::ext::MultiaddrExt;
+use crate::keep_alive;
 use delegate::delegate;
 use either::Either;
-use futures_lite::FutureExt;
+use futures::FutureExt;
 use futures_timer::Delay;
 use libp2p::core::transport::PortUse;
 use libp2p::core::{ConnectedPoint, Endpoint};
@@ -362,7 +363,7 @@ impl NetworkBehaviour for Behaviour {
         }
 
         // retry connecting to all mDNS peers periodically (fails safely if already connected)
-        if self.retry_delay.poll(cx).is_ready() {
+        if self.retry_delay.poll_unpin(cx).is_ready() {
             for (p, mas) in self.mdns_discovered.clone() {
                 for ma in mas {
                     self.dial(p, ma)

rust/networking/src/keep_alive.rs

@@ -0,0 +1,44 @@
+use delegate::delegate;
+use libp2p::swarm::handler::ConnectionEvent;
+use libp2p::swarm::{ConnectionHandlerEvent, SubstreamProtocol, dummy, handler};
+use std::task::{Context, Poll};
+
+/// An implementation of [`ConnectionHandler`] that doesn't handle any protocols, but it keeps
+/// the connection alive.
+#[derive(Clone)]
+#[repr(transparent)]
+pub struct ConnectionHandler(dummy::ConnectionHandler);
+
+impl ConnectionHandler {
+    pub fn new() -> Self {
+        ConnectionHandler(dummy::ConnectionHandler)
+    }
+}
+
+impl handler::ConnectionHandler for ConnectionHandler {
+    // delegate types and implementation mostly to dummy handler
+    type FromBehaviour = <dummy::ConnectionHandler as handler::ConnectionHandler>::FromBehaviour;
+    type ToBehaviour = <dummy::ConnectionHandler as handler::ConnectionHandler>::ToBehaviour;
+    type InboundProtocol =
+        <dummy::ConnectionHandler as handler::ConnectionHandler>::InboundProtocol;
+    type OutboundProtocol =
+        <dummy::ConnectionHandler as handler::ConnectionHandler>::OutboundProtocol;
+    type InboundOpenInfo =
+        <dummy::ConnectionHandler as handler::ConnectionHandler>::InboundOpenInfo;
+    type OutboundOpenInfo =
+        <dummy::ConnectionHandler as handler::ConnectionHandler>::OutboundOpenInfo;
+
+    delegate! {
+        to self.0 {
+            fn listen_protocol(&self) -> SubstreamProtocol<Self::InboundProtocol, Self::InboundOpenInfo>;
+            fn poll(&mut self, cx: &mut Context<'_>) -> Poll<ConnectionHandlerEvent<Self::OutboundProtocol, Self::OutboundOpenInfo, Self::ToBehaviour>>;
+            fn on_behaviour_event(&mut self, event: Self::FromBehaviour);
+            fn on_connection_event(&mut self, event: ConnectionEvent<Self::InboundProtocol, Self::OutboundProtocol, Self::InboundOpenInfo, Self::OutboundOpenInfo>);
+        }
+    }
+
+    // specifically override this to force connection to stay alive
+    fn connection_keep_alive(&self) -> bool {
+        true
+    }
+}

rust/networking/src/lib.rs

@@ -3,7 +3,19 @@
 //! this is here as a placeholder documentation
 //!
 //!
+
+// enable Rust-unstable features for convenience
+#![feature(trait_alias)]
+// #![feature(stmt_expr_attributes)]
+// #![feature(unboxed_closures)]
+// #![feature(assert_matches)]
+// #![feature(async_fn_in_dyn_trait)]
+// #![feature(async_for_loop)]
+// #![feature(auto_traits)]
+// #![feature(negative_impls)]
+
 pub mod discovery;
+pub mod keep_alive;
 pub mod swarm;
 
 /// Namespace for all the type/trait aliases used by this crate.
@@ -42,3 +54,11 @@ pub(crate) mod ext {
         }
     }
 }
+
+pub(crate) mod private {
+    #![allow(dead_code)]
+
+    /// Sealed traits support
+    pub trait Sealed {}
+    impl<T: ?Sized> Sealed for T {}
+}

rust/networking/src/swarm.rs

@@ -31,7 +31,7 @@ pub fn create_swarm(keypair: identity::Keypair) -> alias::AnyResult<Swarm> {
 mod transport {
     use crate::alias;
     use crate::swarm::{NETWORK_VERSION, OVERRIDE_VERSION_ENV_VAR};
-    use futures_lite::{AsyncRead, AsyncWrite};
+    use futures::{AsyncRead, AsyncWrite};
     use keccak_const::Sha3_256;
     use libp2p::core::muxing;
     use libp2p::core::transport::Boxed;

rust/parts.nix

@@ -1,10 +1,11 @@
 { inputs, ... }:
 {
   perSystem =
-    { inputs', pkgs, lib, ... }:
+    { config, self', inputs', pkgs, lib, ... }:
     let
       # Fenix nightly toolchain with all components
-      rustToolchain = inputs'.fenix.packages.stable.withComponents [
+      fenixPkgs = inputs'.fenix.packages;
+      rustToolchain = fenixPkgs.complete.withComponents [
         "cargo"
         "rustc"
         "clippy"

rust/rust-toolchain.toml

@@ -0,0 +1,2 @@
+[toolchain]
+channel = "nightly"

src/exo/download/coordinator.py

@@ -14,7 +14,6 @@ from exo.download.download_utils import (
     map_repo_download_progress_to_download_progress_data,
 )
 from exo.download.shard_downloader import ShardDownloader
-from exo.shared.constants import EXO_MODELS_DIR
 from exo.shared.models.model_cards import ModelId
 from exo.shared.types.commands import (
     CancelDownload,
@@ -47,7 +46,6 @@ class DownloadCoordinator:
     download_command_receiver: Receiver[ForwarderDownloadCommand]
     local_event_sender: Sender[ForwarderEvent]
     event_index_counter: Iterator[int]
-    offline: bool = False
 
     # Local state
     download_status: dict[ModelId, DownloadProgress] = field(default_factory=dict)
@@ -63,13 +61,8 @@ class DownloadCoordinator:
 
     def __post_init__(self) -> None:
         self.event_sender, self.event_receiver = channel[Event]()
-        if self.offline:
-            self.shard_downloader.set_internet_connection(False)
         self.shard_downloader.on_progress(self._download_progress_callback)
 
-    def _model_dir(self, model_id: ModelId) -> str:
-        return str(EXO_MODELS_DIR / model_id.normalize())
-
     async def _download_progress_callback(
         self, callback_shard: ShardMetadata, progress: RepoDownloadProgress
     ) -> None:
@@ -81,7 +74,6 @@ class DownloadCoordinator:
                 shard_metadata=callback_shard,
                 node_id=self.node_id,
                 total_bytes=progress.total_bytes,
-                model_directory=self._model_dir(model_id),
             )
             self.download_status[model_id] = completed
             await self.event_sender.send(
@@ -101,7 +93,6 @@ class DownloadCoordinator:
                 download_progress=map_repo_download_progress_to_download_progress_data(
                     progress
                 ),
-                model_directory=self._model_dir(model_id),
             )
             self.download_status[model_id] = ongoing
             await self.event_sender.send(
@@ -110,17 +101,13 @@ class DownloadCoordinator:
             self._last_progress_time[model_id] = current_time()
 
     async def run(self) -> None:
-        logger.info(
-            f"Starting DownloadCoordinator{' (offline mode)' if self.offline else ''}"
-        )
-        if not self.offline:
-            self._test_internet_connection()
+        logger.info("Starting DownloadCoordinator")
+        self._test_internet_connection()
         async with self._tg as tg:
             tg.start_soon(self._command_processor)
             tg.start_soon(self._forward_events)
             tg.start_soon(self._emit_existing_download_progress)
-            if not self.offline:
-                tg.start_soon(self._check_internet_connection)
+            tg.start_soon(self._check_internet_connection)
 
     def _test_internet_connection(self) -> None:
         try:
@@ -183,11 +170,7 @@ class DownloadCoordinator:
                 return
 
         # Emit pending status
-        progress = DownloadPending(
-            shard_metadata=shard,
-            node_id=self.node_id,
-            model_directory=self._model_dir(model_id),
-        )
+        progress = DownloadPending(shard_metadata=shard, node_id=self.node_id)
         self.download_status[model_id] = progress
         await self.event_sender.send(NodeDownloadProgress(download_progress=progress))
 
@@ -201,7 +184,6 @@ class DownloadCoordinator:
                 shard_metadata=shard,
                 node_id=self.node_id,
                 total_bytes=initial_progress.total_bytes,
-                model_directory=self._model_dir(model_id),
             )
             self.download_status[model_id] = completed
             await self.event_sender.send(
@@ -209,20 +191,6 @@ class DownloadCoordinator:
             )
             return
 
-        if self.offline:
-            logger.warning(
-                f"Offline mode: model {model_id} is not fully available locally, cannot download"
-            )
-            failed = DownloadFailed(
-                shard_metadata=shard,
-                node_id=self.node_id,
-                error_message=f"Model files not found locally in offline mode: {model_id}",
-                model_directory=self._model_dir(model_id),
-            )
-            self.download_status[model_id] = failed
-            await self.event_sender.send(NodeDownloadProgress(download_progress=failed))
-            return
-
         # Start actual download
         self._start_download_task(shard, initial_progress)
 
@@ -238,7 +206,6 @@ class DownloadCoordinator:
             download_progress=map_repo_download_progress_to_download_progress_data(
                 initial_progress
             ),
-            model_directory=self._model_dir(model_id),
         )
         self.download_status[model_id] = status
         self.event_sender.send_nowait(NodeDownloadProgress(download_progress=status))
@@ -252,7 +219,6 @@ class DownloadCoordinator:
                     shard_metadata=shard,
                     node_id=self.node_id,
                     error_message=str(e),
-                    model_directory=self._model_dir(model_id),
                 )
                 self.download_status[model_id] = failed
                 await self.event_sender.send(
@@ -287,7 +253,6 @@ class DownloadCoordinator:
             pending = DownloadPending(
                 shard_metadata=current_status.shard_metadata,
                 node_id=self.node_id,
-                model_directory=self._model_dir(model_id),
             )
             await self.event_sender.send(
                 NodeDownloadProgress(download_progress=pending)
@@ -330,18 +295,11 @@ class DownloadCoordinator:
                             node_id=self.node_id,
                             shard_metadata=progress.shard,
                             total_bytes=progress.total_bytes,
-                            model_directory=self._model_dir(
-                                progress.shard.model_card.model_id
-                            ),
                         )
                     elif progress.status in ["in_progress", "not_started"]:
                         if progress.downloaded_bytes_this_session.in_bytes == 0:
                             status = DownloadPending(
-                                node_id=self.node_id,
-                                shard_metadata=progress.shard,
-                                model_directory=self._model_dir(
-                                    progress.shard.model_card.model_id
-                                ),
+                                node_id=self.node_id, shard_metadata=progress.shard
                             )
                         else:
                             status = DownloadOngoing(
@@ -350,9 +308,6 @@ class DownloadCoordinator:
                                 download_progress=map_repo_download_progress_to_download_progress_data(
                                     progress
                                 ),
-                                model_directory=self._model_dir(
-                                    progress.shard.model_card.model_id
-                                ),
                             )
                     else:
                         continue

src/exo/download/download_utils.py

@@ -448,13 +448,12 @@ async def download_file_with_retry(
     target_dir: Path,
     on_progress: Callable[[int, int, bool], None] = lambda _, __, ___: None,
     on_connection_lost: Callable[[], None] = lambda: None,
-    skip_internet: bool = False,
 ) -> Path:
     n_attempts = 3
     for attempt in range(n_attempts):
         try:
             return await _download_file(
-                model_id, revision, path, target_dir, on_progress, skip_internet
+                model_id, revision, path, target_dir, on_progress
             )
         except HuggingFaceAuthenticationError:
             raise
@@ -488,14 +487,10 @@ async def _download_file(
     path: str,
     target_dir: Path,
     on_progress: Callable[[int, int, bool], None] = lambda _, __, ___: None,
-    skip_internet: bool = False,
 ) -> Path:
     target_path = target_dir / path
 
     if await aios.path.exists(target_path):
-        if skip_internet:
-            return target_path
-
         local_size = (await aios.stat(target_path)).st_size
 
         # Try to verify against remote, but allow offline operation
@@ -515,11 +510,6 @@ async def _download_file(
             )
             return target_path
 
-    if skip_internet:
-        raise FileNotFoundError(
-            f"File {path} not found locally and cannot download in offline mode"
-        )
-
     await aios.makedirs((target_dir / path).parent, exist_ok=True)
     length, etag = await file_meta(model_id, revision, path)
     remote_hash = etag[:-5] if etag.endswith("-gzip") else etag
@@ -824,7 +814,6 @@ async def download_shard(
                     file, curr_bytes, total_bytes, is_renamed
                 ),
                 on_connection_lost=on_connection_lost,
-                skip_internet=skip_internet,
             )
 
     if not skip_download:

src/exo/download/tests/test_offline_mode.py

@@ -1,230 +0,0 @@
-"""Tests for offline/air-gapped mode."""
-
-from collections.abc import AsyncIterator
-from pathlib import Path
-from unittest.mock import AsyncMock, patch
-
-import aiofiles
-import aiofiles.os as aios
-import pytest
-
-from exo.download.download_utils import (
-    _download_file,  # pyright: ignore[reportPrivateUsage]
-    download_file_with_retry,
-    fetch_file_list_with_cache,
-)
-from exo.shared.types.common import ModelId
-from exo.shared.types.worker.downloads import FileListEntry
-
-
-@pytest.fixture
-def model_id() -> ModelId:
-    return ModelId("test-org/test-model")
-
-
-@pytest.fixture
-async def temp_models_dir(tmp_path: Path) -> AsyncIterator[Path]:
-    models_dir = tmp_path / "models"
-    await aios.makedirs(models_dir, exist_ok=True)
-    with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
-        yield models_dir
-
-
-class TestDownloadFileOffline:
-    """Tests for _download_file with skip_internet=True."""
-
-    async def test_returns_local_file_without_http_verification(
-        self, model_id: ModelId, tmp_path: Path
-    ) -> None:
-        """When skip_internet=True and file exists locally, return it immediately
-        without making any HTTP calls (no file_meta verification)."""
-        target_dir = tmp_path / "downloads"
-        await aios.makedirs(target_dir, exist_ok=True)
-
-        local_file = target_dir / "model.safetensors"
-        async with aiofiles.open(local_file, "wb") as f:
-            await f.write(b"model weights data")
-
-        with patch(
-            "exo.download.download_utils.file_meta",
-            new_callable=AsyncMock,
-        ) as mock_file_meta:
-            result = await _download_file(
-                model_id,
-                "main",
-                "model.safetensors",
-                target_dir,
-                skip_internet=True,
-            )
-
-            assert result == local_file
-            mock_file_meta.assert_not_called()
-
-    async def test_raises_file_not_found_for_missing_file(
-        self, model_id: ModelId, tmp_path: Path
-    ) -> None:
-        """When skip_internet=True and file does NOT exist locally,
-        raise FileNotFoundError instead of attempting download."""
-        target_dir = tmp_path / "downloads"
-        await aios.makedirs(target_dir, exist_ok=True)
-
-        with pytest.raises(FileNotFoundError, match="offline mode"):
-            await _download_file(
-                model_id,
-                "main",
-                "missing_model.safetensors",
-                target_dir,
-                skip_internet=True,
-            )
-
-    async def test_returns_local_file_in_subdirectory(
-        self, model_id: ModelId, tmp_path: Path
-    ) -> None:
-        """When skip_internet=True and file exists in a subdirectory,
-        return it without HTTP calls."""
-        target_dir = tmp_path / "downloads"
-        subdir = target_dir / "transformer"
-        await aios.makedirs(subdir, exist_ok=True)
-
-        local_file = subdir / "diffusion_pytorch_model.safetensors"
-        async with aiofiles.open(local_file, "wb") as f:
-            await f.write(b"weights")
-
-        with patch(
-            "exo.download.download_utils.file_meta",
-            new_callable=AsyncMock,
-        ) as mock_file_meta:
-            result = await _download_file(
-                model_id,
-                "main",
-                "transformer/diffusion_pytorch_model.safetensors",
-                target_dir,
-                skip_internet=True,
-            )
-
-            assert result == local_file
-            mock_file_meta.assert_not_called()
-
-
-class TestDownloadFileWithRetryOffline:
-    """Tests for download_file_with_retry with skip_internet=True."""
-
-    async def test_propagates_skip_internet_to_download_file(
-        self, model_id: ModelId, tmp_path: Path
-    ) -> None:
-        """Verify skip_internet is passed through to _download_file."""
-        target_dir = tmp_path / "downloads"
-        await aios.makedirs(target_dir, exist_ok=True)
-
-        local_file = target_dir / "config.json"
-        async with aiofiles.open(local_file, "wb") as f:
-            await f.write(b'{"model_type": "qwen2"}')
-
-        with patch(
-            "exo.download.download_utils.file_meta",
-            new_callable=AsyncMock,
-        ) as mock_file_meta:
-            result = await download_file_with_retry(
-                model_id,
-                "main",
-                "config.json",
-                target_dir,
-                skip_internet=True,
-            )
-
-            assert result == local_file
-            mock_file_meta.assert_not_called()
-
-    async def test_file_not_found_does_not_retry(
-        self, model_id: ModelId, tmp_path: Path
-    ) -> None:
-        """FileNotFoundError from offline mode should not trigger retries."""
-        target_dir = tmp_path / "downloads"
-        await aios.makedirs(target_dir, exist_ok=True)
-
-        with pytest.raises(FileNotFoundError):
-            await download_file_with_retry(
-                model_id,
-                "main",
-                "nonexistent.safetensors",
-                target_dir,
-                skip_internet=True,
-            )
-
-
-class TestFetchFileListOffline:
-    """Tests for fetch_file_list_with_cache with skip_internet=True."""
-
-    async def test_uses_cached_file_list(
-        self, model_id: ModelId, temp_models_dir: Path
-    ) -> None:
-        """When skip_internet=True and cache file exists, use it without network."""
-        from pydantic import TypeAdapter
-
-        cache_dir = temp_models_dir / "caches" / model_id.normalize()
-        await aios.makedirs(cache_dir, exist_ok=True)
-
-        cached_list = [
-            FileListEntry(type="file", path="model.safetensors", size=1000),
-            FileListEntry(type="file", path="config.json", size=200),
-        ]
-        cache_file = cache_dir / f"{model_id.normalize()}--main--file_list.json"
-        async with aiofiles.open(cache_file, "w") as f:
-            await f.write(
-                TypeAdapter(list[FileListEntry]).dump_json(cached_list).decode()
-            )
-
-        with patch(
-            "exo.download.download_utils.fetch_file_list_with_retry",
-            new_callable=AsyncMock,
-        ) as mock_fetch:
-            result = await fetch_file_list_with_cache(
-                model_id, "main", skip_internet=True
-            )
-
-            assert result == cached_list
-            mock_fetch.assert_not_called()
-
-    async def test_falls_back_to_local_directory_scan(
-        self, model_id: ModelId, temp_models_dir: Path
-    ) -> None:
-        """When skip_internet=True and no cache but local files exist,
-        build file list from local directory."""
-        import json
-
-        model_dir = temp_models_dir / model_id.normalize()
-        await aios.makedirs(model_dir, exist_ok=True)
-
-        async with aiofiles.open(model_dir / "config.json", "w") as f:
-            await f.write('{"model_type": "qwen2"}')
-
-        index_data = {
-            "metadata": {},
-            "weight_map": {"model.layers.0.weight": "model.safetensors"},
-        }
-        async with aiofiles.open(model_dir / "model.safetensors.index.json", "w") as f:
-            await f.write(json.dumps(index_data))
-
-        async with aiofiles.open(model_dir / "model.safetensors", "wb") as f:
-            await f.write(b"x" * 500)
-
-        with patch(
-            "exo.download.download_utils.fetch_file_list_with_retry",
-            new_callable=AsyncMock,
-        ) as mock_fetch:
-            result = await fetch_file_list_with_cache(
-                model_id, "main", skip_internet=True
-            )
-
-            mock_fetch.assert_not_called()
-            paths = {entry.path for entry in result}
-            assert "config.json" in paths
-            assert "model.safetensors" in paths
-
-    async def test_raises_when_no_cache_and_no_local_files(
-        self, model_id: ModelId, temp_models_dir: Path
-    ) -> None:
-        """When skip_internet=True and neither cache nor local files exist,
-        raise FileNotFoundError."""
-        with pytest.raises(FileNotFoundError, match="No internet"):
-            await fetch_file_list_with_cache(model_id, "main", skip_internet=True)

src/exo/main.py

@@ -39,13 +39,12 @@ class Node:
 
     node_id: NodeId
     event_index_counter: Iterator[int]
-    offline: bool
     _tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)
 
     @classmethod
     async def create(cls, args: "Args") -> "Self":
         keypair = get_node_id_keypair()
-        node_id = NodeId(keypair.to_node_id())
+        node_id = NodeId(keypair.to_peer_id().to_base58())
         session_id = SessionId(master_node_id=node_id, election_clock=0)
         router = Router.create(keypair)
         await router.register_topic(topics.GLOBAL_EVENTS)
@@ -69,7 +68,6 @@ class Node:
                 download_command_receiver=router.receiver(topics.DOWNLOAD_COMMANDS),
                 local_event_sender=router.sender(topics.LOCAL_EVENTS),
                 event_index_counter=event_index_counter,
-                offline=args.offline,
             )
         else:
             download_coordinator = None
@@ -134,13 +132,10 @@ class Node:
             api,
             node_id,
             event_index_counter,
-            args.offline,
         )
 
     async def run(self):
         async with self._tg as tg:
-            signal.signal(signal.SIGINT, lambda _, __: self.shutdown())
-            signal.signal(signal.SIGTERM, lambda _, __: self.shutdown())
             tg.start_soon(self.router.run)
             tg.start_soon(self.election.run)
             if self.download_coordinator:
@@ -152,6 +147,8 @@ class Node:
             if self.api:
                 tg.start_soon(self.api.run)
             tg.start_soon(self._elect_loop)
+            signal.signal(signal.SIGINT, lambda _, __: self.shutdown())
+            signal.signal(signal.SIGTERM, lambda _, __: self.shutdown())
 
     def shutdown(self):
         # if this is our second call to shutdown, just sys.exit
@@ -225,7 +222,6 @@ class Node:
                             ),
                             local_event_sender=self.router.sender(topics.LOCAL_EVENTS),
                             event_index_counter=self.event_index_counter,
-                            offline=self.offline,
                         )
                         self._tg.start_soon(self.download_coordinator.run)
                     if self.worker:
@@ -264,9 +260,6 @@ def main():
     logger.info("Starting EXO")
     logger.info(f"EXO_LIBP2P_NAMESPACE: {os.getenv('EXO_LIBP2P_NAMESPACE')}")
 
-    if args.offline:
-        logger.info("Running in OFFLINE mode — no internet checks, local models only")
-
     # Set FAST_SYNCH override env var for runner subprocesses
     if args.fast_synch is True:
         os.environ["EXO_FAST_SYNCH"] = "on"
@@ -289,7 +282,6 @@ class Args(CamelCaseModel):
     tb_only: bool = False
     no_worker: bool = False
     no_downloads: bool = False
-    offline: bool = False
     fast_synch: bool | None = None  # None = auto, True = force on, False = force off
 
     @classmethod
@@ -337,11 +329,6 @@ class Args(CamelCaseModel):
             action="store_true",
             help="Disable the download coordinator (node won't download models)",
         )
-        parser.add_argument(
-            "--offline",
-            action="store_true",
-            help="Run in offline/air-gapped mode: skip internet checks, use only pre-staged local models",
-        )
         fast_synch_group = parser.add_mutually_exclusive_group()
         fast_synch_group.add_argument(
             "--fast-synch",

src/exo/master/adapters/chat_completions.py

@@ -17,9 +17,13 @@ from exo.shared.types.api import (
     LogprobsContentItem,
     StreamingChoiceResponse,
     ToolCall,
-    Usage,
 )
-from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
+from exo.shared.types.chunks import (
+    ErrorChunk,
+    PrefillProgressData,
+    TokenChunk,
+    ToolCallChunk,
+)
 from exo.shared.types.common import CommandId
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 
@@ -123,70 +127,71 @@ def chunk_to_response(
 
 async def generate_chat_stream(
     command_id: CommandId,
-    chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
+    event_stream: AsyncGenerator[
+        PrefillProgressData | ErrorChunk | ToolCallChunk | TokenChunk, None
+    ],
 ) -> AsyncGenerator[str, None]:
-    """Generate Chat Completions API streaming events from chunks."""
-    last_usage: Usage | None = None
+    """Generate Chat Completions API streaming events from StreamEvents.
 
-    async for chunk in chunk_stream:
-        if isinstance(chunk, ErrorChunk):
-            error_response = ErrorResponse(
-                error=ErrorInfo(
-                    message=chunk.error_message or "Internal server error",
-                    type="InternalServerError",
-                    code=500,
-                )
-            )
-            yield f"data: {error_response.model_dump_json()}\n\n"
-            yield "data: [DONE]\n\n"
-            return
+    Handles PrefillProgressData, ErrorChunk, ToolCallChunk, and TokenChunk.
+    """
+    async for event in event_stream:
+        match event:
+            case PrefillProgressData():
+                yield f"event: prefill_progress\ndata: {event.model_dump_json()}\n\n"
 
-        last_usage = chunk.usage or last_usage
-
-        if isinstance(chunk, ToolCallChunk):
-            tool_call_deltas = [
-                ToolCall(
-                    id=tool.id,
-                    index=i,
-                    function=tool,
-                )
-                for i, tool in enumerate(chunk.tool_calls)
-            ]
-            tool_response = ChatCompletionResponse(
-                id=command_id,
-                created=int(time.time()),
-                model=chunk.model,
-                choices=[
-                    StreamingChoiceResponse(
-                        index=0,
-                        delta=ChatCompletionMessage(
-                            role="assistant",
-                            tool_calls=tool_call_deltas,
-                        ),
-                        finish_reason="tool_calls",
+            case ErrorChunk():
+                error_response = ErrorResponse(
+                    error=ErrorInfo(
+                        message=event.error_message or "Internal server error",
+                        type="InternalServerError",
+                        code=500,
                     )
-                ],
-                usage=last_usage,
-            )
-            yield f"data: {tool_response.model_dump_json()}\n\n"
-            yield "data: [DONE]\n\n"
-            return
+                )
+                yield f"data: {error_response.model_dump_json()}\n\n"
+                yield "data: [DONE]\n\n"
+                return
 
-        chunk_response = chunk_to_response(chunk, command_id)
-        if chunk.finish_reason is not None:
-            chunk_response = chunk_response.model_copy(update={"usage": last_usage})
-        yield f"data: {chunk_response.model_dump_json()}\n\n"
+            case ToolCallChunk():
+                tool_call_deltas = [
+                    ToolCall(
+                        id=tool.id,
+                        index=i,
+                        function=tool,
+                    )
+                    for i, tool in enumerate(event.tool_calls)
+                ]
+                tool_response = ChatCompletionResponse(
+                    id=command_id,
+                    created=int(time.time()),
+                    model=event.model,
+                    choices=[
+                        StreamingChoiceResponse(
+                            index=0,
+                            delta=ChatCompletionMessage(
+                                role="assistant",
+                                tool_calls=tool_call_deltas,
+                            ),
+                            finish_reason="tool_calls",
+                        )
+                    ],
+                )
+                yield f"data: {tool_response.model_dump_json()}\n\n"
+                yield "data: [DONE]\n\n"
+                return
 
-        if chunk.finish_reason is not None:
-            yield "data: [DONE]\n\n"
+            case TokenChunk():
+                chunk_response = chunk_to_response(event, command_id)
+                yield f"data: {chunk_response.model_dump_json()}\n\n"
+
+                if event.finish_reason is not None:
+                    yield "data: [DONE]\n\n"
 
 
 async def collect_chat_response(
     command_id: CommandId,
     chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
-) -> AsyncGenerator[str]:
-    # This is an AsyncGenerator[str] rather than returning a ChatCompletionReponse because
-    # FastAPI handles the cancellation better but wouldn't auto-serialize for some reason
+) -> ChatCompletionResponse:
     """Collect all token chunks and return a single ChatCompletionResponse."""
     text_parts: list[str] = []
     tool_calls: list[ToolCall] = []
@@ -194,7 +199,6 @@ async def collect_chat_response(
     model: str | None = None
     finish_reason: FinishReason | None = None
     error_message: str | None = None
-    last_usage: Usage | None = None
 
     async for chunk in chunk_stream:
         if isinstance(chunk, ErrorChunk):
@@ -204,8 +208,6 @@ async def collect_chat_response(
         if model is None:
             model = chunk.model
 
-        last_usage = chunk.usage or last_usage
-
         if isinstance(chunk, TokenChunk):
             text_parts.append(chunk.text)
             if chunk.logprob is not None:
@@ -236,7 +238,7 @@ async def collect_chat_response(
     combined_text = "".join(text_parts)
     assert model is not None
 
-    yield ChatCompletionResponse(
+    return ChatCompletionResponse(
         id=command_id,
         created=int(time.time()),
         model=model,
@@ -254,6 +256,4 @@ async def collect_chat_response(
                 finish_reason=finish_reason,
             )
         ],
-        usage=last_usage,
-    ).model_dump_json()
-    return
+    )

src/exo/master/adapters/claude.py

@@ -4,7 +4,7 @@ import json
 from collections.abc import AsyncGenerator
 from typing import Any
 
-from exo.shared.types.api import FinishReason, Usage
+from exo.shared.types.api import FinishReason
 from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
 from exo.shared.types.claude_api import (
     ClaudeContentBlock,
@@ -161,14 +161,12 @@ async def collect_claude_response(
     command_id: CommandId,
     model: str,
     chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
-) -> AsyncGenerator[str]:
-    # This is an AsyncGenerator[str] rather than returning a ChatCompletionReponse because
-    # FastAPI handles the cancellation better but wouldn't auto-serialize for some reason
+) -> ClaudeMessagesResponse:
     """Collect all token chunks and return a single ClaudeMessagesResponse."""
     text_parts: list[str] = []
     tool_use_blocks: list[ClaudeToolUseBlock] = []
     stop_reason: ClaudeStopReason | None = None
-    last_usage: Usage | None = None
+    last_stats = None
     error_message: str | None = None
 
     async for chunk in chunk_stream:
@@ -176,8 +174,6 @@ async def collect_claude_response(
             error_message = chunk.error_message or "Internal server error"
             break
 
-        last_usage = chunk.usage or last_usage
-
         if isinstance(chunk, ToolCallChunk):
             for tool in chunk.tool_calls:
                 tool_use_blocks.append(
@@ -187,10 +183,12 @@ async def collect_claude_response(
                         input=json.loads(tool.arguments),  # pyright: ignore[reportAny]
                     )
                 )
+            last_stats = chunk.stats or last_stats
             stop_reason = "tool_use"
             continue
 
         text_parts.append(chunk.text)
+        last_stats = chunk.stats or last_stats
 
         if chunk.finish_reason is not None:
             stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason)
@@ -210,11 +208,11 @@ async def collect_claude_response(
     if not content:
         content.append(ClaudeTextBlock(text=""))
 
-    # Use actual usage data if available
-    input_tokens = last_usage.prompt_tokens if last_usage else 0
-    output_tokens = last_usage.completion_tokens if last_usage else 0
+    # Use actual usage data from stats if available
+    input_tokens = last_stats.prompt_tokens if last_stats else 0
+    output_tokens = last_stats.generation_tokens if last_stats else 0
 
-    yield ClaudeMessagesResponse(
+    return ClaudeMessagesResponse(
         id=f"msg_{command_id}",
         model=model,
         content=content,
@@ -223,8 +221,7 @@ async def collect_claude_response(
             input_tokens=input_tokens,
             output_tokens=output_tokens,
         ),
-    ).model_dump_json()
-    return
+    )
 
 
 async def generate_claude_stream(
@@ -252,7 +249,7 @@ async def generate_claude_stream(
 
     output_tokens = 0
     stop_reason: ClaudeStopReason | None = None
-    last_usage: Usage | None = None
+    last_stats = None
     next_block_index = 1  # text block is 0, tool blocks start at 1
 
     async for chunk in chunk_stream:
@@ -260,9 +257,8 @@ async def generate_claude_stream(
             # Close text block and bail
             break
 
-        last_usage = chunk.usage or last_usage
-
         if isinstance(chunk, ToolCallChunk):
+            last_stats = chunk.stats or last_stats
             stop_reason = "tool_use"
 
             # Emit tool_use content blocks
@@ -294,6 +290,7 @@ async def generate_claude_stream(
             continue
 
         output_tokens += 1  # Count each chunk as one token
+        last_stats = chunk.stats or last_stats
 
         # content_block_delta
         delta_event = ClaudeContentBlockDeltaEvent(
@@ -305,9 +302,9 @@ async def generate_claude_stream(
         if chunk.finish_reason is not None:
             stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason)
 
-    # Use actual token count from usage if available
-    if last_usage is not None:
-        output_tokens = last_usage.completion_tokens
+    # Use actual token count from stats if available
+    if last_stats is not None:
+        output_tokens = last_stats.generation_tokens
 
     # content_block_stop for text block
     block_stop = ClaudeContentBlockStopEvent(index=0)

src/exo/master/adapters/responses.py

@@ -4,7 +4,6 @@ from collections.abc import AsyncGenerator
 from itertools import count
 from typing import Any
 
-from exo.shared.types.api import Usage
 from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
 from exo.shared.types.common import CommandId
 from exo.shared.types.openai_responses import (
@@ -122,15 +121,13 @@ async def collect_responses_response(
     command_id: CommandId,
     model: str,
     chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
-) -> AsyncGenerator[str]:
-    # This is an AsyncGenerator[str] rather than returning a ChatCompletionReponse because
-    # FastAPI handles the cancellation better but wouldn't auto-serialize for some reason
+) -> ResponsesResponse:
     """Collect all token chunks and return a single ResponsesResponse."""
     response_id = f"resp_{command_id}"
     item_id = f"item_{command_id}"
     accumulated_text = ""
     function_call_items: list[ResponseFunctionCallItem] = []
-    last_usage: Usage | None = None
+    last_stats = None
     error_message: str | None = None
 
     async for chunk in chunk_stream:
@@ -138,32 +135,32 @@ async def collect_responses_response(
             error_message = chunk.error_message or "Internal server error"
             break
 
-        last_usage = chunk.usage or last_usage
-
         if isinstance(chunk, ToolCallChunk):
             for tool in chunk.tool_calls:
                 function_call_items.append(
                     ResponseFunctionCallItem(
-                        id=tool.id,
-                        call_id=tool.id,
+                        id=f"fc_{tool.id}",
+                        call_id=f"call_{tool.id}",
                         name=tool.name,
                         arguments=tool.arguments,
                     )
                 )
+            last_stats = chunk.stats or last_stats
             continue
 
         accumulated_text += chunk.text
+        last_stats = chunk.stats or last_stats
 
     if error_message is not None:
         raise ValueError(error_message)
 
-    # Create usage from usage data if available
+    # Create usage from stats if available
     usage = None
-    if last_usage is not None:
+    if last_stats is not None:
         usage = ResponseUsage(
-            input_tokens=last_usage.prompt_tokens,
-            output_tokens=last_usage.completion_tokens,
-            total_tokens=last_usage.total_tokens,
+            input_tokens=last_stats.prompt_tokens,
+            output_tokens=last_stats.generation_tokens,
+            total_tokens=last_stats.prompt_tokens + last_stats.generation_tokens,
         )
 
     output: list[ResponseItem] = [
@@ -175,15 +172,14 @@ async def collect_responses_response(
     ]
     output.extend(function_call_items)
 
-    yield ResponsesResponse(
+    return ResponsesResponse(
         id=response_id,
         model=model,
         status="completed",
         output=output,
         output_text=accumulated_text,
         usage=usage,
-    ).model_dump_json()
-    return
+    )
 
 
 async def generate_responses_stream(
@@ -239,16 +235,15 @@ async def generate_responses_stream(
 
     accumulated_text = ""
     function_call_items: list[ResponseFunctionCallItem] = []
-    last_usage: Usage | None = None
+    last_stats = None
     next_output_index = 1  # message item is at 0
 
     async for chunk in chunk_stream:
         if isinstance(chunk, ErrorChunk):
             break
 
-        last_usage = chunk.usage or last_usage
-
         if isinstance(chunk, ToolCallChunk):
+            last_stats = chunk.stats or last_stats
             for tool in chunk.tool_calls:
                 fc_id = f"fc_{tool.id}"
                 call_id = f"call_{tool.id}"
@@ -307,6 +302,7 @@ async def generate_responses_stream(
             continue
 
         accumulated_text += chunk.text
+        last_stats = chunk.stats or last_stats
 
         # response.output_text.delta
         delta_event = ResponseTextDeltaEvent(
@@ -350,13 +346,13 @@ async def generate_responses_stream(
     )
     yield f"event: response.output_item.done\ndata: {item_done.model_dump_json()}\n\n"
 
-    # Create usage from usage data if available
+    # Create usage from stats if available
     usage = None
-    if last_usage is not None:
+    if last_stats is not None:
         usage = ResponseUsage(
-            input_tokens=last_usage.prompt_tokens,
-            output_tokens=last_usage.completion_tokens,
-            total_tokens=last_usage.total_tokens,
+            input_tokens=last_stats.prompt_tokens,
+            output_tokens=last_stats.generation_tokens,
+            total_tokens=last_stats.prompt_tokens + last_stats.generation_tokens,
         )
 
     # response.completed

src/exo/master/api.py

@@ -85,7 +85,6 @@ from exo.shared.types.api import (
     ImageGenerationTaskParams,
     ImageListItem,
     ImageListResponse,
-    ImageSize,
     ModelList,
     ModelListModel,
     PlaceInstanceParams,
@@ -101,12 +100,12 @@ from exo.shared.types.api import (
     TraceRankStats,
     TraceResponse,
     TraceStatsResponse,
-    normalize_image_size,
 )
 from exo.shared.types.chunks import (
     ErrorChunk,
     ImageChunk,
     InputImageChunk,
+    PrefillProgressData,
     TokenChunk,
     ToolCallChunk,
 )
@@ -127,7 +126,6 @@ from exo.shared.types.commands import (
     PlaceInstance,
     SendInputChunk,
     StartDownload,
-    TaskCancelled,
     TaskFinished,
     TextGeneration,
 )
@@ -137,6 +135,7 @@ from exo.shared.types.events import (
     Event,
     ForwarderEvent,
     IndexedEvent,
+    PrefillProgress,
     TracesMerged,
 )
 from exo.shared.types.memory import Memory
@@ -220,7 +219,8 @@ class API:
         )
 
         self._text_generation_queues: dict[
-            CommandId, Sender[TokenChunk | ErrorChunk | ToolCallChunk]
+            CommandId,
+            Sender[TokenChunk | ErrorChunk | ToolCallChunk | PrefillProgressData],
         ] = {}
         self._image_generation_queues: dict[
             CommandId, Sender[ImageChunk | ErrorChunk]
@@ -524,36 +524,51 @@ class API:
             instance_id=instance_id,
         )
 
-    async def _token_chunk_stream(
+    async def _stream_events(
         self, command_id: CommandId
-    ) -> AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None]:
-        """Yield chunks for a given command until completion.
+    ) -> AsyncGenerator[
+        TokenChunk | ErrorChunk | ToolCallChunk | PrefillProgressData, None
+    ]:
+        """Yield stream events for a command.
 
         This is the internal low-level stream used by all API adapters.
         """
         try:
             self._text_generation_queues[command_id], recv = channel[
-                ErrorChunk | ToolCallChunk | TokenChunk
+                TokenChunk | ErrorChunk | ToolCallChunk | PrefillProgressData
             ]()
 
-            with recv as token_chunks:
-                async for chunk in token_chunks:
-                    yield chunk
-                    if chunk.finish_reason is not None:
+            with recv as events:
+                async for event in events:
+                    yield event
+                    if (
+                        isinstance(event, TokenChunk)
+                        and event.finish_reason is not None
+                    ):
                         break
 
         except anyio.get_cancelled_exc_class():
-            command = TaskCancelled(cancelled_command_id=command_id)
-            with anyio.CancelScope(shield=True):
-                await self.command_sender.send(
-                    ForwarderCommand(origin=self.node_id, command=command)
-                )
+            # TODO: TaskCancelled
+            """
+            self.command_sender.send_nowait(
+                ForwarderCommand(origin=self.node_id, command=command)
+            )
+            """
             raise
         finally:
-            await self._send(TaskFinished(finished_command_id=command_id))
+            command = TaskFinished(finished_command_id=command_id)
+            await self._send(command)
             if command_id in self._text_generation_queues:
                 del self._text_generation_queues[command_id]
 
+    async def _chunk_stream(
+        self, command_id: CommandId
+    ) -> AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None]:
+        """Yield chunks, filtering out prefill progress events."""
+        async for event in self._stream_events(command_id):
+            if not isinstance(event, PrefillProgressData):
+                yield event
+
     async def _collect_text_generation_with_stats(
         self, command_id: CommandId
     ) -> BenchChatCompletionResponse:
@@ -564,7 +579,7 @@ class API:
 
         stats: GenerationStats | None = None
 
-        async for chunk in self._token_chunk_stream(command_id):
+        async for chunk in self._chunk_stream(command_id):
             if chunk.finish_reason == "error":
                 raise HTTPException(
                     status_code=500,
@@ -636,7 +651,7 @@ class API:
             return StreamingResponse(
                 generate_chat_stream(
                     command.command_id,
-                    self._token_chunk_stream(command.command_id),
+                    self._stream_events(command.command_id),
                 ),
                 media_type="text/event-stream",
                 headers={
@@ -645,14 +660,14 @@ class API:
                     "X-Accel-Buffering": "no",
                 },
             )
-        else:
-            return StreamingResponse(
-                collect_chat_response(
-                    command.command_id,
-                    self._token_chunk_stream(command.command_id),
-                ),
-                media_type="application/json",
+
+        try:
+            return await collect_chat_response(
+                command.command_id,
+                self._chunk_stream(command.command_id),
             )
+        except ValueError as e:
+            raise HTTPException(status_code=500, detail=str(e)) from e
 
     async def bench_chat_completions(
         self, payload: BenchChatCompletionRequest
@@ -668,7 +683,8 @@ class API:
         command = TextGeneration(task_params=task_params)
         await self._send(command)
 
-        return await self._collect_text_generation_with_stats(command.command_id)
+        response = await self._collect_text_generation_with_stats(command.command_id)
+        return response
 
     async def _resolve_and_validate_text_model(self, model_id: ModelId) -> ModelId:
         """Validate a text model exists and return the resolved model ID.
@@ -753,11 +769,9 @@ class API:
         When stream=True and partial_images > 0, returns a StreamingResponse
         with SSE-formatted events for partial and final images.
         """
+        payload.model = await self._validate_image_model(ModelId(payload.model))
         payload = payload.model_copy(
-            update={
-                "model": await self._validate_image_model(ModelId(payload.model)),
-                "advanced_params": _ensure_seed(payload.advanced_params),
-            }
+            update={"advanced_params": _ensure_seed(payload.advanced_params)}
         )
 
         command = ImageGeneration(
@@ -888,11 +902,6 @@ class API:
                         del image_metadata[key]
 
         except anyio.get_cancelled_exc_class():
-            command = TaskCancelled(cancelled_command_id=command_id)
-            with anyio.CancelScope(shield=True):
-                await self.command_sender.send(
-                    ForwarderCommand(origin=self.node_id, command=command)
-                )
             raise
         finally:
             await self._send(TaskFinished(finished_command_id=command_id))
@@ -974,11 +983,6 @@ class API:
 
             return (images, stats if capture_stats else None)
         except anyio.get_cancelled_exc_class():
-            command = TaskCancelled(cancelled_command_id=command_id)
-            with anyio.CancelScope(shield=True):
-                await self.command_sender.send(
-                    ForwarderCommand(origin=self.node_id, command=command)
-                )
             raise
         finally:
             await self._send(TaskFinished(finished_command_id=command_id))
@@ -1013,13 +1017,12 @@ class API:
     async def bench_image_generations(
         self, request: Request, payload: BenchImageGenerationTaskParams
     ) -> BenchImageGenerationResponse:
+        payload.model = await self._validate_image_model(ModelId(payload.model))
+
+        payload.stream = False
+        payload.partial_images = 0
         payload = payload.model_copy(
-            update={
-                "model": await self._validate_image_model(ModelId(payload.model)),
-                "stream": False,
-                "partial_images": 0,
-                "advanced_params": _ensure_seed(payload.advanced_params),
-            }
+            update={"advanced_params": _ensure_seed(payload.advanced_params)}
         )
 
         command = ImageGeneration(
@@ -1040,7 +1043,7 @@ class API:
         prompt: str,
         model: ModelId,
         n: int,
-        size: ImageSize,
+        size: str,
         response_format: Literal["url", "b64_json"],
         input_fidelity: Literal["low", "high"],
         stream: bool,
@@ -1110,7 +1113,7 @@ class API:
         prompt: str = Form(...),
         model: str = Form(...),
         n: int = Form(1),
-        size: str | None = Form(None),
+        size: str = Form("1024x1024"),
         response_format: Literal["url", "b64_json"] = Form("b64_json"),
         input_fidelity: Literal["low", "high"] = Form("low"),
         stream: str = Form("false"),
@@ -1136,7 +1139,7 @@ class API:
             prompt=prompt,
             model=ModelId(model),
             n=n,
-            size=normalize_image_size(size),
+            size=size,
             response_format=response_format,
             input_fidelity=input_fidelity,
             stream=stream_bool,
@@ -1172,7 +1175,7 @@ class API:
         prompt: str = Form(...),
         model: str = Form(...),
         n: int = Form(1),
-        size: str | None = Form(None),
+        size: str = Form("1024x1024"),
         response_format: Literal["url", "b64_json"] = Form("b64_json"),
         input_fidelity: Literal["low", "high"] = Form("low"),
         quality: Literal["high", "medium", "low"] = Form("medium"),
@@ -1192,7 +1195,7 @@ class API:
             prompt=prompt,
             model=ModelId(model),
             n=n,
-            size=normalize_image_size(size),
+            size=size,
             response_format=response_format,
             input_fidelity=input_fidelity,
             stream=False,
@@ -1228,7 +1231,7 @@ class API:
                 generate_claude_stream(
                     command.command_id,
                     payload.model,
-                    self._token_chunk_stream(command.command_id),
+                    self._chunk_stream(command.command_id),
                 ),
                 media_type="text/event-stream",
                 headers={
@@ -1237,15 +1240,15 @@ class API:
                     "X-Accel-Buffering": "no",
                 },
             )
-        else:
-            return StreamingResponse(
-                collect_claude_response(
-                    command.command_id,
-                    payload.model,
-                    self._token_chunk_stream(command.command_id),
-                ),
-                media_type="application/json",
+
+        try:
+            return await collect_claude_response(
+                command.command_id,
+                payload.model,
+                self._chunk_stream(command.command_id),
             )
+        except ValueError as e:
+            raise HTTPException(status_code=500, detail=str(e)) from e
 
     async def openai_responses(
         self, payload: ResponsesRequest
@@ -1263,7 +1266,7 @@ class API:
                 generate_responses_stream(
                     command.command_id,
                     payload.model,
-                    self._token_chunk_stream(command.command_id),
+                    self._chunk_stream(command.command_id),
                 ),
                 media_type="text/event-stream",
                 headers={
@@ -1273,15 +1276,14 @@ class API:
                 },
             )
 
-        else:
-            return StreamingResponse(
-                collect_responses_response(
-                    command.command_id,
-                    payload.model,
-                    self._token_chunk_stream(command.command_id),
-                ),
-                media_type="application/json",
+        try:
+            return await collect_responses_response(
+                command.command_id,
+                payload.model,
+                self._chunk_stream(command.command_id),
             )
+        except ValueError as e:
+            raise HTTPException(status_code=500, detail=str(e)) from e
 
     def _calculate_total_available_memory(self) -> Memory:
         """Calculate total available memory across all nodes in bytes."""
@@ -1435,6 +1437,20 @@ class API:
                             except BrokenResourceError:
                                 self._text_generation_queues.pop(event.command_id, None)
 
+                    elif isinstance(event, PrefillProgress):
+                        if queue := self._text_generation_queues.get(
+                            event.command_id, None
+                        ):
+                            try:
+                                await queue.send(
+                                    PrefillProgressData(
+                                        processed_tokens=event.processed_tokens,
+                                        total_tokens=event.total_tokens,
+                                    )
+                                )
+                            except BrokenResourceError:
+                                self._text_generation_queues.pop(event.command_id, None)
+
                     if isinstance(event, TracesMerged):
                         self._save_merged_trace(event)
 

src/exo/master/main.py

@@ -24,7 +24,6 @@ from exo.shared.types.commands import (
     PlaceInstance,
     RequestEventLog,
     SendInputChunk,
-    TaskCancelled,
     TaskFinished,
     TestCommand,
     TextGeneration,
@@ -40,7 +39,6 @@ from exo.shared.types.events import (
     NodeTimedOut,
     TaskCreated,
     TaskDeleted,
-    TaskStatusUpdated,
     TraceEventData,
     TracesCollected,
     TracesMerged,
@@ -281,7 +279,7 @@ class Master:
                         case DeleteInstance():
                             placement = delete_instance(command, self.state.instances)
                             transition_events = get_transition_events(
-                                self.state.instances, placement, self.state.tasks
+                                self.state.instances, placement
                             )
                             for cmd in cancel_unnecessary_downloads(
                                 placement, self.state.downloads
@@ -301,7 +299,7 @@ class Master:
                                 self.state.node_network,
                             )
                             transition_events = get_transition_events(
-                                self.state.instances, placement, self.state.tasks
+                                self.state.instances, placement
                             )
                             generated_events.extend(transition_events)
                         case CreateInstance():
@@ -311,7 +309,7 @@ class Master:
                                 self.state.instances,
                             )
                             transition_events = get_transition_events(
-                                self.state.instances, placement, self.state.tasks
+                                self.state.instances, placement
                             )
                             generated_events.extend(transition_events)
                         case SendInputChunk(chunk=chunk):
@@ -321,18 +319,6 @@ class Master:
                                     chunk=chunk,
                                 )
                             )
-                        case TaskCancelled():
-                            if (
-                                task_id := self.command_task_mapping.get(
-                                    command.cancelled_command_id
-                                )
-                            ) is not None:
-                                generated_events.append(
-                                    TaskStatusUpdated(
-                                        task_status=TaskStatus.Cancelled,
-                                        task_id=task_id,
-                                    )
-                                )
                         case TaskFinished():
                             generated_events.append(
                                 TaskDeleted(
@@ -341,9 +327,10 @@ class Master:
                                     ]
                                 )
                             )
-                            self.command_task_mapping.pop(
-                                command.finished_command_id, None
-                            )
+                            if command.finished_command_id in self.command_task_mapping:
+                                del self.command_task_mapping[
+                                    command.finished_command_id
+                                ]
                         case RequestEventLog():
                             # We should just be able to send everything, since other buffers will ignore old messages
                             # rate limit to 1000 at a time

src/exo/master/placement.py

@@ -22,15 +22,9 @@ from exo.shared.types.commands import (
     PlaceInstance,
 )
 from exo.shared.types.common import NodeId
-from exo.shared.types.events import (
-    Event,
-    InstanceCreated,
-    InstanceDeleted,
-    TaskStatusUpdated,
-)
+from exo.shared.types.events import Event, InstanceCreated, InstanceDeleted
 from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
-from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import (
     DownloadOngoing,
     DownloadProgress,
@@ -192,7 +186,6 @@ def delete_instance(
 def get_transition_events(
     current_instances: Mapping[InstanceId, Instance],
     target_instances: Mapping[InstanceId, Instance],
-    tasks: Mapping[TaskId, Task],
 ) -> Sequence[Event]:
     events: list[Event] = []
 
@@ -208,18 +201,6 @@ def get_transition_events(
     # find instances to delete
     for instance_id in current_instances:
         if instance_id not in target_instances:
-            for task in tasks.values():
-                if task.instance_id == instance_id and task.task_status in [
-                    TaskStatus.Pending,
-                    TaskStatus.Running,
-                ]:
-                    events.append(
-                        TaskStatusUpdated(
-                            task_status=TaskStatus.Cancelled,
-                            task_id=task.task_id,
-                        )
-                    )
-
             events.append(
                 InstanceDeleted(
                     instance_id=instance_id,

src/exo/master/tests/test_claude_tool_use.py

@@ -4,11 +4,7 @@ import json
 from collections.abc import AsyncGenerator
 from typing import Any, cast
 
-from exo.master.adapters.claude import (
-    ClaudeMessagesResponse,
-    collect_claude_response,
-    generate_claude_stream,
-)
+from exo.master.adapters.claude import collect_claude_response, generate_claude_stream
 from exo.shared.types.api import ToolCallItem
 from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
 from exo.shared.types.common import CommandId, ModelId
@@ -21,18 +17,6 @@ async def _chunks_to_stream(
         yield chunk
 
 
-async def _collect_response(
-    command_id: CommandId,
-    model: str,
-    chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
-) -> ClaudeMessagesResponse:
-    """Helper to consume the async generator and parse the JSON response."""
-    parts: list[str] = []
-    async for part in collect_claude_response(command_id, model, chunk_stream):
-        parts.append(part)
-    return ClaudeMessagesResponse.model_validate_json("".join(parts))
-
-
 MODEL = ModelId("test-model")
 COMMAND_ID = CommandId("cmd_test123")
 
@@ -63,7 +47,7 @@ class TestCollectClaudeResponseToolUse:
                 ],
             ),
         ]
-        response = await _collect_response(
+        response = await collect_claude_response(
             COMMAND_ID, "test-model", _chunks_to_stream(chunks)
         )
 
@@ -93,7 +77,7 @@ class TestCollectClaudeResponseToolUse:
                 ],
             ),
         ]
-        response = await _collect_response(
+        response = await collect_claude_response(
             COMMAND_ID, "test-model", _chunks_to_stream(chunks)
         )
 
@@ -118,7 +102,7 @@ class TestCollectClaudeResponseToolUse:
                 ],
             ),
         ]
-        response = await _collect_response(
+        response = await collect_claude_response(
             COMMAND_ID, "test-model", _chunks_to_stream(chunks)
         )
 
@@ -132,7 +116,7 @@ class TestCollectClaudeResponseToolUse:
 
     async def test_no_content_produces_empty_text_block(self):
         chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = []
-        response = await _collect_response(
+        response = await collect_claude_response(
             COMMAND_ID, "test-model", _chunks_to_stream(chunks)
         )
         assert len(response.content) == 1

src/exo/master/tests/test_master.py

@@ -42,7 +42,7 @@ from exo.utils.channels import channel
 @pytest.mark.asyncio
 async def test_master():
     keypair = get_node_id_keypair()
-    node_id = NodeId(keypair.to_node_id())
+    node_id = NodeId(keypair.to_peer_id().to_base58())
     session_id = SessionId(master_node_id=node_id, election_clock=0)
 
     ge_sender, global_event_receiver = channel[ForwarderEvent]()

src/exo/master/tests/test_placement.py

@@ -239,7 +239,7 @@ def test_get_transition_events_no_change(instance: Instance):
     target_instances = {instance_id: instance}
 
     # act
-    events = get_transition_events(current_instances, target_instances, {})
+    events = get_transition_events(current_instances, target_instances)
 
     # assert
     assert len(events) == 0
@@ -252,7 +252,7 @@ def test_get_transition_events_create_instance(instance: Instance):
     target_instances: dict[InstanceId, Instance] = {instance_id: instance}
 
     # act
-    events = get_transition_events(current_instances, target_instances, {})
+    events = get_transition_events(current_instances, target_instances)
 
     # assert
     assert len(events) == 1
@@ -266,7 +266,7 @@ def test_get_transition_events_delete_instance(instance: Instance):
     target_instances: dict[InstanceId, Instance] = {}
 
     # act
-    events = get_transition_events(current_instances, target_instances, {})
+    events = get_transition_events(current_instances, target_instances)
 
     # assert
     assert len(events) == 1

src/exo/routing/connection_message.py

@@ -30,7 +30,7 @@ class ConnectionMessage(CamelCaseModel):
     @classmethod
     def from_update(cls, update: ConnectionUpdate) -> "ConnectionMessage":
         return cls(
-            node_id=NodeId(update.peer_id),
+            node_id=NodeId(update.peer_id.to_base58()),
             connection_type=ConnectionMessageType.from_update_type(update.update_type),
             remote_ipv4=update.remote_ipv4,
             remote_tcp_port=update.remote_tcp_port,

src/exo/routing/router.py

@@ -221,7 +221,7 @@ def get_node_id_keypair(
     Obtain the :class:`PeerId` by from it.
     """
     # TODO(evan): bring back node id persistence once we figure out how to deal with duplicates
-    return Keypair.generate()
+    return Keypair.generate_ed25519()
 
     def lock_path(path: str | bytes | PathLike[str] | PathLike[bytes]) -> Path:
         return Path(str(path) + ".lock")
@@ -235,12 +235,12 @@ def get_node_id_keypair(
                 protobuf_encoded = f.read()
 
                 try:  # if decoded successfully, save & return
-                    return Keypair.from_bytes(protobuf_encoded)
+                    return Keypair.from_protobuf_encoding(protobuf_encoded)
                 except ValueError as e:  # on runtime error, assume corrupt file
                     logger.warning(f"Encountered error when trying to get keypair: {e}")
 
         # if no valid credentials, create new ones and persist
         with open(path, "w+b") as f:
             keypair = Keypair.generate_ed25519()
-            f.write(keypair.to_bytes())
+            f.write(keypair.to_protobuf_encoding())
             return keypair

src/exo/shared/apply.py

@@ -15,6 +15,7 @@ from exo.shared.types.events import (
     NodeDownloadProgress,
     NodeGatheredInfo,
     NodeTimedOut,
+    PrefillProgress,
     RunnerDeleted,
     RunnerStatusUpdated,
     TaskAcknowledged,
@@ -64,6 +65,7 @@ def event_apply(event: Event, state: State) -> State:
             | ChunkGenerated()
             | TaskAcknowledged()
             | InputChunkReceived()
+            | PrefillProgress()
             | TracesCollected()
             | TracesMerged()
         ):  # Pass-through events that don't modify state
@@ -218,6 +220,11 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
         key: value for key, value in state.downloads.items() if key != event.node_id
     }
     # Clean up all granular node mappings
+    node_identities = {
+        key: value
+        for key, value in state.node_identities.items()
+        if key != event.node_id
+    }
     node_memory = {
         key: value for key, value in state.node_memory.items() if key != event.node_id
     }
@@ -258,6 +265,7 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
             "downloads": downloads,
             "topology": topology,
             "last_seen": last_seen,
+            "node_identities": node_identities,
             "node_memory": node_memory,
             "node_disk": node_disk,
             "node_system": node_system,

src/exo/shared/models/model_cards.py

@@ -44,8 +44,7 @@ async def _refresh_card_cache():
         async for toml_file in path.rglob("*.toml"):
             try:
                 card = await ModelCard.load_from_path(toml_file)
-                if card.model_id not in _card_cache:
-                    _card_cache[card.model_id] = card
+                _card_cache[card.model_id] = card
             except (ValidationError, TOMLKitError):
                 pass
 
@@ -183,7 +182,6 @@ class ConfigData(BaseModel):
     def supports_tensor(self) -> bool:
         return self.architectures in [
             ["Glm4MoeLiteForCausalLM"],
-            ["GlmMoeDsaForCausalLM"],
             ["DeepseekV32ForCausalLM"],
             ["DeepseekV3ForCausalLM"],
             ["Qwen3NextForCausalLM"],

src/exo/shared/tests/test_node_id_persistence.py

@@ -23,7 +23,7 @@ def _get_keypair_concurrent_subprocess_task(
     sem.release()
     # wait to be told to begin simultaneous read
     ev.wait()
-    queue.put(get_node_id_keypair().to_bytes())
+    queue.put(get_node_id_keypair().to_protobuf_encoding())
 
 
 def _get_keypair_concurrent(num_procs: int) -> bytes:

src/exo/shared/types/api.py

@@ -1,9 +1,10 @@
 import time
 from collections.abc import Generator
-from typing import Annotated, Any, Literal, get_args
+from typing import Annotated, Any, Literal
 from uuid import uuid4
 
 from pydantic import BaseModel, Field, field_validator
+from pydantic_core import PydanticUseDefault
 
 from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.common import CommandId, NodeId
@@ -227,6 +228,13 @@ class PlaceInstanceParams(BaseModel):
     instance_meta: InstanceMeta = InstanceMeta.MlxRing
     min_nodes: int = 1
 
+    @field_validator("sharding", "instance_meta", mode="plain")
+    @classmethod
+    def use_default(cls, v: object):
+        if not v or not isinstance(v, (Sharding, InstanceMeta)):
+            raise PydanticUseDefault()
+        return v
+
 
 class CreateInstanceParams(BaseModel):
     instance: Instance
@@ -262,27 +270,6 @@ class DeleteInstanceResponse(BaseModel):
     instance_id: InstanceId
 
 
-ImageSize = Literal[
-    "auto",
-    "512x512",
-    "768x768",
-    "1024x768",
-    "768x1024",
-    "1024x1024",
-    "1024x1536",
-    "1536x1024",
-]
-
-
-def normalize_image_size(v: object) -> ImageSize:
-    """Shared validator for ImageSize fields: maps None → "auto" and rejects invalid values."""
-    if v is None:
-        return "auto"
-    if v not in get_args(ImageSize):
-        raise ValueError(f"Invalid size: {v!r}. Must be one of {get_args(ImageSize)}")
-    return v  # pyright: ignore[reportReturnType]
-
-
 class AdvancedImageParams(BaseModel):
     seed: Annotated[int, Field(ge=0)] | None = None
     num_inference_steps: Annotated[int, Field(ge=1, le=100)] | None = None
@@ -302,7 +289,7 @@ class ImageGenerationTaskParams(BaseModel):
     partial_images: int | None = 0
     quality: Literal["high", "medium", "low"] | None = "medium"
     response_format: Literal["url", "b64_json"] | None = "b64_json"
-    size: ImageSize = "auto"
+    size: str | None = "1024x1024"
     stream: bool | None = False
     style: str | None = "vivid"
     user: str | None = None
@@ -310,11 +297,6 @@ class ImageGenerationTaskParams(BaseModel):
     # Internal flag for benchmark mode - set by API, preserved through serialization
     bench: bool = False
 
-    @field_validator("size", mode="before")
-    @classmethod
-    def normalize_size(cls, v: object) -> ImageSize:
-        return normalize_image_size(v)
-
 
 class BenchImageGenerationTaskParams(ImageGenerationTaskParams):
     bench: bool = True
@@ -331,18 +313,13 @@ class ImageEditsTaskParams(BaseModel):
     quality: Literal["high", "medium", "low"] | None = "medium"
     output_format: Literal["png", "jpeg", "webp"] = "png"
     response_format: Literal["url", "b64_json"] | None = "b64_json"
-    size: ImageSize = "auto"
+    size: str | None = "1024x1024"
     image_strength: float | None = 0.7
     stream: bool = False
     partial_images: int | None = 0
     advanced_params: AdvancedImageParams | None = None
     bench: bool = False
 
-    @field_validator("size", mode="before")
-    @classmethod
-    def normalize_size(cls, v: object) -> ImageSize:
-        return normalize_image_size(v)
-
     def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
         for name, value in super().__repr_args__():  # pyright: ignore[reportAny]
             if name == "image_data":

src/exo/shared/types/chunks.py

@@ -77,3 +77,13 @@ class InputImageChunk(BaseChunk):
 
 
 GenerationChunk = TokenChunk | ImageChunk | ToolCallChunk | ErrorChunk
+
+
+class PrefillProgressData(TaggedModel):
+    """Data class for prefill progress events during streaming."""
+
+    processed_tokens: int
+    total_tokens: int
+
+
+StreamEvent = TokenChunk | PrefillProgressData

src/exo/shared/types/commands.py

@@ -48,10 +48,6 @@ class DeleteInstance(BaseCommand):
     instance_id: InstanceId
 
 
-class TaskCancelled(BaseCommand):
-    cancelled_command_id: CommandId
-
-
 class TaskFinished(BaseCommand):
     finished_command_id: CommandId
 
@@ -93,7 +89,6 @@ Command = (
     | PlaceInstance
     | CreateInstance
     | DeleteInstance
-    | TaskCancelled
     | TaskFinished
     | SendInputChunk
 )

src/exo/shared/types/events.py

@@ -102,6 +102,12 @@ class InputChunkReceived(BaseEvent):
     chunk: InputImageChunk
 
 
+class PrefillProgress(BaseEvent):
+    command_id: CommandId
+    processed_tokens: int
+    total_tokens: int
+
+
 class TopologyEdgeCreated(BaseEvent):
     conn: Connection
 
@@ -148,6 +154,7 @@ Event = (
     | NodeDownloadProgress
     | ChunkGenerated
     | InputChunkReceived
+    | PrefillProgress
     | TopologyEdgeCreated
     | TopologyEdgeDeleted
     | TracesCollected

src/exo/shared/types/tasks.py

@@ -24,7 +24,6 @@ class TaskStatus(str, Enum):
     Complete = "Complete"
     TimedOut = "TimedOut"
     Failed = "Failed"
-    Cancelled = "Cancelled"
 
 
 class BaseTask(TaggedModel):
@@ -61,11 +60,6 @@ class TextGeneration(BaseTask):  # emitted by Master
     error_message: str | None = Field(default=None)
 
 
-class CancelTask(BaseTask):
-    cancelled_task_id: TaskId
-    runner_id: RunnerId
-
-
 class ImageGeneration(BaseTask):  # emitted by Master
     command_id: CommandId
     task_params: ImageGenerationTaskParams
@@ -93,7 +87,6 @@ Task = (
     | LoadModel
     | StartWarmup
     | TextGeneration
-    | CancelTask
     | ImageGeneration
     | ImageEdits
     | Shutdown

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

@@ -26,7 +26,6 @@ class DownloadProgressData(CamelCaseModel):
 class BaseDownloadProgress(TaggedModel):
     node_id: NodeId
     shard_metadata: ShardMetadata
-    model_directory: str = ""
 
 
 class DownloadPending(BaseDownloadProgress):

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

@@ -62,8 +62,12 @@ class PartialImageResponse(BaseRunnerResponse):
 class ToolCallResponse(BaseRunnerResponse):
     tool_calls: list[ToolCallItem]
     usage: Usage | None
-    stats: GenerationStats | None = None
 
 
 class FinishedResponse(BaseRunnerResponse):
     pass
+
+
+class PrefillProgressResponse(BaseRunnerResponse):
+    processed_tokens: int
+    total_tokens: int

src/exo/utils/banner.py

@@ -1,7 +1,5 @@
-import sys
-
-
 def print_startup_banner(port: int) -> None:
+    """Print a prominent startup banner with API endpoint information."""
     dashboard_url = f"http://localhost:{port}"
     banner = f"""
 ╔═══════════════════════════════════════════════════════════════════════╗
@@ -29,4 +27,4 @@ def print_startup_banner(port: int) -> None:
 
 """
 
-    print(banner, file=sys.stderr)
+    print(banner)

src/exo/utils/channels.py

@@ -125,9 +125,7 @@ class MpSender[T]:
             self._state.buffer.put(item, block=True)
 
     async def send_async(self, item: T) -> None:
-        await to_thread.run_sync(
-            self.send, item, limiter=CapacityLimiter(1), abandon_on_cancel=True
-        )
+        await to_thread.run_sync(self.send, item, limiter=CapacityLimiter(1))
 
     def close(self) -> None:
         if not self._state.closed.is_set():

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

@@ -14,7 +14,6 @@ from exo.shared.types.api import (
     ImageEditsTaskParams,
     ImageGenerationStats,
     ImageGenerationTaskParams,
-    ImageSize,
 )
 from exo.shared.types.memory import Memory
 from exo.shared.types.worker.runner_response import (
@@ -24,9 +23,9 @@ from exo.shared.types.worker.runner_response import (
 from exo.worker.engines.image.distributed_model import DistributedImageModel
 
 
-def parse_size(size_str: ImageSize) -> tuple[int, int]:
+def parse_size(size_str: str | None) -> tuple[int, int]:
     """Parse size parameter like '1024x1024' to (width, height) tuple."""
-    if size_str == "auto":
+    if not size_str:
         return (1024, 1024)
 
     try:
@@ -110,9 +109,6 @@ def generate_image(
             # Decode base64 image data and save to temp file
             image_path = Path(tmpdir) / "input.png"
             image_path.write_bytes(base64.b64decode(task.image_data))
-            if task.size == "auto":
-                with Image.open(image_path) as img:
-                    width, height = img.size
 
         for image_num in range(num_images):
             # Increment seed for each image to ensure unique results

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

@@ -163,14 +163,11 @@ class PipelineLastLayer(CustomMlxLayer):
                 output, (self.r + 1) % self.s, group=self.group
             )
             if cache is not None:
-                # CacheList (used by MLA models like DeepSeekV32, GLM MoE DSA)
-                # doesn't have .keys directly; access via first sub-cache.
-                _cache = cache[0] if hasattr(cache, "caches") else cache  # type: ignore
-                _cache.keys = mx.depends(_cache.keys, output)  # type: ignore
+                cache.keys = mx.depends(cache.keys, output)  # type: ignore[reportUnknownMemberType]
             if self.is_prefill:
                 mx.eval(output)
                 if cache is not None:
-                    mx.eval(_cache.keys)  # type: ignore
+                    mx.eval(cache.keys)  # type: ignore
 
         if not self.is_prefill:
             output = mx.distributed.all_gather(output, group=self.group)[
@@ -310,9 +307,7 @@ def patch_pipeline_model[T](model: T, group: mx.distributed.Group) -> T:
 
         # Add dependency to last cache entry to ensure distributed ops are evaluated
         if cache is not None:
-            last = cache[-1]  # type: ignore
-            dep_cache = last[0] if hasattr(last, "caches") else last  # type: ignore
-            dep_cache.keys = mx.depends(dep_cache.keys, logits)  # type: ignore
+            cache[-1].state = mx.depends(cache[-1].state, logits)  # type: ignore
 
         return logits
 
@@ -338,9 +333,7 @@ def patch_tensor_model[T](model: T) -> T:
 
         # Add dependency to last cache entry to ensure distributed ops are evaluated
         if cache is not None and len(cache) > 0:  # pyright: ignore[reportAny]
-            last = cache[-1]  # pyright: ignore[reportAny]
-            dep_cache = last[0] if hasattr(last, "caches") else last  # pyright: ignore[reportAny]
-            dep_cache.keys = mx.depends(dep_cache.keys, logits)  # pyright: ignore[reportAny,reportUnknownMemberType]
+            cache[-1].state = mx.depends(cache[-1].state, logits)  # pyright: ignore[reportAny,reportUnknownMemberType]
 
         return logits
 
@@ -554,12 +547,10 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
         on_timeout: TimeoutCallback | None,
     ) -> nn.Module:
         model = cast(DeepseekV3Model, model)
-
         for layer in model.layers:
             eval_with_timeout(
                 layer.parameters(), timeout_seconds / len(model.layers), on_timeout
             )
-
             # Shard the self attention
             if layer.self_attn.q_lora_rank is None:
                 layer.self_attn.q_proj = self.all_to_sharded_linear(
@@ -590,18 +581,12 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
                 layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj)
                 layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
 
-            # Shard the MoE.
+            # Shard the MoE. Shard in place since the MoE should be responsible
+            # for aggregating the results.
             else:
-                if getattr(layer.mlp, "shared_experts", None) is not None:
-                    self.all_to_sharded_linear_in_place(
-                        layer.mlp.shared_experts.gate_proj
-                    )
-                    self.sharded_to_all_linear_in_place(
-                        layer.mlp.shared_experts.down_proj
-                    )
-                    self.all_to_sharded_linear_in_place(
-                        layer.mlp.shared_experts.up_proj
-                    )
+                self.all_to_sharded_linear_in_place(layer.mlp.shared_experts.gate_proj)
+                self.sharded_to_all_linear_in_place(layer.mlp.shared_experts.down_proj)
+                self.all_to_sharded_linear_in_place(layer.mlp.shared_experts.up_proj)
                 self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj)
                 self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj)
                 self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj)
@@ -794,7 +779,8 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
 
             layer.self_attn = WrappedMiniMaxAttention(layer.self_attn, self.group)  # pyright: ignore[reportAttributeAccessIssue,reportArgumentType]
 
-            # Shard the MoE.
+            # Shard the MoE. Shard in place since the MoE should be responsible
+            # for aggregating the results.
             self.all_to_sharded_linear_in_place(
                 layer.block_sparse_moe.switch_mlp.gate_proj
             )
@@ -907,7 +893,8 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                     layer.self_attn.num_attention_heads //= self.N
                     layer.self_attn.num_key_value_heads //= self.N
 
-            # Shard the MoE.
+            # Shard the MoE. Shard in place since the MoE should be responsible
+            # for aggregating the results.
             if isinstance(layer.mlp, (Qwen3MoeSparseMoeBlock, Qwen3NextSparseMoeBlock)):
                 self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj)
                 self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj)

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

@@ -57,7 +57,6 @@ def prefill(
     sampler: Callable[[mx.array], mx.array],
     prompt_tokens: mx.array,
     cache: KVCacheType,
-    group: mx.distributed.Group | None = None,
 ) -> tuple[float, int, list[CacheSnapshot]]:
     """Prefill the KV cache with prompt tokens.
 
@@ -87,9 +86,6 @@ def prefill(
 
     set_pipeline_prefill(model, is_prefill=True)
 
-    mx_barrier(group)
-    logger.info("Starting prefill")
-
     # Use max_tokens=1 because max_tokens=0 does not work.
     # We just throw away the generated token - we only care about filling the cache
     for _ in stream_generate(
@@ -164,7 +160,7 @@ def warmup_inference(
         max_tokens=50,
         sampler=sampler,
         prompt_cache=cache,
-        prefill_step_size=2048,
+        prefill_step_size=1024,
         kv_group_size=KV_GROUP_SIZE,
         kv_bits=KV_BITS,
     ):
@@ -256,6 +252,7 @@ def mlx_generate(
     task: TextGenerationTaskParams,
     prompt: str,
     kv_prefix_cache: KVPrefixCache | None = None,
+    on_prefill_progress: Callable[[int, int], None] | None = None,
     group: mx.distributed.Group | None = None,
 ) -> Generator[GenerationResponse]:
     # Ensure that generation stats only contains peak memory for this generation
@@ -309,9 +306,16 @@ def mlx_generate(
     )
     max_stop_len = max((len(s) for s in stop_sequences), default=0)
 
+    mx_barrier(group)
+    logger.info("Ready to prefill")
+
     # Prefill cache with all tokens except the last one
     prefill_tps, prefill_tokens, ssm_snapshots_list = prefill(
-        model, tokenizer, sampler, prompt_tokens[:-1], caches, group
+        model,
+        tokenizer,
+        sampler,
+        prompt_tokens[:-1],
+        caches,
     )
     cache_snapshots: list[CacheSnapshot] | None = ssm_snapshots_list or None
 
@@ -328,7 +332,6 @@ def mlx_generate(
     think_start = tokenizer.think_start
     think_end = tokenizer.think_end
 
-    logger.info("Starting decode")
     mx_barrier(group)
 
     for completion_tokens, out in enumerate(
@@ -343,6 +346,7 @@ def mlx_generate(
             prefill_step_size=1,
             kv_group_size=KV_GROUP_SIZE,
             kv_bits=KV_BITS,
+            prompt_progress_callback=on_prefill_progress,
         ),
         start=1,
     ):
@@ -391,11 +395,10 @@ def mlx_generate(
                     f"Model generated unexpected finish_reason: {out.finish_reason}"
                 )
 
-            total_prompt_tokens = len(all_prompt_tokens)
             usage = Usage(
-                prompt_tokens=total_prompt_tokens,
+                prompt_tokens=int(out.prompt_tokens),
                 completion_tokens=completion_tokens,
-                total_tokens=total_prompt_tokens + completion_tokens,
+                total_tokens=int(out.prompt_tokens) + completion_tokens,
                 prompt_tokens_details=PromptTokensDetails(
                     cached_tokens=prefix_hit_length
                 ),

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

@@ -64,6 +64,8 @@ from exo.worker.runner.bootstrap import logger
 Group = mx.distributed.Group
 
 
+# TODO: Test this
+#  ALSO https://github.com/exo-explore/exo/pull/233#discussion_r2549683673
 def get_weights_size(model_shard_meta: ShardMetadata) -> Memory:
     return Memory.from_float_kb(
         (model_shard_meta.end_layer - model_shard_meta.start_layer)
@@ -81,6 +83,30 @@ class ModelLoadingTimeoutError(Exception):
     pass
 
 
+def mx_barrier(group: Group | None = None):
+    mx.eval(
+        mx.distributed.all_sum(
+            mx.array(1.0),
+            stream=mx.default_stream(mx.Device(mx.cpu)),
+            group=group,
+        )
+    )
+
+
+def broadcast_from_zero(value: int, group: Group | None = None):
+    if group is None:
+        return value
+
+    if group.rank() == 0:
+        a = mx.array([value], dtype=mx.int32)
+    else:
+        a = mx.array([0], dtype=mx.int32)
+
+    m = mx.distributed.all_sum(a, stream=mx.Device(mx.DeviceType.cpu), group=group)
+    mx.eval(m)
+    return int(m.item())
+
+
 class HostList(RootModel[list[str]]):
     @classmethod
     def from_hosts(cls, hosts: list[Host]) -> "HostList":
@@ -285,12 +311,10 @@ def get_eos_token_ids_for_model(model_id: ModelId) -> list[int] | None:
     model_id_lower = model_id.lower()
     if "kimi-k2" in model_id_lower:
         return [163586]
-    elif "glm-5" in model_id_lower or "glm-4.7" in model_id_lower:
-        # For GLM-5 and GLM-4.7
+    elif "glm-4.7-flash" in model_id_lower:
         # 154820: <|endoftext|>, 154827: <|user|>, 154829: <|observation|>
         return [154820, 154827, 154829]
     elif "glm" in model_id_lower:
-        # For GLM-4.5 and older
         return [151336, 151329, 151338]
     return None
 
@@ -355,13 +379,7 @@ def load_tokenizer_for_model_id(
             return list(hf_tokenizer.model.encode(text, allowed_special="all"))  # pyright: ignore[reportUnknownMemberType,reportUnknownArgumentType]
 
         hf_tokenizer.encode = _patched_encode
-        return TokenizerWrapper(
-            hf_tokenizer,
-            eos_token_ids=eos_token_ids,
-            tool_call_start="<|tool_calls_section_begin|>",
-            tool_call_end="<|tool_calls_section_end|>",
-            tool_parser=_parse_kimi_tool_calls,
-        )
+        return TokenizerWrapper(hf_tokenizer, eos_token_ids=eos_token_ids)
 
     tokenizer = load_tokenizer(
         model_path,
@@ -573,61 +591,3 @@ def mlx_cleanup(
     import gc
 
     gc.collect()
-
-
-def mx_any(bool_: bool, group: Group | None) -> bool:
-    if group is None:
-        return bool_
-    num_true = mx.distributed.all_sum(
-        mx.array(bool_), group=group, stream=mx.default_stream(mx.Device(mx.cpu))
-    )
-    mx.eval(num_true)
-    return num_true.item() > 0
-
-
-def mx_barrier(group: Group | None):
-    if group is None:
-        return
-    mx.eval(
-        mx.distributed.all_sum(
-            mx.array(1.0), group=group, stream=mx.default_stream(mx.Device(mx.cpu))
-        )
-    )
-
-
-def _parse_kimi_tool_calls(text: str):
-    import regex as re
-
-    # kimi has a fixed function naming scheme, with a json formatted arg
-    #   functions.multiply:0<|tool_call_argument_begin|>{"a": 2, "b": 3}
-    _func_name_regex = re.compile(
-        r"^\s*((?:functions\.)?(.+?):\d+)\s*<\|tool_call_argument_begin\|>", re.DOTALL
-    )
-    _func_arg_regex = re.compile(r"<\|tool_call_argument_begin\|>\s*(.*)\s*", re.DOTALL)
-    _tool_call_split_regex = re.compile(
-        r"<\|tool_call_begin\|>(.*?)<\|tool_call_end\|>", re.DOTALL
-    )
-
-    def _parse_single_tool(text: str) -> dict[str, Any]:
-        func_name_match = _func_name_regex.search(text)
-        if func_name_match is None:
-            raise ValueError("No tool call found.")
-        tool_call_id = func_name_match.group(1)  # e.g. "functions.get_weather:0"
-        func_name = func_name_match.group(2)  # e.g. "get_weather"
-
-        func_args_match = _func_arg_regex.search(text)
-        if func_args_match is None:
-            raise ValueError("No tool call arguments found.")
-        func_args = func_args_match.group(1)
-        try:
-            arg_dct = json.loads(func_args)  # pyright: ignore[reportAny]
-        except Exception:
-            arg_dct = None
-
-        return dict(id=tool_call_id, name=func_name, arguments=arg_dct)
-
-    tool_matches = _tool_call_split_regex.findall(text)
-    if tool_matches:
-        return [_parse_single_tool(match) for match in tool_matches]  # pyright: ignore[reportAny]
-    else:
-        return [_parse_single_tool(text)]

src/exo/worker/main.py

@@ -33,7 +33,6 @@ from exo.shared.types.events import (
 from exo.shared.types.multiaddr import Multiaddr
 from exo.shared.types.state import State
 from exo.shared.types.tasks import (
-    CancelTask,
     CreateRunner,
     DownloadModel,
     ImageEdits,
@@ -225,22 +224,15 @@ class Worker:
                         )
                     )
                 case Shutdown(runner_id=runner_id):
-                    runner = self.runners.pop(runner_id)
                     try:
                         with fail_after(3):
-                            await runner.start_task(task)
+                            await self.runners.pop(runner_id).start_task(task)
                     except TimeoutError:
                         await self.event_sender.send(
                             TaskStatusUpdated(
                                 task_id=task.task_id, task_status=TaskStatus.TimedOut
                             )
                         )
-                    finally:
-                        runner.shutdown()
-                case CancelTask(
-                    cancelled_task_id=cancelled_task_id, runner_id=runner_id
-                ):
-                    await self.runners[runner_id].cancel_task(cancelled_task_id)
                 case ImageEdits() if task.task_params.total_input_chunks > 0:
                     # Assemble image from chunks and inject into task
                     cmd_id = task.command_id
@@ -278,18 +270,18 @@ class Worker:
                         del self.input_chunk_buffer[cmd_id]
                     if cmd_id in self.input_chunk_counts:
                         del self.input_chunk_counts[cmd_id]
-                    await self._start_runner_task(modified_task)
+                    await self.runners[self._task_to_runner_id(task)].start_task(
+                        modified_task
+                    )
                 case task:
-                    await self._start_runner_task(task)
+                    await self.runners[self._task_to_runner_id(task)].start_task(task)
 
     def shutdown(self):
         self._tg.cancel_scope.cancel()
 
-    async def _start_runner_task(self, task: Task):
-        if (instance := self.state.instances.get(task.instance_id)) is not None:
-            await self.runners[
-                instance.shard_assignments.node_to_runner[self.node_id]
-            ].start_task(task)
+    def _task_to_runner_id(self, task: Task):
+        instance = self.state.instances[task.instance_id]
+        return instance.shard_assignments.node_to_runner[self.node_id]
 
     async def _nack_request(self, since_idx: int) -> None:
         # We request all events after (and including) the missing index.
@@ -328,6 +320,8 @@ class Worker:
             for event in self.out_for_delivery.copy().values():
                 await self.local_event_sender.send(event)
 
+    ## Op Executors
+
     def _create_supervisor(self, task: CreateRunner) -> RunnerSupervisor:
         """Creates and stores a new AssignedRunner with initial downloading status."""
         runner = RunnerSupervisor.create(