small refactor

CONTRIBUTING.md

@@ -29,7 +29,6 @@ To run EXO from source:
 git clone https://github.com/exo-explore/exo.git
 cd exo/dashboard
 npm install && npm run build && cd ..
-uv sync --extra mlx
 uv run exo
 ```
 

Cargo.lock

@@ -1007,6 +1007,7 @@ dependencies = [
  "serde_json",
  "tokio",
  "zenoh",
+ "zenoh-ext",
 ]
 
 [[package]]
@@ -1842,6 +1843,12 @@ dependencies = [
  "spin 0.9.8",
 ]
 
+[[package]]
+name = "leb128"
+version = "0.2.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6cc46bac87ef8093eed6f272babb833b6443374399985ac8ed28471ee0918545"
+
 [[package]]
 name = "leb128fmt"
 version = "0.1.0"
@@ -5107,7 +5114,7 @@ dependencies = [
 [[package]]
 name = "zenoh"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "ahash",
  "arc-swap",
@@ -5157,7 +5164,7 @@ dependencies = [
 [[package]]
 name = "zenoh-buffers"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "zenoh-collections",
 ]
@@ -5165,7 +5172,7 @@ dependencies = [
 [[package]]
 name = "zenoh-codec"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "tracing",
  "uhlc",
@@ -5176,7 +5183,7 @@ dependencies = [
 [[package]]
 name = "zenoh-collections"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "ahash",
 ]
@@ -5184,7 +5191,7 @@ dependencies = [
 [[package]]
 name = "zenoh-config"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "json5",
  "nonempty-collections",
@@ -5209,7 +5216,7 @@ dependencies = [
 [[package]]
 name = "zenoh-core"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "lazy_static",
  "tokio",
@@ -5220,7 +5227,7 @@ dependencies = [
 [[package]]
 name = "zenoh-crypto"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "aes",
  "hmac",
@@ -5230,10 +5237,29 @@ dependencies = [
  "zenoh-result",
 ]
 
+[[package]]
+name = "zenoh-ext"
+version = "1.9.0"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
+dependencies = [
+ "async-trait",
+ "bincode",
+ "flume",
+ "futures",
+ "leb128",
+ "serde",
+ "tokio",
+ "tracing",
+ "uhlc",
+ "zenoh",
+ "zenoh-macros",
+ "zenoh-util",
+]
+
 [[package]]
 name = "zenoh-keyexpr"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "getrandom 0.2.17",
  "hashbrown 0.16.1",
@@ -5248,7 +5274,7 @@ dependencies = [
 [[package]]
 name = "zenoh-link"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "zenoh-config",
  "zenoh-link-commons",
@@ -5266,7 +5292,7 @@ dependencies = [
 [[package]]
 name = "zenoh-link-commons"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "async-trait",
  "base64",
@@ -5302,7 +5328,7 @@ dependencies = [
 [[package]]
 name = "zenoh-link-quic"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "async-trait",
  "rustls-webpki",
@@ -5318,7 +5344,7 @@ dependencies = [
 [[package]]
 name = "zenoh-link-quic_datagram"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "async-trait",
  "rustls-webpki",
@@ -5334,7 +5360,7 @@ dependencies = [
 [[package]]
 name = "zenoh-link-tcp"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "async-trait",
  "socket2 0.5.10",
@@ -5351,7 +5377,7 @@ dependencies = [
 [[package]]
 name = "zenoh-link-tls"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "async-trait",
  "base64",
@@ -5380,7 +5406,7 @@ dependencies = [
 [[package]]
 name = "zenoh-link-udp"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "async-trait",
  "libc",
@@ -5402,7 +5428,7 @@ dependencies = [
 [[package]]
 name = "zenoh-link-unixsock_stream"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "async-trait",
  "nix 0.29.0",
@@ -5420,7 +5446,7 @@ dependencies = [
 [[package]]
 name = "zenoh-link-ws"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "async-trait",
  "futures-util",
@@ -5440,7 +5466,7 @@ dependencies = [
 [[package]]
 name = "zenoh-macros"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -5451,8 +5477,7 @@ dependencies = [
 [[package]]
 name = "zenoh-plugin-storage-manager"
 version = "1.9.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b1d9fd55c490af4dc4abb6d51b620ecf6f824a9285a1f9446b59eeec6d66ccf7"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "async-trait",
  "bincode",
@@ -5479,7 +5504,7 @@ dependencies = [
 [[package]]
 name = "zenoh-plugin-trait"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "git-version",
  "libloading",
@@ -5496,7 +5521,7 @@ dependencies = [
 [[package]]
 name = "zenoh-protocol"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "const_format",
  "rand 0.8.6",
@@ -5511,7 +5536,7 @@ dependencies = [
 [[package]]
 name = "zenoh-result"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "anyhow",
 ]
@@ -5519,7 +5544,7 @@ dependencies = [
 [[package]]
 name = "zenoh-runtime"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "lazy_static",
  "ron",
@@ -5533,7 +5558,7 @@ dependencies = [
 [[package]]
 name = "zenoh-sync"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "arc-swap",
  "event-listener",
@@ -5547,7 +5572,7 @@ dependencies = [
 [[package]]
 name = "zenoh-task"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "futures",
  "tokio",
@@ -5560,7 +5585,7 @@ dependencies = [
 [[package]]
 name = "zenoh-transport"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "async-trait",
  "crossbeam-utils",
@@ -5594,7 +5619,7 @@ dependencies = [
 [[package]]
 name = "zenoh-util"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "async-trait",
  "const_format",
@@ -5620,7 +5645,7 @@ dependencies = [
 [[package]]
 name = "zenoh_backend_traits"
 version = "1.9.0"
-source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
+source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
 dependencies = [
  "async-trait",
  "const_format",

Cargo.toml

@@ -24,8 +24,8 @@ networking = { path = "rust/networking" }
 # pyo3
 pyo3 = "0.28.3"
 pyo3-async-runtimes = "0.28.0"
-pyo3-log = "0.13.3"
-pyo3-stub-gen = "0.22.3"
+pyo3-log = "0.13.2"
+pyo3-stub-gen = "0.22.2"
 
 # util
 extend = "1.2"
@@ -42,6 +42,7 @@ pin-project = "1.1.10"
 serde_json = "1.0.149"
 rand = "0.10.1"
 parking_lot = "0.12.5"
+pidfile-rs = "0.3.1"
 
 # Tracing/logging
 log = "0.4"
@@ -49,18 +50,15 @@ env_logger = "0.11.10"
 
 # networking
 zenoh = "=1.9.0"
+zenoh-ext = "=1.9.0"
 zenoh-plugin-storage-manager = { version = "=1.9.0", default-features = false }
 zenoh-plugin-trait = "=1.9.0"
 netwatcher = "0.6.0"
 bytemuck = "1.25.0"
-blake3 = "1.8.5"
-smol = "2.0.2"
-socket2 = "0.6.4"
-tracing = "0.1.44"
-pidfile-rs = { git = "https://github.com/AndreiCravtov/pidfile-rs" }
 
 [patch.crates-io]
 zenoh = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
+zenoh-ext = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
 zenoh-buffers = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
 zenoh-codec = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
 zenoh-collections = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
@@ -86,6 +84,7 @@ zenoh-sync = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
 zenoh-task = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
 zenoh-transport = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
 zenoh-util = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
+zenoh-plugin-storage-manager = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
 zenoh_backend_traits = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
 
 [workspace.lints.rust]

README.md

@@ -118,7 +118,7 @@ Then restart the Nix daemon: `sudo launchctl kickstart -k system/org.nixos.nix-d
     --force
   ```
 
-Clone the repo, build the dashboard, install the dependencies, and run exo:
+Clone the repo, build the dashboard, and run exo:
 
 ```bash
 # Clone exo
@@ -127,9 +127,6 @@ git clone https://github.com/exo-explore/exo
 # Build dashboard
 cd exo/dashboard && npm install && npm run build && cd ..
 
-# Install Python dependencies, including the MLX backend
-uv sync --extra mlx
-
 # Run exo
 uv run exo
 ```
@@ -179,7 +176,7 @@ rustup toolchain install nightly
 
 **Note:** The `macmon` package is macOS-only and not required for Linux.
 
-Clone the repo, build the dashboard, install the dependencies, and run exo:
+Clone the repo, build the dashboard, and run exo:
 
 ```bash
 # Clone exo
@@ -188,10 +185,6 @@ git clone https://github.com/exo-explore/exo
 # Build dashboard
 cd exo/dashboard && npm install && npm run build && cd ..
 
-# Install Python dependencies with the MLX backend for your hardware
-# (NVIDIA: --extra mlx-cuda13 or --extra mlx-cuda12)
-uv sync --extra mlx-cpu
-
 # Run exo
 uv run exo
 ```
@@ -236,12 +229,6 @@ The macOS app requires macOS Tahoe 26.2 or later.
 
 Download the latest build here: [EXO-latest.dmg](https://assets.exolabs.net/EXO-latest.dmg).
 
-You can also install the latest build with Homebrew:
-
-```bash
-brew install --cask exo
-```
-
 The app will ask for permission to modify system settings and install a new Network profile. Improvements to this are being worked on.
 
 **Custom Namespace for Cluster Isolation:**

TODO.md

@@ -1,13 +1,14 @@
-1. EXO_BOOTSTRAP_PEERS is currently broken
-
+3. Task cancellation. When API http request gets cancelled, it should cancel corresponding task.
 4. I'd like to see profiled network latency / bandwidth.
 5. I'd like to see how much bandwidth each link is using.
 7. Solve the problem of in continuous batching when a new prompt comes in, it will block decode of the current batch until the prefill is complete.
 8. We want people to be able to copy models over to a new device without ever connecting EXO to the internet. Right now EXO require internet connection once to cache some files to check if a download is complete. Instead, we should simply check if there is a non-empty model folder locally with no .partial files. This indicates it's a fully downloaded model that can be loaded.
 13. Memory pressure instead of memory used.
 14. Show the type of each connection (TB5, Ethernet, etc.) in the UI. Refer to old exo: https://github.com/exo-explore/exo/blob/56f783b38dc6b08ce606b07a5386dc40dae00330/exo/helpers.py#L251
+15. Prioritise certain connection types (or by latency). TB5 > Ethernet > WiFi. Refer to old exo: https://github.com/exo-explore/exo/blob/56f783b38dc6b08ce606b07a5386dc40dae00330/exo/helpers.py#L251
 16. Dynamically switch to higher priority connection when it becomes available. Probably bring back InstanceReplacedAtomically.
 17. Faster model loads by streaming model from other devices in cluster.
 18. Add support for specifying the type of network connection to use in a test. Depends on 15/16.
+25. Rethink retry logic
 27. Log cleanup - per-module log filters and default to DEBUG log levels
 28. Validate RDMA connections with ibv_devinfo in the info gatherer

dashboard/package-lock.json

@@ -8,7 +8,6 @@
 			"name": "exo-dashboard",
 			"version": "1.0.0",
 			"dependencies": {
-				"devalue": "^5.6.2",
 				"highlight.js": "^11.11.1",
 				"katex": "^0.16.27",
 				"marked": "^17.0.1",
@@ -2332,9 +2331,9 @@
 			}
 		},
 		"node_modules/devalue": {
-			"version": "5.6.2",
-			"resolved": "https://registry.npmjs.org/devalue/-/devalue-5.6.2.tgz",
-			"integrity": "sha512-nPRkjWzzDQlsejL1WVifk5rvcFi/y1onBRxjaFMjZeR9mFpqu2gmAZ9xUB9/IEanEP/vBtGeGganC/GO1fmufg==",
+			"version": "5.5.0",
+			"resolved": "https://registry.npmjs.org/devalue/-/devalue-5.5.0.tgz",
+			"integrity": "sha512-69sM5yrHfFLJt0AZ9QqZXGCPfJ7fQjvpln3Rq5+PS03LD32Ost1Q9N+eEnaQwGRIriKkMImXD56ocjQmfjbV3w==",
 			"license": "MIT"
 		},
 		"node_modules/enhanced-resolve": {

dashboard/package.json

@@ -11,6 +11,8 @@
 		"check": "svelte-kit sync && svelte-check --tsconfig ./tsconfig.json"
 	},
 	"devDependencies": {
+		"prettier": "^3.4.2",
+		"prettier-plugin-svelte": "^3.3.3",
 		"@sveltejs/adapter-static": "^3.0.10",
 		"@sveltejs/kit": "^2.48.4",
 		"@sveltejs/vite-plugin-svelte": "^5.0.0",
@@ -18,8 +20,6 @@
 		"@types/d3": "^7.4.3",
 		"@types/node": "^22",
 		"d3": "^7.9.0",
-		"prettier": "^3.4.2",
-		"prettier-plugin-svelte": "^3.3.3",
 		"svelte": "^5.0.0",
 		"svelte-check": "^4.0.0",
 		"tailwindcss": "^4.0.0",
@@ -28,7 +28,6 @@
 		"vite": "^6.0.0"
 	},
 	"dependencies": {
-		"devalue": "^5.6.2",
 		"highlight.js": "^11.11.1",
 		"katex": "^0.16.27",
 		"marked": "^17.0.1",

dashboard/src/lib/components/ChatSidebar.svelte

@@ -12,6 +12,8 @@
     toggleDebugMode,
     topologyOnlyMode,
     toggleTopologyOnlyMode,
+    getInstanceFirstShard,
+    type Instance,
   } from "$lib/stores/app.svelte";
 
   interface Props {
@@ -186,7 +188,7 @@
   function extractInstanceModelId(instanceWrapped: unknown): string | null {
     const [, instance] = getTaggedValue(instanceWrapped);
     if (!instance || typeof instance !== "object") return null;
-    const inst = instance as { shardAssignments?: { modelId?: string } };
+    const inst = instance as Instance;
     return inst.shardAssignments?.modelId ?? null;
   }
 
@@ -204,11 +206,7 @@
     else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
 
     let sharding: string | null = null;
-    const inst = instance as {
-      shardAssignments?: { runnerToShard?: Record<string, unknown> };
-    };
-    const runnerToShard = inst.shardAssignments?.runnerToShard || {};
-    const firstShardWrapped = Object.values(runnerToShard)[0];
+    const firstShardWrapped = getInstanceFirstShard(instance as Instance);
     if (firstShardWrapped) {
       const [shardTag] = getTaggedValue(firstShardWrapped);
       if (shardTag === "PipelineShardMetadata") sharding = "Pipeline";

dashboard/src/lib/components/PrefillDecodeDisaggregation.svelte

@@ -9,6 +9,7 @@
     createInstanceLink,
     updateInstanceLink,
     deleteInstanceLink,
+    getInstanceNodeIds,
     type Instance,
   } from "$lib/stores/app.svelte";
   import { deriveBaseModel, deriveFamily } from "$lib/utils/model_family";
@@ -16,7 +17,6 @@
   type InstanceWrapper = {
     MlxRingInstance?: Instance;
     MlxJacclInstance?: Instance;
-    VllmInstance?: Instance;
   };
 
   let interval: ReturnType<typeof setInterval> | null = null;
@@ -43,13 +43,9 @@
     const ids = nodeIdentities();
     for (const [id, raw] of Object.entries(instances())) {
       const wrapper = raw as InstanceWrapper;
-      const inst =
-        wrapper.MlxRingInstance ??
-        wrapper.MlxJacclInstance ??
-        wrapper.VllmInstance;
+      const inst = wrapper.MlxRingInstance ?? wrapper.MlxJacclInstance;
       const modelId = inst?.shardAssignments?.modelId ?? "";
-      const nodeToRunner = inst?.shardAssignments?.nodeToRunner ?? {};
-      const nodeIds = Object.keys(nodeToRunner);
+      const nodeIds = getInstanceNodeIds(inst);
       const nodeNames = nodeIds
         .map((nodeId) => ids[nodeId]?.friendlyName ?? nodeId.slice(0, 6))
         .filter((name) => !!name);

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

@@ -66,12 +66,40 @@ export interface TopologyData {
   edges: TopologyEdge[];
 }
 
+export type InstanceShard = [nodeId: string, runnerId: string, shard: unknown];
+
+export interface ShardAssignments {
+  modelId: string;
+  shards: InstanceShard[];
+  primaryOutputNode: number;
+}
+
 export interface Instance {
-  shardAssignments?: {
-    modelId?: string;
-    runnerToShard?: Record<string, unknown>;
-    nodeToRunner?: Record<string, string>;
-  };
+  shardAssignments: ShardAssignments;
+}
+
+export function getInstanceShards(
+  instance: Instance | null | undefined,
+): InstanceShard[] {
+  return instance?.shardAssignments.shards ?? [];
+}
+
+export function getInstanceRunnerIds(
+  instance: Instance | null | undefined,
+): string[] {
+  return getInstanceShards(instance).map(([, runnerId]) => runnerId);
+}
+
+export function getInstanceNodeIds(
+  instance: Instance | null | undefined,
+): string[] {
+  return [...new Set(getInstanceShards(instance).map(([nodeId]) => nodeId))];
+}
+
+export function getInstanceFirstShard(
+  instance: Instance | null | undefined,
+): unknown {
+  return getInstanceShards(instance)[0]?.[2];
 }
 
 export interface RawInstanceLink {
@@ -918,7 +946,7 @@ class AppStore {
   private extractInstanceModelId(instanceWrapped: unknown): string | null {
     const [, instance] = this.getTaggedValue(instanceWrapped);
     if (!instance || typeof instance !== "object") return null;
-    const inst = instance as { shardAssignments?: { modelId?: string } };
+    const inst = instance as Instance;
     return inst.shardAssignments?.modelId ?? null;
   }
 
@@ -936,11 +964,8 @@ class AppStore {
     else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
 
     let sharding: string | null = null;
-    const inst = instance as {
-      shardAssignments?: { runnerToShard?: Record<string, unknown> };
-    };
-    const runnerToShard = inst.shardAssignments?.runnerToShard || {};
-    const firstShardWrapped = Object.values(runnerToShard)[0];
+    const inst = instance as Instance;
+    const firstShardWrapped = getInstanceFirstShard(inst);
     if (firstShardWrapped) {
       const [shardTag] = this.getTaggedValue(firstShardWrapped);
       if (shardTag === "PipelineShardMetadata") sharding = "Pipeline";
@@ -2262,7 +2287,7 @@ class AppStore {
         if (keys.length === 1) {
           const instance = (instanceWrapper as Record<string, unknown>)[
             keys[0]
-          ] as { shardAssignments?: { modelId?: string } };
+          ] as Instance;
           const instanceModelId = instance?.shardAssignments?.modelId;
 
           // ensure to only return requestedModelId that matches an instance

dashboard/src/routes/+page.svelte

@@ -65,6 +65,11 @@
     nodeThunderboltBridge,
     nodeIdentities,
     isConnected,
+    getInstanceFirstShard,
+    getInstanceNodeIds,
+    getInstanceRunnerIds,
+    getInstanceShards,
+    type Instance,
     type DownloadProgress,
     type PlacementPreview,
   } from "$lib/stores/app.svelte";
@@ -998,11 +1003,7 @@
     if (keys.length !== 1) return new Set();
     const instance = (instanceWrapped as Record<string, unknown>)[keys[0]];
     if (!instance || typeof instance !== "object") return new Set();
-    const inst = instance as {
-      shardAssignments?: { nodeToRunner?: Record<string, string> };
-    };
-    if (!inst.shardAssignments?.nodeToRunner) return new Set();
-    return new Set(Object.keys(inst.shardAssignments.nodeToRunner));
+    return new Set(getInstanceNodeIds(instance as Instance));
   }
 
   function toggleInstanceDownloadDetails(nodeId: string): void {
@@ -1784,13 +1785,7 @@
       };
     }
 
-    const inst = instance as {
-      shardAssignments?: {
-        nodeToRunner?: Record<string, string>;
-        runnerToShard?: Record<string, unknown>;
-        modelId?: string;
-      };
-    };
+    const inst = instance as Instance;
     const instanceModelId = inst.shardAssignments?.modelId;
 
     if (!instanceModelId) {
@@ -1805,16 +1800,7 @@
       };
     }
 
-    // Get node IDs assigned to this instance
-    const nodeToRunner = inst.shardAssignments?.nodeToRunner || {};
-    const runnerToShard = inst.shardAssignments?.runnerToShard || {};
-    const runnerToNode: Record<string, string> = {};
-    for (const [nodeId, runnerId] of Object.entries(nodeToRunner)) {
-      runnerToNode[runnerId] = nodeId;
-    }
-    const instanceNodeIds = Object.keys(runnerToShard)
-      .map((runnerId) => runnerToNode[runnerId])
-      .filter(Boolean);
+    const instanceNodeIds = getInstanceNodeIds(inst);
 
     const result = collectDownloadStatus(instanceModelId, instanceNodeIds);
 
@@ -1858,6 +1844,7 @@
       case "FAILED":
         return "text-red-400";
       case "SHUTDOWN":
+      case "SHUTTING DOWN":
         return "text-gray-400";
       case "DOWNLOADING":
         return "text-blue-400";
@@ -1865,6 +1852,7 @@
       case "WARMING UP":
       case "WAITING":
       case "INITIALIZING":
+      case "CONNECTING":
         return "text-yellow-400";
       case "RUNNING":
         return "text-teal-400";
@@ -1887,10 +1875,7 @@
       return { statusText: "PREPARING", statusClass: "inactive" };
     }
 
-    const inst = instance as {
-      shardAssignments?: { runnerToShard?: Record<string, unknown> };
-    };
-    const runnerIds = Object.keys(inst.shardAssignments?.runnerToShard || {});
+    const runnerIds = getInstanceRunnerIds(instance as Instance);
 
     const statuses = runnerIds
       .map((rid) => {
@@ -1898,14 +1883,15 @@
         if (!r) return null;
         const [kind] = getTagged(r);
         const statusMap: Record<string, string> = {
-          RunnerWaitingForInitialization: "WaitingForInitialization",
-          RunnerInitializingBackend: "InitializingBackend",
-          RunnerWaitingForModel: "WaitingForModel",
+          RunnerIdle: "Idle",
+          RunnerConnecting: "Connecting",
+          RunnerConnected: "Connected",
           RunnerLoading: "Loading",
           RunnerLoaded: "Loaded",
           RunnerWarmingUp: "WarmingUp",
           RunnerReady: "Ready",
           RunnerRunning: "Running",
+          RunnerShuttingDown: "ShuttingDown",
           RunnerShutdown: "Shutdown",
           RunnerFailed: "Failed",
         };
@@ -1959,14 +1945,15 @@
       return { statusText: "RUNNING", statusClass: "running" };
     if (has("Ready")) return { statusText: "READY", statusClass: "loaded" };
     if (has("Loaded")) return { statusText: "LOADED", statusClass: "loaded" };
-    if (has("WaitingForModel"))
-      return { statusText: "WAITING", statusClass: "starting" };
-    if (has("InitializingBackend"))
-      return { statusText: "INITIALIZING", statusClass: "starting" };
-    if (has("WaitingForInitialization"))
+    if (has("Connected"))
       return { statusText: "INITIALIZING", statusClass: "starting" };
+    if (has("Connecting"))
+      return { statusText: "CONNECTING", statusClass: "starting" };
+    if (has("Idle")) return { statusText: "WAITING", statusClass: "starting" };
+    if (has("ShuttingDown"))
+      return { statusText: "SHUTTING DOWN", statusClass: "inactive" };
 
-    return { statusText: "RUNNING", statusClass: "active" };
+    return { statusText: "PREPARING", statusClass: "inactive" };
   }
 
   function getBytes(value: unknown): number {
@@ -2039,7 +2026,7 @@
   function getInstanceModelId(instanceWrapped: unknown): string {
     const [, instance] = getTagged(instanceWrapped);
     if (!instance || typeof instance !== "object") return "Unknown";
-    const inst = instance as { shardAssignments?: { modelId?: string } };
+    const inst = instance as Instance;
     return inst.shardAssignments?.modelId || "Unknown Model";
   }
 
@@ -2067,17 +2054,11 @@
     if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
     else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
 
-    const inst = instance as {
-      shardAssignments?: {
-        nodeToRunner?: Record<string, string>;
-        runnerToShard?: Record<string, unknown>;
-      };
-    };
+    const inst = instance as Instance;
 
     // Sharding strategy from first shard
     let sharding = "Unknown";
-    const runnerToShard = inst.shardAssignments?.runnerToShard || {};
-    const firstShardWrapped = Object.values(runnerToShard)[0];
+    const firstShardWrapped = getInstanceFirstShard(inst);
     if (firstShardWrapped) {
       const [shardTag] = getTagged(firstShardWrapped);
       if (shardTag === "PipelineShardMetadata") sharding = "Pipeline";
@@ -2087,8 +2068,7 @@
     }
 
     // Node names from topology
-    const nodeToRunner = inst.shardAssignments?.nodeToRunner || {};
-    const nodeIds = Object.keys(nodeToRunner);
+    const nodeIds = getInstanceNodeIds(inst);
     const nodeNames = nodeIds.map((nodeId) => {
       const node = data?.nodes?.[nodeId];
       return node?.friendly_name || nodeId.slice(0, 8);
@@ -2192,35 +2172,19 @@
   }
 
   function getOrderedRunnerNodes(
-    instance: Record<string, unknown>,
+    instance: Instance,
     shardType: "Pipeline" | "Tensor",
   ) {
-    const runnerToShard =
-      (
-        instance.shardAssignments as
-          | { runnerToShard?: Record<string, unknown> }
-          | undefined
-      )?.runnerToShard || {};
-    const nodeToRunner =
-      (
-        instance.shardAssignments as
-          | { nodeToRunner?: Record<string, string> }
-          | undefined
-      )?.nodeToRunner || {};
-    const runnerEntries = Object.entries(runnerToShard).map(
-      ([runnerId, shardWrapped]) => {
+    const runnerEntries = getInstanceShards(instance).map(
+      ([nodeId, runnerId, shardWrapped]) => {
         const [tag, shard] = getTagged(shardWrapped);
         const meta = shard as
           | {
-              modelMeta?: {
-                worldSize?: number;
-                nLayers?: number;
-                deviceRank?: number;
-              };
+              deviceRank?: number;
             }
           | undefined;
-        const deviceRank = meta?.modelMeta?.deviceRank ?? 0;
-        return { runnerId, tag, deviceRank };
+        const deviceRank = meta?.deviceRank ?? 0;
+        return { nodeId, runnerId, tag, deviceRank };
       },
     );
 
@@ -2231,13 +2195,11 @@
           : r.tag === "TensorShardMetadata",
       )
       .sort((a, b) => a.deviceRank - b.deviceRank)
-      .map((r, idx) => {
-        const nodeId = Object.entries(nodeToRunner).find(
-          ([, rid]) => rid === r.runnerId,
-        )?.[0];
-        return { nodeId, runnerId: r.runnerId, order: idx };
-      })
-      .filter((item) => item.nodeId);
+      .map((r, idx) => ({
+        nodeId: r.nodeId,
+        runnerId: r.runnerId,
+        order: idx,
+      }));
 
     return ordered as Array<{
       nodeId: string;
@@ -2281,10 +2243,7 @@
 
     // Jaccl (RDMA) – show RDMA interfaces from ibvDevices
     if (instanceTag === "MlxJacclInstance") {
-      const ordered = getOrderedRunnerNodes(
-        instance as Record<string, unknown>,
-        "Tensor",
-      );
+      const ordered = getOrderedRunnerNodes(instance as Instance, "Tensor");
       const ibvDevices =
         (instance as { ibvDevices?: Array<Array<string | null>> }).ibvDevices ||
         [];
@@ -2316,10 +2275,7 @@
 
     // Ring – derive ring order from pipeline shard ranks and pick host IPs from hostsByNode
     if (instanceTag === "MlxRingInstance") {
-      const ordered = getOrderedRunnerNodes(
-        instance as Record<string, unknown>,
-        "Pipeline",
-      );
+      const ordered = getOrderedRunnerNodes(instance as Instance, "Pipeline");
       const hostsByNode =
         (
           instance as {
@@ -2606,6 +2562,7 @@
         status.statusText === "WARMING UP" ||
         status.statusText === "WAITING" ||
         status.statusText === "INITIALIZING" ||
+        status.statusText === "CONNECTING" ||
         status.statusText === "PREPARING"
       ) {
         chatLaunchState = "launching";
@@ -5108,7 +5065,10 @@
                   {@const isFailed = statusText === "FAILED"}
                   {@const isLoading = statusText === "LOADING"}
                   {@const isWarmingUp =
-                    statusText === "WARMING UP" || statusText === "WAITING"}
+                    statusText === "WARMING UP" ||
+                    statusText === "WAITING" ||
+                    statusText === "INITIALIZING" ||
+                    statusText === "CONNECTING"}
                   {@const isReady =
                     statusText === "READY" || statusText === "LOADED"}
                   {@const isRunning = statusText === "RUNNING"}
@@ -6244,7 +6204,10 @@
                     {@const isFailed = statusText === "FAILED"}
                     {@const isLoading = statusText === "LOADING"}
                     {@const isWarmingUp =
-                      statusText === "WARMING UP" || statusText === "WAITING"}
+                      statusText === "WARMING UP" ||
+                      statusText === "WAITING" ||
+                      statusText === "INITIALIZING" ||
+                      statusText === "CONNECTING"}
                     {@const isReady =
                       statusText === "READY" || statusText === "LOADED"}
                     {@const isRunning = statusText === "RUNNING"}

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

@@ -3,7 +3,11 @@
   import { fade } from "svelte/transition";
   import HeaderNav from "$lib/components/HeaderNav.svelte";
   import IntegrationCard from "$lib/components/IntegrationCard.svelte";
-  import { instances, refreshState } from "$lib/stores/app.svelte";
+  import {
+    instances,
+    refreshState,
+    type Instance,
+  } from "$lib/stores/app.svelte";
   import { onMount } from "svelte";
 
   const apiUrl = browser
@@ -24,9 +28,7 @@
         if (values.length > 0) {
           const instance = values[0];
           if (instance && typeof instance === "object") {
-            const inst = instance as {
-              shardAssignments?: { modelId?: string };
-            };
+            const inst = instance as Instance;
             const modelId = inst.shardAssignments?.modelId;
             if (modelId && !models.includes(modelId)) {
               models.push(modelId);

docs/architecture.md

@@ -0,0 +1,84 @@
+# EXO Architecture overview
+
+EXO uses an _Event Sourcing_ architecture, and Erlang-style _message passing_. To facilitate this, we've written a channel library extending anyio channels with inspiration from tokio::sync::mpsc. 
+
+Each logical module - designed to be functional independently of the others - communicates with the rest of the system by sending messages on topics.
+
+## Systems
+
+There are currently 5 major systems:
+
+- Master
+    
+    Executes placement and orders events through a single writer
+
+- Worker
+    
+    Schedules work on a node, gathers system information, etc.#
+
+- Runner
+    
+    Executes inference jobs (for now) in an isolated process from the worker for fault-tolerance.
+
+- API
+    
+    Runs a python webserver for exposing state and commands to client applications
+
+- Election
+    
+    Implements a distributed algorithm for master election in unstable networking conditions
+
+## API Layer
+
+The API system uses multiple adapters to support multiple API formats, converting them to a single request / response type.
+
+### Adapter Pattern
+
+Adapters convert between external API formats and EXO's internal types:
+
+```
+Chat Completions → [adapter] → TextGenerationTaskParams → Application
+Claude Messages  → [adapter] → TextGenerationTaskParams → Application
+Responses API    → [adapter] → TextGenerationTaskParams → Application
+Ollama API       → [adapter] → TextGenerationTaskParams → Application
+```
+
+Each adapter implements two key functions:
+1. **Request conversion**: Converts API-specific requests to `TextGenerationTaskParams`
+2. **Response generation**: Converts internal `TokenChunk` streams back to API-specific formats (streaming and non-streaming)
+
+
+## Topics
+
+There are currently 5 topics:
+
+- Commands
+
+    The API and Worker instruct the master when the event log isn't sufficient. Namely placement and catchup requests go through Commands atm.
+
+- Local Events
+
+    All nodes write events here, the master reads those events and orders them
+
+- Global Events
+
+    The master writes events here, all nodes read from this topic and fold the produced events into their `State`
+
+- Election Messages
+
+    Before establishing a cluster, nodes communicate here to negotiate a master node.
+
+- Connection Messages
+
+    The networking system write mdns-discovered hardware connections here.
+
+
+## Event Sourcing
+
+Lots has been written about event sourcing, but it lets us centralize faulty connections and message ACKing with the following model.
+
+Whenever a device produces side effects, it captures those side effects in an `Event`. `Event`s are then "applied" to their model of `State`, which is globally distributed across the cluster. Whenever a command is received, it is combined with state to produce side effects, captured in yet more events. The rule of thumb is "`Event`s are past tense, `Command`s are imperative". Telling a node to perform some action like "place this model" or "Give me a copy of the event log" is represented by a command (The worker's `Task`s are also commands), while "this node is using 300GB of ram" is an event. Notably, `Event`s SHOULD never cause side effects on their own. There are a few exceptions to this, we're working out the specifics of generalizing the distributed event sourcing model to make it better suit our needs
+
+## Purity
+
+A significant goal of the current design is to make data flow explicit. Classes should either represent simple data (`FrozenModel`s typically, and `TaggedModel`s for unions) or active `System`s (Erlang `Actor`s), with all transformations of that data being "referentially transparent" - destructure and construct new data, don't mutate in place. We have had varying degrees of success with this, and are still exploring where purity makes sense.

flake.nix

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

justfile

@@ -16,10 +16,10 @@ check:
     uv run basedpyright --project pyproject.toml
 
 sync:
-    uv sync --all-packages --extra mlx
+    uv sync --all-packages
 
 sync-clean:
-    uv sync --all-packages --extra mlx --force-reinstall --no-cache
+    uv sync --all-packages --force-reinstall --no-cache
 
 rust-rebuild:
     PYO3_PYTHON="$(uv run python -c 'import sys; print(sys.executable)')" cargo run --bin stub_gen

resources/inference_model_cards/moonshotai--Kimi-K2.7-Code.toml

@@ -1,36 +0,0 @@
-model_id = "moonshotai/Kimi-K2.7-Code"
-n_layers = 61
-hidden_size = 7168
-num_key_value_heads = 64
-supports_tensor = true
-tasks = ["TextGeneration"]
-family = "kimi"
-quantization = ""
-base_model = "Kimi K2.7 Code"
-capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
-context_length = 262144
-backends = ["MlxMetal", "MlxCuda", "MlxCpu"]
-[storage_size]
-in_bytes = 595204986173
-
-# Vision tower + mm_projector extracted unmodified (bf16) from the official
-# repo, in the same format as exolabs/Kimi-K2.6-vision; extraction script
-# included in the weights repo. Vision config is identical to Kimi-K2.6's.
-[vision]
-image_token_id = 163605
-model_type = "kimi_vl"
-weights_repo = "aidiffuser/Kimi-K2.7-Code-vision"
-processor_repo = "moonshotai/Kimi-K2.7-Code"
-
-# Source: https://huggingface.co/moonshotai/Kimi-K2.7-Code
-# (recommends temperature 1.0 / top_p 0.95 for thinking mode, same as K2.6)
-[sampling_defaults]
-temperature = 1.0
-top_p = 0.95
-min_p = 0.01
-
-[sampling_defaults.non_thinking]
-temperature = 0.6
-top_p = 0.95
-min_p = 0.01

rust/exo_rs/Cargo.toml

@@ -35,7 +35,7 @@ pyo3-async-runtimes = { workspace = true, features = [
 ] }
 pyo3-log.workspace = true
 
-pidfile-rs = { workspace = true }
+pidfile-rs = { git = "https://github.com/AndreiCravtov/pidfile-rs" }
 
 # async runtime
 tokio = { workspace = true, features = ["full"] }
@@ -48,6 +48,7 @@ env_logger.workspace = true
 
 # Networking
 zenoh.workspace = true
+zenoh-ext = { workspace = true, features = ["unstable"] }
 rand.workspace = true
 serde_json.workspace = true
 parking_lot.workspace = true

rust/exo_rs/exo_rs.pyi

@@ -2,40 +2,39 @@
 # ruff: noqa: E501, F401, F403, F405
 
 import builtins
+import collections.abc
 import os
 import pathlib
 import typing
 __all__ = [
-    "FromSwarm",
+    "LVAggregator",
+    "LVPublisher",
+    "LVSubscriber",
     "NetworkingHandle",
     "Pidfile",
     "PidfileError",
+    "PyFromSwarm",
+    "SessionHandle",
+    "Storage",
+    "StorageGetter",
 ]
 
-class FromSwarm:
-    @typing.final
-    class Connection(FromSwarm):
-        __match_args__ = ("connected",)
-        @property
-        def connected(self) -> builtins.bool: ...
-        def __new__(cls, connected: builtins.bool) -> FromSwarm.Connection: ...
-    
-    @typing.final
-    class Message(FromSwarm):
-        __match_args__ = ("topic", "data",)
-        @property
-        def topic(self) -> builtins.str: ...
-        @property
-        def data(self) -> bytes: ...
-        def __new__(cls, topic: builtins.str, data: bytes) -> FromSwarm.Message: ...
-    
-    ...
+@typing.final
+class LVAggregator:
+    def dump(self) -> builtins.dict[builtins.str, builtins.str]: ...
+
+@typing.final
+class LVPublisher:
+    def put(self, data: builtins.str) -> collections.abc.Awaitable[None]: ...
+
+@typing.final
+class LVSubscriber:
+    def recv(self) -> collections.abc.Awaitable[tuple[str, str] | None]: ...
 
 @typing.final
 class NetworkingHandle:
     @staticmethod
-    def new(identity: builtins.str, namespace: builtins.str, listen_port: builtins.int, discovery_service_port: builtins.int) -> NetworkingHandle: ...
-    def recv(self) -> typing.Awaitable[FromSwarm]: ...
+    def new(identity: builtins.str, listen_port: builtins.int, discovery_service_port: builtins.int) -> NetworkingHandle: ...
     async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
         r"""
         Subscribe to a `GossipSub` topic.
@@ -54,6 +53,7 @@ class NetworkingHandle:
         
         If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
         """
+    async def recv(self) -> PyFromSwarm: ...
 
 @typing.final
 class Pidfile:
@@ -80,6 +80,7 @@ class Pidfile:
     def __new__(cls, path: builtins.str | os.PathLike | pathlib.Path, mode: builtins.int) -> Pidfile:
         r"""
         Creates a new PID file and locks it.
+        Writes the current process ID to the PID file.
         
         If the PID file cannot be locked, returns `PidfileError::AlreadyRunning` with
         a PID of the already running process, or `None` if no PID has been written to
@@ -111,3 +112,43 @@ class PidfileError(builtins.Exception):
     def __repr__(self) -> builtins.str: ...
     def __str__(self) -> builtins.str: ...
 
+class PyFromSwarm:
+    @typing.final
+    class Connection(PyFromSwarm):
+        __match_args__ = ("connected",)
+        @property
+        def connected(self) -> builtins.bool: ...
+        def __new__(cls, connected: builtins.bool) -> PyFromSwarm.Connection: ...
+    
+    @typing.final
+    class Message(PyFromSwarm):
+        __match_args__ = ("topic", "data",)
+        @property
+        def topic(self) -> builtins.str: ...
+        @property
+        def data(self) -> bytes: ...
+        def __new__(cls, topic: builtins.str, data: bytes) -> PyFromSwarm.Message: ...
+    
+    ...
+
+@typing.final
+class SessionHandle:
+    @staticmethod
+    def new(identity: builtins.str, listen_port: builtins.int, discovery_service_port: builtins.int) -> tuple[SessionHandle, NetworkingHandle]: ...
+    def last_value_aggregator(self, prefix: builtins.str) -> LVAggregator: ...
+    def last_value_subscriber(self, kexpr: builtins.str) -> LVSubscriber: ...
+    def last_value_publisher(self, kexpr: builtins.str) -> LVPublisher: ...
+    def storage_interface(self) -> Storage: ...
+
+@typing.final
+class Storage:
+    def get(self, key: builtins.str) -> collections.abc.Awaitable[str | None]: ...
+    def get_many(self, key: builtins.str) -> StorageGetter: ...
+    def put(self, key: builtins.str, data: builtins.str) -> collections.abc.Awaitable[None]: ...
+    def delete(self, key: builtins.str) -> collections.abc.Awaitable[None]: ...
+    def dump(self, prefix: builtins.str) -> collections.abc.Awaitable[dict[str, str]]: ...
+
+@typing.final
+class StorageGetter:
+    def recv(self) -> collections.abc.Awaitable[tuple[str, str] | None]: ...
+

rust/exo_rs/pyproject.toml

@@ -28,6 +28,3 @@ generate-init-py = true
 log_cli = true
 log_cli_level = "INFO"
 asyncio_mode = "auto"
-
-[tool.uv]
-cache-keys = [{ file = "src/**/*.rs" }]

rust/exo_rs/src/last_value.rs

@@ -0,0 +1,181 @@
+use networking::{Session, liveliness_aggregator::LivelinessAggregator};
+use std::{
+    collections::{HashMap, HashSet},
+    sync::Arc,
+};
+use zenoh::{Result as ZResult, Wait};
+
+use parking_lot::Mutex;
+use pyo3::{
+    exceptions::{PyConnectionError, PyRuntimeError},
+    prelude::*,
+};
+use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
+use zenoh::{
+    handlers::FifoChannelHandler,
+    sample::{Sample, SampleKind},
+};
+use zenoh_ext::{
+    AdvancedPublisher, AdvancedSubscriber, AdvancedSubscriberBuilderExt, HistoryConfig,
+};
+
+#[gen_stub_pyclass]
+#[pyclass]
+pub struct LVAggregator {
+    pub prefix: Arc<str>,
+    pub store: Arc<Mutex<HashMap<String, String>>>,
+    pub current_live: LivelinessAggregator,
+}
+
+pub fn spawn_lv_aggregator_onto(session: &Session, prefix: Arc<str>) -> ZResult<LVAggregator> {
+    // nota bene: config must be kept in line with SessionHandle::last_value_receiver
+    let store = Arc::new(Mutex::new(HashMap::default()));
+    session
+        .z
+        //assuming all LV aggregators are prefix/node_id/atomic_json
+        .declare_subscriber(format!("{prefix}/*/*"))
+        .advanced()
+        .history(
+            HistoryConfig::default()
+                .max_samples(1)
+                .detect_late_publishers(),
+        )
+        .callback({
+            let store = Arc::clone(&store);
+            let prefix = Arc::clone(&prefix);
+            move |sample| {
+                if let Some(s) = sample
+                    .key_expr()
+                    .to_string()
+                    .strip_prefix(&*prefix)
+                    .and_then(|it| it.strip_prefix('/'))
+                {
+                    let s = s.to_string();
+                    match sample.kind() {
+                        SampleKind::Put => {
+                            store.lock().insert(
+                                s,
+                                sample
+                                    .payload()
+                                    .try_to_string()
+                                    .expect("we only use utf8 encoded strings. someone messed up")
+                                    .to_string(),
+                            );
+                        }
+                        SampleKind::Delete => {
+                            store.lock().remove(&s);
+                        }
+                    }
+                };
+            }
+        })
+        .background()
+        .wait()?;
+    Ok(LVAggregator {
+        prefix,
+        store,
+        current_live: session.liveliness_aggregator.clone(),
+    })
+}
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl LVAggregator {
+    pub fn dump(&self) -> HashMap<String, String> {
+        let mut store = self.store.lock();
+        let currently_alive: HashSet<String> = self.current_live.dump();
+        // remove any keys that are no longer live
+        store.retain(|key, _| {
+            currently_alive.iter().any(|node_id| {
+                key.strip_prefix(node_id)
+                    .is_some_and(|rest| rest.starts_with("/"))
+            })
+        });
+        store.clone()
+    }
+}
+
+#[gen_stub_pyclass]
+#[pyclass]
+pub struct LVSubscriber {
+    pub subscriber: AdvancedSubscriber<FifoChannelHandler<Sample>>,
+}
+#[gen_stub_pymethods]
+#[pymethods]
+impl LVSubscriber {
+    #[gen_stub(override_return_type(
+        type_repr="collections.abc.Awaitable[tuple[str, str] | None]",
+        imports=("collections.abc")
+    ))]
+    pub fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        pyo3_async_runtimes::tokio::future_into_py(py, {
+            if self.subscriber.receiver_count() != 1 {
+                return Err(PyRuntimeError::new_err(
+                    "tried to receive twice on the same receiver",
+                ));
+            }
+            let subscriber = self.subscriber.clone();
+            async move {
+                loop {
+                    match subscriber.recv_async().await {
+                        Ok(sample) if sample.kind() == SampleKind::Delete => continue,
+                        Err(_) => {
+                            return Ok(None);
+                        }
+                        Ok(sample) => {
+                            return Ok(Some((
+                                sample.key_expr().to_string(),
+                                sample
+                                    .payload()
+                                    .try_to_string()
+                                    .expect("we only use utf8 encoded strings. someone messed up")
+                                    .to_string(),
+                            )));
+                        }
+                    }
+                }
+            }
+        })
+    }
+}
+
+#[gen_stub_pyclass]
+#[pyclass]
+pub struct LVPublisher {
+    pub state: Arc<AdvancedPublisher<'static>>,
+}
+impl LVPublisher {
+    pub fn new(publisher: AdvancedPublisher<'static>) -> Self {
+        Self {
+            state: Arc::new(publisher),
+        }
+    }
+}
+#[gen_stub_pymethods]
+#[pymethods]
+impl LVPublisher {
+    #[gen_stub(override_return_type(
+        type_repr="collections.abc.Awaitable[None]",
+        imports=("collections.abc")
+    ))]
+    pub fn put<'py>(&'py self, py: Python<'py>, data: String) -> PyResult<Bound<'py, PyAny>> {
+        let state = Arc::clone(&self.state);
+        pyo3_async_runtimes::tokio::future_into_py(py, {
+            // clone the data so py can have it back
+            async move {
+                state
+                    .put(data)
+                    .await
+                    .map_err(|e| PyConnectionError::new_err(e.to_string()))?;
+                Ok(())
+            }
+        })
+    }
+}
+
+pub fn lv_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<LVPublisher>()?;
+    m.add_class::<LVSubscriber>()?;
+    m.add_class::<LVAggregator>()?;
+    Ok(())
+}

rust/exo_rs/src/lib.rs

@@ -5,12 +5,18 @@
 //!
 
 mod allow_threading;
-// mod ident;
-mod networking;
 mod pidfile;
+// mod ident;
+pub mod last_value;
+mod networking;
+pub mod session;
+mod storage;
 
+use crate::last_value::lv_submodule;
 use crate::networking::networking_submodule;
 use crate::pidfile::pidfile_submodule;
+use crate::session::session_submodule;
+use crate::storage::storage_submodule;
 use pyo3::prelude::PyModule;
 use pyo3::{Bound, PyResult, pymodule};
 use pyo3_stub_gen::define_stub_info_gatherer;
@@ -146,7 +152,7 @@ pub(crate) mod ext {
 /// 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.
-#[pymodule(name = "exo_rs", gil_used = true)]
+#[pymodule(name = "exo_rs")]
 fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     // install logger
     pyo3_log::init();
@@ -154,15 +160,12 @@ fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     builder.enable_all();
     pyo3_async_runtimes::tokio::init(builder);
 
-    // TODO: for now this is all NOT a submodule, but figure out how to make the submodule system
-    //       work with maturin, where the types generate correctly, in the right folder, without
-    //       too many importing issues...
+    // TODO: for now this is all NOT a submodule. KISS
     pidfile_submodule(m)?;
-    // m.add_class::<PyKeypair>()?;
     networking_submodule(m)?;
-
-    // top-level constructs
-    // TODO: ...
+    lv_submodule(m)?;
+    session_submodule(m)?;
+    storage_submodule(m)?;
 
     Ok(())
 }

rust/exo_rs/src/networking.rs

@@ -4,13 +4,15 @@ use std::sync::Arc;
 use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
 use crate::ext::{ResultExt as _, TokioMpscSenderExt as _};
 use futures_lite::{Stream, StreamExt as _};
+use networking::Session;
 use networking::swarm::{FromSwarm, Swarm, ToSwarm, create_swarm};
-use networking::{Session, is_valid_zid};
 use pyo3::exceptions::{PyRuntimeError, PyValueError};
 use pyo3::prelude::*;
 use pyo3::types::PyBytes;
 use pyo3::{Bound, Py, PyAny, PyErr, PyResult, Python, pymethods};
-use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods};
+use pyo3_stub_gen::derive::{
+    gen_methods_from_python, gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods,
+};
 use tokio::sync::{Mutex, mpsc, oneshot};
 
 #[gen_stub_pyclass]
@@ -22,7 +24,7 @@ pub struct PyNetworkingHandle {
 }
 
 #[gen_stub_pyclass_complex_enum]
-#[pyclass(name = "FromSwarm")]
+#[pyclass]
 pub enum PyFromSwarm {
     Connection { connected: bool },
     Message { topic: String, data: Py<PyBytes> },
@@ -64,21 +66,24 @@ impl PyNetworkingHandle {
     // ---- Lifecycle management methods ----
 
     #[staticmethod]
-    pub fn new(
+    pub fn new<'py>(
         identity: &str,
-        namespace: &str,
         listen_port: u16,
         discovery_service_port: u16,
     ) -> PyResult<PyNetworkingHandle> {
         // todo: zenoh self assigned peers
         if listen_port == 0 {
-            todo!("cannot listen on port 0 yet");
+            todo!();
         }
         // create communication channels
         let (to_swarm, from_client) = mpsc::channel(1024);
 
         // get identity
-        if !is_valid_zid(identity) {
+        if !identity
+            .chars()
+            .all(|c| ('0'..='9').contains(&c) || ('a'..='f').contains(&c))
+            || identity.len() > 32
+        {
             return Err(PyValueError::new_err(format!(
                 "{identity} is not a valid zenoh identity"
             )));
@@ -88,7 +93,6 @@ impl PyNetworkingHandle {
         let swarm = pyo3_async_runtimes::tokio::get_runtime()
             .block_on(create_swarm(
                 identity,
-                namespace,
                 from_client,
                 listen_port,
                 discovery_service_port,
@@ -101,9 +105,7 @@ impl PyNetworkingHandle {
         })
     }
 
-    #[gen_stub(override_return_type(
-        type_repr="typing.Awaitable[FromSwarm]", imports=("typing")
-    ))]
+    #[gen_stub(skip)]
     pub fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         let swarm = Arc::clone(&self.swarm);
         pyo3_async_runtimes::tokio::future_into_py(py, async move {
@@ -189,6 +191,15 @@ impl PyNetworkingHandle {
     }
 }
 
+pyo3_stub_gen::inventory::submit! {
+    gen_methods_from_python! {
+        r#"
+            class PyNetworkingHandle:
+                async def recv() -> PyFromSwarm: ...
+        "#
+    }
+}
+
 pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<PyNetworkingHandle>()?;
     m.add_class::<PyFromSwarm>()?;

rust/exo_rs/src/pidfile.rs

@@ -77,6 +77,7 @@ impl PyPidfile {
 #[pymethods]
 impl PyPidfile {
     /// Creates a new PID file and locks it.
+    /// Writes the current process ID to the PID file.
     ///
     /// If the PID file cannot be locked, returns `PidfileError::AlreadyRunning` with
     /// a PID of the already running process, or `None` if no PID has been written to

rust/exo_rs/src/session.rs

@@ -0,0 +1,111 @@
+use networking::Session;
+use pyo3::{
+    exceptions::{PyConnectionError, PyRuntimeError, PyValueError},
+    prelude::*,
+};
+use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
+use zenoh::Wait;
+use zenoh_ext::{
+    AdvancedPublisherBuilderExt, AdvancedSubscriberBuilderExt, CacheConfig, HistoryConfig,
+    MissDetectionConfig,
+};
+
+use crate::{
+    last_value::{LVAggregator, LVPublisher, LVSubscriber, spawn_lv_aggregator_onto},
+    networking::PyNetworkingHandle,
+    storage::Storage,
+};
+
+#[gen_stub_pyclass]
+#[pyclass]
+pub struct SessionHandle {
+    pub session: Session,
+}
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl SessionHandle {
+    // NOTE: `async fn`s here that use `.await` will wrap the future in `.allow_threads_py()`
+    //       immediately beforehand to release the interpreter.
+    //       SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
+
+    // ---- Lifecycle management methods ----
+
+    #[staticmethod]
+    pub fn new<'py>(
+        identity: &str,
+        listen_port: u16,
+        discovery_service_port: u16,
+    ) -> PyResult<(SessionHandle, PyNetworkingHandle)> {
+        // get identity
+        if !identity
+            .chars()
+            .all(|c| ('0'..='9').contains(&c) || ('a'..='f').contains(&c))
+            || identity.len() > 32
+        {
+            return Err(PyValueError::new_err(format!(
+                "{identity} is not a valid zenoh identity"
+            )));
+        }
+
+        let cfg = networking::cfg(identity, listen_port).map_err(|e| {
+            PyValueError::new_err(format!("failed to write config: {}", e.to_string()))
+        })?;
+        let session = pyo3_async_runtimes::tokio::get_runtime()
+            .block_on(networking::open(cfg, listen_port, discovery_service_port))
+            .map_err(|e| {
+                PyRuntimeError::new_err(format!(
+                    "failed to spawn networking on tokio runtime: {}",
+                    e.to_string()
+                ))
+            })?;
+        let legacy = PyNetworkingHandle::from_session(session.clone());
+        Ok((Self { session }, legacy))
+    }
+
+    pub fn last_value_aggregator(&self, prefix: String) -> PyResult<LVAggregator> {
+        spawn_lv_aggregator_onto(&self.session, prefix.into()).map_err(|e| {
+            PyConnectionError::new_err(format!("failed to spawn liveliness aggregator: {e}"))
+        })
+    }
+
+    pub fn last_value_subscriber(&self, kexpr: &str) -> PyResult<LVSubscriber> {
+        // nota bene: config must be kept in track with the LVAggregator
+        self.session
+            .z
+            .declare_subscriber(kexpr)
+            .advanced()
+            .history(
+                HistoryConfig::default()
+                    .max_samples(1)
+                    .detect_late_publishers(),
+            )
+            .wait()
+            .map_err(|e| PyConnectionError::new_err(format!("failed to declare subscriber: {e}")))
+            .map(|subscriber| LVSubscriber { subscriber })
+    }
+
+    pub fn last_value_publisher(&self, kexpr: String) -> PyResult<LVPublisher> {
+        self.session
+            .z
+            .declare_publisher(kexpr)
+            .advanced()
+            .publisher_detection()
+            .sample_miss_detection(MissDetectionConfig::default())
+            .cache(CacheConfig::default().max_samples(1))
+            .wait()
+            .map_err(|e| PyConnectionError::new_err(format!("failed to declare publisher: {e}")))
+            .map(LVPublisher::new)
+    }
+
+    pub fn storage_interface(&self) -> Storage {
+        Storage {
+            session: self.session.z.clone(),
+        }
+    }
+}
+
+pub fn session_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<SessionHandle>()?;
+    Ok(())
+}

rust/exo_rs/src/storage.rs

@@ -0,0 +1,199 @@
+use std::{collections::HashMap, time::Duration};
+
+use networking::STORAGE_PREFIX;
+use pyo3::{
+    exceptions::{PyConnectionError, PyRuntimeError, PyValueError},
+    prelude::*,
+};
+use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
+use zenoh::{
+    Session as ZSession, Wait, handlers::FifoChannelHandler, query::Reply, sample::SampleKind,
+};
+
+#[gen_stub_pyclass]
+#[pyclass]
+pub struct Storage {
+    pub session: ZSession,
+}
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl Storage {
+    #[gen_stub(override_return_type(
+        type_repr="collections.abc.Awaitable[str | None]",
+        imports=("collections.abc")
+    ))]
+    pub fn get<'py>(&'py self, py: Python<'py>, key: String) -> PyResult<Bound<'py, PyAny>> {
+        if key.contains('*') {
+            return Err(PyValueError::new_err(format!(
+                "{key} is invalid -- Storage.get only supports fixed keys"
+            )));
+        }
+        let session = self.session.clone();
+        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+            let subscriber = session
+                .get(format!("{STORAGE_PREFIX}/{key}"))
+                //.allowed_destination(Locality::SessionLocal)
+                .await
+                .map_err(|e| PyConnectionError::new_err(format!("failed to query storage: {e}")))?;
+            tokio::select! {
+                _ = tokio::time::sleep(Duration::from_secs(1)) => {
+                    Ok(None)
+                }
+                reply = subscriber.recv_async() => {
+                    Ok(reply.ok()
+                        .and_then(|reply| reply.into_result().ok())
+                        .and_then(|sample| {
+                            if sample.kind() == SampleKind::Put {
+                                Some(sample
+                                    .payload()
+                                    .try_to_string()
+                                    .expect("we only use utf8 encoded strings. someone messed up").to_string()
+                                )
+                            } else { None }
+                        })
+                    )
+                }
+            }
+        })
+    }
+    pub fn get_many(&self, key: String) -> PyResult<StorageGetter> {
+        self.session
+            .get(key)
+            .wait()
+            .map_err(|e| PyConnectionError::new_err(format!("failed to query storage: {e}")))
+            .map(StorageGetter)
+    }
+
+    #[gen_stub(override_return_type(
+        type_repr="collections.abc.Awaitable[None]",
+        imports=("collections.abc")
+    ))]
+    pub fn put<'py>(
+        &'py self,
+        py: Python<'py>,
+        key: String,
+        data: String,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        if key.contains('*') {
+            return Err(PyValueError::new_err(format!(
+                "{key} is invalid -- Storage.put only supports fixed keys"
+            )));
+        }
+        let session = self.session.clone();
+        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+            session
+                .put(format!("{STORAGE_PREFIX}/{key}"), data)
+                .await
+                .map_err(|e| PyConnectionError::new_err(format!("failed to query storage: {e}")))
+        })
+    }
+
+    #[gen_stub(override_return_type(
+        type_repr="collections.abc.Awaitable[None]",
+        imports=("collections.abc")
+    ))]
+    pub fn delete<'py>(&'py self, py: Python<'py>, key: String) -> PyResult<Bound<'py, PyAny>> {
+        if key.contains('*') {
+            return Err(PyValueError::new_err(format!(
+                "{key} is invalid -- Storage.delete only supports fixed keys"
+            )));
+        }
+        let session = self.session.clone();
+        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+            session
+                .delete(format!("{STORAGE_PREFIX}/{key}"))
+                .await
+                .map_err(|e| PyConnectionError::new_err(format!("failed to query storage: {e}")))
+        })
+    }
+
+    #[gen_stub(override_return_type(
+        type_repr="collections.abc.Awaitable[dict[str, str]]",
+        imports=("collections.abc")
+    ))]
+    pub fn dump<'py>(&'py self, py: Python<'py>, prefix: String) -> PyResult<Bound<'py, PyAny>> {
+        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+            Ok(networking::read_raw_memory_storage()
+                .await
+                .into_iter()
+                .filter_map(|(key, value)| {
+                    Some((
+                        key?.as_str().strip_prefix(prefix.as_str())?.to_string(),
+                        value
+                            .payload
+                            .try_to_string()
+                            .expect("we only use utf8 encoded strings. someone messed up")
+                            .to_string(),
+                    ))
+                })
+                .collect::<HashMap<String, String>>())
+        })
+    }
+}
+
+#[gen_stub_pyclass]
+#[pyclass]
+pub struct StorageGetter(FifoChannelHandler<Reply>);
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl StorageGetter {
+    #[gen_stub(override_return_type(
+        type_repr="collections.abc.Awaitable[tuple[str, str] | None]",
+        imports=("collections.abc")
+    ))]
+    fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        if self.0.receiver_count() != 1 {
+            return Err(PyRuntimeError::new_err(
+                "Tried to call StorageGetter.recv twice concurrently",
+            ));
+        }
+        let dupe = self.0.clone();
+        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+            let sample = loop {
+                match dupe.recv_async().await {
+                    Err(_) => return Ok(None),
+                    Ok(reply) => match reply.into_result() {
+                        Err(e) => {
+                            log::warn!("Ignoring reply error: {e}");
+                            continue;
+                        }
+                        Ok(sample) => match sample.kind() {
+                            SampleKind::Put => break sample,
+                            SampleKind::Delete => {
+                                log::warn!(
+                                    "Received unexpected DELETE from queryable: {}",
+                                    sample.key_expr()
+                                );
+                                continue;
+                            }
+                        },
+                    },
+                };
+            };
+
+            let key = sample
+                .key_expr()
+                .to_string()
+                .strip_prefix(format!("{STORAGE_PREFIX}/").as_str())
+                .expect("invalid storage format encountered")
+                .to_string();
+
+            Ok(Some((
+                key,
+                sample
+                    .payload()
+                    .try_to_string()
+                    .expect("we only use utf8 encoded strings. someone messed up")
+                    .to_string(),
+            )))
+        })
+    }
+}
+
+pub fn storage_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<Storage>()?;
+    m.add_class::<StorageGetter>()?;
+    Ok(())
+}

rust/exo_rs/tests/adv_subs.rs

@@ -0,0 +1,168 @@
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
+
+use zenoh_ext::{
+    AdvancedPublisherBuilderExt, AdvancedSubscriber, AdvancedSubscriberBuilderExt, CacheConfig,
+    HistoryConfig, MissDetectionConfig,
+};
+
+use zenoh::{handlers::FifoChannelHandler, sample::Sample};
+
+// Adjust these imports to your crate/module paths.
+use exo_rs::{
+    last_value::{LVPublisher, LVSubscriber},
+    session::SessionHandle,
+};
+async fn expect_two_values(
+    sub: &AdvancedSubscriber<FifoChannelHandler<Sample>>,
+    key_a: &str,
+    val_a: &str,
+    key_b: &str,
+    val_b: &str,
+) {
+    use std::collections::HashMap;
+    use tokio::time::{Duration, Instant, timeout};
+    use zenoh::sample::SampleKind;
+
+    let deadline = Instant::now() + Duration::from_secs(10);
+    let mut seen: HashMap<String, String> = HashMap::new();
+
+    loop {
+        if seen.get(key_a).map(String::as_str) == Some(val_a)
+            && seen.get(key_b).map(String::as_str) == Some(val_b)
+        {
+            return;
+        }
+
+        let remaining = deadline.saturating_duration_since(Instant::now());
+        assert!(
+            !remaining.is_zero(),
+            "timed out waiting for both historical samples; expected {key_a}={val_a}, {key_b}={val_b}; seen = {seen:?}"
+        );
+
+        match timeout(remaining.min(Duration::from_millis(750)), sub.recv_async()).await {
+            Ok(Ok(sample)) => {
+                if sample.kind() == SampleKind::Delete {
+                    continue;
+                }
+
+                let key = sample.key_expr().to_string();
+                let value = sample
+                    .payload()
+                    .try_to_string()
+                    .expect("payload should be UTF-8")
+                    .to_string();
+
+                if key == key_a || key == key_b {
+                    eprintln!("received relevant {key} = {value}");
+                    seen.insert(key, value);
+                } else {
+                    eprintln!("received unrelated {key} = {value}");
+                }
+            }
+            Ok(Err(e)) => panic!("subscriber receive failed: {e}"),
+            Err(_) => {}
+        }
+    }
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn lv_subscriber_receives_last_value_from_multiple_publishers() {
+    let cfg =
+        networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414).expect("create config");
+    let n_session = networking::open(cfg, 52414, 52413)
+        .await
+        .expect("open session");
+
+    let session = SessionHandle { session: n_session };
+
+    let run_id = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap()
+        .as_nanos();
+
+    let base = format!("zenoh_advanced_history_test/{run_id}");
+
+    let key_a = format!("{base}/a");
+    let key_b = format!("{base}/b");
+    let sub_key = format!("{base}/*");
+
+    let pub1: LVPublisher = session
+        .last_value_publisher(key_a.clone())
+        .expect("declare LV publisher a");
+
+    pub1.state.put("aa").await.expect("publish aa");
+
+    let pub2: LVPublisher = session
+        .last_value_publisher(key_b.clone())
+        .expect("declare LV publisher b");
+
+    pub2.state.put("bb").await.expect("publish bb");
+
+    // Let publisher detection / cache metadata settle before the late subscriber joins.
+    tokio::time::sleep(Duration::from_millis(250)).await;
+
+    let sub: LVSubscriber = session
+        .last_value_subscriber(&*sub_key)
+        .expect("declare LV subscriber");
+
+    expect_two_values(&sub.subscriber, &*key_a, "aa", &*key_b, "bb").await
+}
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn wildcard_advanced_subscriber_receives_history_from_both_publishers() {
+    let cfg =
+        networking::cfg(&format!("{:x}", rand::random::<u128>()), 52412).expect("create config");
+    let n_session = networking::open(cfg, 52412, 52411)
+        .await
+        .expect("open session");
+    let session = n_session.z.clone();
+
+    // Unique prefix so the wildcard subscriber cannot accidentally see unrelated traffic.
+    let run_id = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap()
+        .as_nanos();
+
+    let base = format!("zenoh_advanced_history_test/{run_id}");
+    let key_a = format!("{base}/a");
+    let key_b = format!("{base}/b");
+    let sub_key = format!("{base}/*");
+
+    let pub1 = session
+        .declare_publisher(key_a.clone())
+        .advanced()
+        .publisher_detection()
+        .sample_miss_detection(MissDetectionConfig::default())
+        .cache(CacheConfig::default().max_samples(1))
+        .await
+        .expect("declare advanced publisher a");
+
+    pub1.put("aa").await.expect("publish aa");
+
+    let pub2 = session
+        .declare_publisher(key_b.clone())
+        .advanced()
+        .sample_miss_detection(MissDetectionConfig::default())
+        .publisher_detection()
+        .cache(CacheConfig::default().max_samples(1))
+        .await
+        .expect("declare advanced publisher b");
+
+    pub2.put("bb").await.expect("publish bb");
+
+    // Give liveliness/cache declarations a brief chance to settle before declaring
+    // the late-joining advanced subscriber.
+    tokio::time::sleep(Duration::from_millis(250)).await;
+
+    let sub = session
+        .declare_subscriber(sub_key)
+        .advanced()
+        .history(
+            HistoryConfig::default()
+                .max_samples(1)
+                .detect_late_publishers(),
+        )
+        .await
+        .expect("declare advanced subscriber");
+
+    expect_two_values(&sub, &*key_a, "aa", &*key_b, "bb").await
+}

rust/exo_rs/tests/dummy.rs

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

rust/exo_rs/tests/test_python.py

@@ -6,14 +6,14 @@ from _pytest.capture import CaptureFixture
 from exo_rs import (
     NetworkingHandle,
     Pidfile,
-    FromSwarm,
+    PyFromSwarm,
 )
 
 
 @pytest.mark.asyncio
 async def test_sleep_on_multiple_items() -> None:
     print("PYTHON: starting handle")
-    h = NetworkingHandle.new(os.urandom(16).hex().lstrip("0"), 52414, 52413)
+    h = NetworkingHandle.new(os.urandom(16).hex().rstrip("0"), 52412, 52411)
     print("PYTHON: handle started")
 
     rt = asyncio.create_task(_await_recv(h))
@@ -36,9 +36,9 @@ async def _await_recv(h: NetworkingHandle):
     while True:
         event = await h.recv()
         match event:
-            case FromSwarm.Connection() as c:
+            case PyFromSwarm.Connection() as c:
                 print(f"PYTHON: connection update: {c}")
-            case FromSwarm.Message() as m:
+            case PyFromSwarm.Message() as m:
                 print(f"PYTHON: message: {m}")
 
 

rust/exo_rs/tests/test_storage_handle.py

@@ -0,0 +1,69 @@
+import os
+import uuid
+
+import pytest
+from exo_rs import SessionHandle
+
+
+ZENOH_PORT = 52414
+DISCOVERY_PORT = 52413
+
+
+@pytest.fixture(scope="module")
+def storage():
+    node_id = os.urandom(16).hex().rstrip("0")
+
+    session_handle, _nh = SessionHandle.new(
+        node_id,
+        ZENOH_PORT,
+        DISCOVERY_PORT,
+    )
+
+    return session_handle.storage_interface()
+
+
+@pytest.mark.asyncio
+async def test_storage_get_missing_key_returns_none(storage):
+    key = f"tests/storage/{uuid.uuid4().hex}/missing"
+
+    value = await storage.get(key)
+    assert value is None
+
+
+@pytest.mark.asyncio
+async def test_storage_put_then_get_returns_value(storage):
+    key = f"tests/storage/{uuid.uuid4().hex}/value"
+    expected = "hello storage"
+
+    await storage.put(key, expected)
+    assert await storage.get(key) == expected
+
+
+@pytest.mark.asyncio
+async def test_storage_put_overwrites_value(storage):
+    key = f"tests/storage/{uuid.uuid4().hex}/overwrite"
+
+    await storage.put(key, "old")
+    await storage.put(key, "new")
+    assert await storage.get(key) == "new"
+
+
+@pytest.mark.asyncio
+async def test_storage_put_overwrites_value(storage):
+    key = f"tests/storage/{uuid.uuid4().hex}/overwrite"
+
+    await storage.put(key, "old")
+    await storage.delete(key)
+    assert await storage.get(key) == None
+
+
+@pytest.mark.asyncio
+async def test_storage_get_rejects_wildcard_key(storage):
+    with pytest.raises(ValueError, match="only supports fixed keys"):
+        await storage.get("tests/storage/*")
+
+
+@pytest.mark.asyncio
+async def test_storage_put_rejects_wildcard_key(storage):
+    with pytest.raises(ValueError, match="only supports fixed keys"):
+        await storage.put("tests/storage/*", "value")

rust/networking/Cargo.toml

@@ -15,13 +15,13 @@ zenoh-plugin-trait.workspace = true
 rand.workspace = true
 log.workspace = true
 bytemuck = { workspace = true, features = ["derive"] }
-socket2.workspace = true
-blake3.workspace = true
+socket2 = "0.6.4"
+blake3 = "1.8.5"
 
 [lints]
 workspace = true
 
 [dev-dependencies]
 env_logger.workspace = true
-smol.workspace = true
-tracing.workspace = true
+smol = "2.0.2"
+tracing = "0.1.44"

rust/networking/examples/do_nothing.rs

@@ -7,7 +7,7 @@ async fn main() -> Result<()> {
     zenoh::init_log_from_env_or("info");
     info!("Opening session...");
     let cfg = networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414)?;
-    let session = networking::open(cfg, "exo", 52414, 52413).await?;
+    let session = networking::open(cfg, 52414, 52413).await?;
     let _tok = session
         .z
         .liveliness()

rust/networking/examples/heartbeat.rs

@@ -8,7 +8,7 @@ async fn main() -> Result<()> {
     env_logger::try_init_from_env(Env::new().default_filter_or("info")).expect("logger failed");
     info!("Opening session...");
     let cfg = networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414)?;
-    let session = networking::open(cfg, "exo", 52414, 52413).await?;
+    let session = networking::open(cfg, 52414, 52413).await?;
     let _tok = session
         .z
         .liveliness()

rust/networking/examples/put_string.rs

@@ -14,7 +14,7 @@ async fn main() -> Result<()> {
         .expect("USAGE: put_string <n> -- pub a string of n bytes into stream/data");
     info!("Opening session...");
     let cfg = networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414)?;
-    let session = networking::open(cfg, "exo", 52414, 52413).await?;
+    let session = networking::open(cfg, 52414, 52413).await?;
     let _tok = session
         .z
         .liveliness()

rust/networking/examples/serve_storage.rs

@@ -10,7 +10,7 @@ async fn main() -> Result<()> {
     env_logger::try_init_from_env(Env::new().default_filter_or("info")).expect("logger failed");
     info!("Opening session...");
     let cfg = networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414)?;
-    let session = networking::open(cfg, "exo", 52414, 52413).await?;
+    let session = networking::open(cfg, 52414, 52413).await?;
     let _tok = session
         .z
         .liveliness()

rust/networking/examples/z_get.rs

@@ -10,7 +10,7 @@ async fn main() -> Result<()> {
     env_logger::try_init_from_env(Env::new().default_filter_or("info")).expect("logger failed");
     info!("Opening session...");
     let cfg = networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414)?;
-    let session = networking::open(cfg, "exo", 52414, 52413).await?;
+    let session = networking::open(cfg, 52414, 52413).await?;
     let other_live = session
         .z
         .liveliness()

rust/networking/src/discovery.rs

@@ -1,5 +1,5 @@
 use std::{
-    io,
+    env, io,
     net::{Ipv6Addr, SocketAddr, SocketAddrV6},
     sync::Arc,
     time::Duration,
@@ -37,12 +37,17 @@ pub struct Discovered {
 }
 
 impl Discovery {
-    pub async fn new(
-        zid: ZenohId,
-        namespace: [u8; 8],
-        listen_port: u16,
-        discovery_port: u16,
-    ) -> io::Result<Self> {
+    pub async fn new(zid: ZenohId, listen_port: u16, discovery_port: u16) -> io::Result<Self> {
+        let namespace: [u8; 8] = {
+            blake3::hash(
+                env::var("EXO_ZENOH_NAMESPACE")
+                    .unwrap_or_else(|_| "exo".to_string())
+                    .as_bytes(),
+            )
+            .as_bytes()[..8]
+                .try_into()
+                .expect("8=8")
+        };
         let sock = socket2::Socket::new(
             socket2::Domain::IPV6,
             socket2::Type::DGRAM,

rust/networking/src/lib.rs

@@ -7,21 +7,24 @@ use zenoh_plugin_trait::PluginsManager;
 
 pub use zenoh::{Config, config::ZenohId};
 
-use crate::discovery::Discovery;
+use crate::{
+    discovery::Discovery,
+    liveliness_aggregator::{LivelinessAggregator, spawn_liveliness_aggregator},
+};
+
+pub use zenoh_plugin_storage_manager::read_raw_memory_storage;
+pub const STORAGE_PREFIX: &str = "storage/mem1";
 
 pub mod discovery;
+pub mod liveliness_aggregator;
 pub mod swarm;
 
-pub fn is_valid_zid(identity: &str) -> bool {
-    let mut iter = identity.chars();
-    iter.next()
-        .is_some_and(|c| ('1'..='9').contains(&c) || ('a'..='f').contains(&c))
-        && iter.all(|c| ('0'..='9').contains(&c) || ('a'..='f').contains(&c))
-        && identity.len() <= 32
-}
-
 pub fn cfg(identity: &str, listen_port: u16) -> Result<zenoh::Config> {
-    assert!(is_valid_zid(identity));
+    assert!(
+        identity
+            .chars()
+            .all(|c| ('0'..='9').contains(&c) || ('a'..='f').contains(&c))
+    );
     assert!(identity.len() <= 32);
     assert!(listen_port != 0, "must used defined listen port");
     let mut cfg = zenoh::Config::default();
@@ -33,36 +36,32 @@ pub fn cfg(identity: &str, listen_port: u16) -> Result<zenoh::Config> {
     cfg.insert_json5("scouting/multicast/autoconnect", "[]")?;
     cfg.insert_json5("scouting/gossip/multihop", "true")?;
     cfg.insert_json5("adminspace/enabled", "true")?;
-    //cfg.insert_json5("transport/link/tx/batch_size", "9216")?;
+    cfg.insert_json5("transport/link/tx/batch_size", "9216")?;
     cfg.insert_json5("transport/link/rx/buffer_size", "16777216")?;
-    //cfg.insert_json5("timestamping/enabled", "true")?;
+    cfg.insert_json5("timestamping/enabled", "true")?;
     cfg.insert_json5("plugins/storage_manager/__required__", "true")?;
     cfg.insert_json5(
-        "plugins/storage_manager/storages/mem1",
-        r#"{
-            key_expr: "storage/mem1/**",
-            strip_prefix: "storage/mem1",
-            volume: "memory",
-            replication: {
-                interval: 2,
-            }
-        }"#,
+        "plugins/storage_manager/storages/mem1/key_expr",
+        &format!("\"{STORAGE_PREFIX}/**\""),
+    )?;
+    cfg.insert_json5(
+        "plugins/storage_manager/storages/mem1/strip_prefix",
+        &format!("\"{STORAGE_PREFIX}\""),
+    )?;
+    cfg.insert_json5("plugins/storage_manager/storages/mem1/volume", "\"memory\"")?;
+    cfg.insert_json5(
+        "plugins/storage_manager/storages/mem1/replication/interval",
+        "2",
     )?;
     Ok(cfg)
 }
 
 pub async fn open(
     cfg: zenoh::Config,
-    namespace: &str,
     listen_port: u16,
     discovery_service_port: u16,
 ) -> Result<Session> {
     assert!(listen_port != 0, "must used defined listen port");
-    let namespace: [u8; 8] = {
-        blake3::hash(namespace.as_bytes()).as_bytes()[..8]
-            .try_into()
-            .expect("8 is equal to 8")
-    };
     let mut plugins = PluginsManager::static_plugins_only();
     plugins.declare_static_plugin::<StoragesPlugin, _>("storage_manager", true);
     let mut runtime = zenoh::internal::runtime::RuntimeBuilder::new(cfg)
@@ -71,8 +70,7 @@ pub async fn open(
         .await?;
     let z = zenoh::session::init(runtime.clone().into()).await?;
     runtime.start().await?;
-    let mut discovery =
-        Discovery::new(z.zid(), namespace, listen_port, discovery_service_port).await?;
+    let mut discovery = Discovery::new(z.zid(), listen_port, discovery_service_port).await?;
     let _jh = Arc::new(AbortOnDrop(tokio::task::spawn(async move {
         loop {
             let Ok(discovered) = discovery.next().await.inspect_err(|e| {
@@ -99,10 +97,15 @@ pub async fn open(
                 .await;
         }
     })));
-    Ok(Session { z, _jh })
+    let liveliness_aggregator = spawn_liveliness_aggregator(&z)?;
+    Ok(Session {
+        z,
+        liveliness_aggregator,
+        _jh,
+    })
 }
 
-struct AbortOnDrop(JoinHandle<()>);
+pub struct AbortOnDrop(pub JoinHandle<()>);
 impl Drop for AbortOnDrop {
     fn drop(&mut self) {
         self.0.abort();
@@ -112,5 +115,6 @@ impl Drop for AbortOnDrop {
 #[derive(Clone)]
 pub struct Session {
     pub z: ZSession,
+    pub liveliness_aggregator: LivelinessAggregator,
     _jh: Arc<AbortOnDrop>,
 }

rust/networking/src/liveliness_aggregator.rs

@@ -0,0 +1,44 @@
+use std::{collections::HashSet, sync::Arc};
+
+use parking_lot::Mutex;
+use zenoh::{Result, Session, Wait, sample::SampleKind};
+
+pub fn spawn_liveliness_aggregator(session: &Session) -> Result<LivelinessAggregator> {
+    let store = Arc::new(Mutex::new(HashSet::default()));
+    session
+        .liveliness()
+        .declare_subscriber("live/*")
+        .history(true)
+        .callback({
+            let store = Arc::clone(&store);
+            move |sample| {
+                let Some(nid) = sample
+                    .key_expr()
+                    .to_string()
+                    .strip_prefix("live/")
+                    .map(str::to_owned)
+                else {
+                    return;
+                };
+                let mut mg = store.lock();
+                match sample.kind() {
+                    SampleKind::Put => mg.insert(nid),
+                    SampleKind::Delete => mg.remove(&nid),
+                };
+            }
+        })
+        .background()
+        .wait()?;
+    Ok(LivelinessAggregator { store })
+}
+
+#[derive(Clone)]
+pub struct LivelinessAggregator {
+    // need two arcs as the sub owns an arc to the store.
+    store: Arc<Mutex<HashSet<String>>>,
+}
+impl LivelinessAggregator {
+    pub fn dump(&self) -> HashSet<String> {
+        self.store.lock().clone()
+    }
+}

rust/networking/src/swarm.rs

@@ -70,14 +70,14 @@ impl Swarm {
                     token = discovery.recv_async() => {
                         if let Ok(token) = token {
                             let key_expr = token.key_expr().as_str().to_owned();
-                            let zid = key_expr.strip_prefix("live/");
+                            let nid = key_expr.strip_prefix("live/");
                             yield match token.kind() {
                                 SampleKind::Put => {
-                                    log::info!("discovered: {zid:?}");
+                                    log::info!("discovered: {nid:?}");
                                     FromSwarm::Discovered {}
                                 }
                                 SampleKind::Delete => {
-                                    log::info!("expired: {zid:?}");
+                                    log::info!("expired: {nid:?}");
                                     FromSwarm::Expired {}
                                 }
                             }
@@ -194,13 +194,12 @@ async fn on_message(
 
 pub async fn create_swarm(
     identity: &str,
-    namespace: &str,
     from_client: mpsc::Receiver<ToSwarm>,
     listen_port: u16,
     discovery_service_port: u16,
 ) -> Result<Swarm> {
     let cfg = crate::cfg(identity, listen_port)?;
-    let session = crate::open(cfg, namespace, listen_port, discovery_service_port).await?;
+    let session = crate::open(cfg, listen_port, discovery_service_port).await?;
     Ok(Swarm {
         session,
         from_client,

rust/networking/tests/runtime_polling.rs

@@ -0,0 +1,99 @@
+use std::time::Duration;
+
+use zenoh::Wait;
+
+fn unique_key(name: &str) -> String {
+    format!("test/zenoh-runtime-polling/{}/{}", std::process::id(), name)
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn get_and_recv_work_on_tokio_baseline() {
+    let session = zenoh::open(zenoh::Config::default())
+        .await
+        .expect("open session");
+
+    let key = unique_key("tokio-baseline");
+    let reply_key = key.clone();
+
+    let _queryable = session
+        .declare_queryable(key.clone())
+        .callback(move |query| {
+            query
+                .reply(reply_key.clone(), "hello-from-queryable")
+                .wait()
+                .expect("reply from queryable");
+        })
+        .await
+        .expect("declare queryable");
+
+    let replies = session.get(key).await.expect("get");
+
+    let reply = tokio::time::timeout(Duration::from_secs(5), replies.recv_async())
+        .await
+        .expect("timed out waiting for reply")
+        .expect("reply channel closed");
+
+    let sample = reply.result().expect("reply result was error");
+    let payload = sample
+        .payload()
+        .try_to_string()
+        .expect("payload should be utf8");
+
+    assert_eq!(payload.as_ref(), "hello-from-queryable");
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn get_and_recv_work_when_polled_by_smol_without_tokio_context() {
+    let session = zenoh::open(zenoh::Config::default())
+        .await
+        .expect("open session under tokio");
+
+    let key = unique_key("smol-no-tokio-context");
+    let reply_key = key.clone();
+
+    let _queryable = session
+        .declare_queryable(key.clone())
+        .callback(move |query| {
+            query
+                .reply(reply_key.clone(), "hello-from-queryable")
+                .wait()
+                .expect("reply from queryable");
+        })
+        .await
+        .expect("declare queryable under tokio");
+
+    let session_for_smol = session.clone();
+
+    let (tx, rx) = std::sync::mpsc::channel();
+
+    std::thread::spawn(move || {
+        // This thread was not entered by Tokio.
+        // If Zenoh's get/recv path requires an ambient Tokio Handle in the polling thread,
+        // this is where it should panic, hang, or error.
+        let result = {
+            smol::block_on(async move {
+                let replies = session_for_smol.get(key).await.expect("get under smol");
+
+                let reply = replies.recv_async().await.expect("reply channel closed");
+
+                let sample = reply.result().expect("reply result was error");
+                let payload = sample
+                    .payload()
+                    .try_to_string()
+                    .expect("payload should be utf8");
+
+                payload.to_string()
+            })
+        };
+
+        tx.send(result).expect("send test result");
+    });
+
+    let result = rx
+        .recv_timeout(Duration::from_secs(5))
+        .expect("smol thread timed out; likely hung waiting for get/reply");
+
+    let payload = result;
+
+    assert_eq!(payload, "hello-from-queryable");
+}

src/exo/__init__.py

@@ -1,3 +0,0 @@
-from importlib.metadata import version
-
-__version__ = version("exo")

src/exo/api/main.py

@@ -13,6 +13,7 @@ from uuid import uuid4
 
 import anyio
 from anyio import BrokenResourceError, ClosedResourceError
+from exo_rs import SessionHandle
 from fastapi import FastAPI, File, Form, HTTPException, Query, Request, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
@@ -22,6 +23,7 @@ from hypercorn.config import Config
 from hypercorn.typing import ASGIFramework
 from hypercorn.utils import LifespanTimeoutError, ShutdownError
 from loguru import logger
+from pydantic import ValidationError
 
 from exo.api.adapters.chat_completions import (
     chat_request_to_text_generation,
@@ -153,14 +155,11 @@ from exo.shared.types.chunks import (
     ToolCallChunk,
 )
 from exo.shared.types.commands import (
-    AddCustomModelCard,
     CancelDownload,
     Command,
     CreateInstance,
-    DeleteCustomModelCard,
     DeleteDownload,
     DeleteInstance,
-    DeleteInstanceLink,
     DownloadCommand,
     ForwarderCommand,
     ForwarderDownloadCommand,
@@ -168,7 +167,6 @@ from exo.shared.types.commands import (
     ImageGeneration,
     PlaceInstance,
     SendInputChunk,
-    SetInstanceLink,
     StartDownload,
     TaskCancelled,
     TaskFinished,
@@ -246,6 +244,7 @@ class API:
         download_command_sender: Sender[ForwarderDownloadCommand],
         # This lets us pause the API if an election is running
         election_receiver: Receiver[ElectionMessage],
+        session_handle: SessionHandle,
     ) -> None:
         self.state = State()
         self._event_log = DiskEventLog(_API_EVENT_LOG_DIR)
@@ -258,6 +257,8 @@ class API:
         self.last_completed_election: int = 0
         self.port = port
         self._sent_image_hashes: set[str] = set()
+        self.aggregator = session_handle.last_value_aggregator("metrics")
+        self.storage = session_handle.storage_interface()
 
         self.paused: bool = False
         self.paused_ev: anyio.Event = anyio.Event()
@@ -409,10 +410,12 @@ class API:
         self.app.post("/onboarding")(self.complete_onboarding)
 
     def get_state(self, path: str = ""):
+        state = self.state.with_aggregator(self.aggregator)
+
         if path == "":
-            return self.state
+            return state
         try:
-            x = self.state.model_dump(by_alias=True)
+            x: Any = state.model_dump(by_alias=True)
             for attr in path.split("/"):
                 if attr != "":
                     if isinstance(x, dict):
@@ -476,6 +479,7 @@ class API:
         model_card = await ModelCard.load(model_id)
 
         try:
+            state = self.state.with_aggregator(self.aggregator)
             placements = get_instance_placements(
                 PlaceInstance(
                     model_card=model_card,
@@ -483,13 +487,13 @@ class API:
                     instance_meta=instance_meta,
                     min_nodes=min_nodes,
                 ),
-                node_memory=self.state.node_memory,
-                node_network=self.state.node_network,
-                node_backends=self.state.node_backends,
-                topology=self.state.topology,
-                current_instances=self.state.instances,
-                download_status=self.state.downloads,
-                node_rdma_ctl=self.state.node_rdma_ctl,
+                node_memory=state.node_memory,
+                node_network=state.node_network,
+                node_backends=state.node_backends,
+                topology=state.topology,
+                current_instances=state.instances,
+                download_status=state.downloads,
+                node_rdma_ctl=state.node_rdma_ctl,
             )
         except ValueError as exc:
             raise HTTPException(status_code=400, detail=str(exc)) from exc
@@ -514,8 +518,9 @@ class API:
         seen: set[tuple[ModelId, Sharding, InstanceMeta, int]] = set()
         previews: list[PlacementPreview] = []
         required_nodes = set(node_ids) if node_ids else None
+        state = self.state.with_aggregator(self.aggregator)
 
-        if len(list(self.state.topology.list_nodes())) == 0:
+        if len(list(state.topology.list_nodes())) == 0:
             return PlacementPreviewResponse(previews=[])
 
         try:
@@ -530,9 +535,7 @@ class API:
                 instance_combinations.extend(
                     [
                         (sharding, instance_meta, i)
-                        for i in range(
-                            1, len(list(self.state.topology.list_nodes())) + 1
-                        )
+                        for i in range(1, len(list(state.topology.list_nodes())) + 1)
                     ]
                 )
         # TODO: PDD
@@ -547,14 +550,14 @@ class API:
                         instance_meta=instance_meta,
                         min_nodes=min_nodes,
                     ),
-                    node_memory=self.state.node_memory,
-                    node_network=self.state.node_network,
-                    node_backends=self.state.node_backends,
-                    topology=self.state.topology,
-                    current_instances=self.state.instances,
+                    node_memory=state.node_memory,
+                    node_network=state.node_network,
+                    node_backends=state.node_backends,
+                    topology=state.topology,
+                    current_instances=state.instances,
                     required_nodes=required_nodes,
-                    download_status=self.state.downloads,
-                    node_rdma_ctl=self.state.node_rdma_ctl,
+                    download_status=state.downloads,
+                    node_rdma_ctl=state.node_rdma_ctl,
                 )
             except ValueError as exc:
                 if (model_card.model_id, sharding, instance_meta, 0) not in seen:
@@ -593,7 +596,7 @@ class API:
 
             instance = new_instances[0]
             shard_assignments = instance.shard_assignments
-            placement_node_ids = list(shard_assignments.node_to_runner.keys())
+            placement_node_ids = list(s.node_id for s in shard_assignments.shards)
 
             memory_delta_by_node: dict[str, int] = {}
             if placement_node_ids:
@@ -696,9 +699,16 @@ class API:
         return {"disaggregation": ENABLE_DISAGGREGATION}
 
     async def list_instance_links(self) -> list[InstanceLink]:
+        links: list[InstanceLink] = []
         if not ENABLE_DISAGGREGATION:
-            return []
-        return list(self.state.instance_links.values())
+            return links
+        for _, value in (await self.storage.dump("custom_model_cards/")).items():
+            try:
+                link = InstanceLink.model_validate_json(value)
+            except ValidationError:
+                continue
+            links.append(link)
+        return links
 
     async def create_instance_link(
         self, body: InstanceLinkBody
@@ -715,25 +725,22 @@ class API:
     async def _set_instance_link(
         self, link_id: InstanceLinkId, body: InstanceLinkBody
     ) -> InstanceLinkResponse:
-        command = SetInstanceLink(
-            link_id=link_id,
-            prefill_instances=list(body.prefill_instances),
-            decode_instances=list(body.decode_instances),
-        )
-        await self._send(command)
-        return InstanceLinkResponse(
-            message="Command received.", command_id=command.command_id
+        await self.storage.put(
+            f"instance_links/{link_id}",
+            InstanceLink(
+                link_id=link_id,
+                prefill_instances=body.prefill_instances,
+                decode_instances=body.decode_instances,
+            ).model_dump_json(),
         )
+        return InstanceLinkResponse(message="Command received.")
 
     async def delete_instance_link(
         self, link_id: InstanceLinkId
     ) -> InstanceLinkResponse:
         _require_disaggregation_enabled()
-        command = DeleteInstanceLink(link_id=link_id)
-        await self._send(command)
-        return InstanceLinkResponse(
-            message="Command received.", command_id=command.command_id
-        )
+        await self.storage.delete(f"instance_links/{link_id}")
+        return InstanceLinkResponse(message="Command received.")
 
     async def cancel_command(self, command_id: CommandId) -> CancelCommandResponse:
         """Cancel an active command by closing its stream and notifying workers."""
@@ -791,7 +798,11 @@ class API:
     async def _collect_text_generation_with_stats(
         self, command_id: CommandId
     ) -> BenchChatCompletionResponse:
-        sampler = PowerSampler(get_node_system=lambda: self.state.node_system)
+        sampler = PowerSampler(
+            get_node_system=lambda: self.state.with_aggregator(
+                self.aggregator
+            ).node_system
+        )
         text_parts: list[str] = []
         tool_calls: list[ToolCall] = []
         model: ModelId | None = None
@@ -1314,7 +1325,11 @@ class API:
         num_images: int,
         response_format: str,
     ) -> BenchImageGenerationResponse:
-        sampler = PowerSampler(get_node_system=lambda: self.state.node_system)
+        sampler = PowerSampler(
+            get_node_system=lambda: self.state.with_aggregator(
+                self.aggregator
+            ).node_system
+        )
         images: list[ImageData] = []
         stats: ImageGenerationStats | None = None
         async with anyio.create_task_group() as tg:
@@ -1609,20 +1624,18 @@ class API:
         return JSONResponse(content="Ollama is running")
 
     async def ollama_chat(
-        self, request: Request
+        self, request: OllamaChatRequest
     ) -> OllamaChatResponse | StreamingResponse:
         """Ollama Chat API — accepts JSON regardless of Content-Type."""
-        body = await request.body()
-        payload = OllamaChatRequest.model_validate_json(body)
-        task_params = ollama_request_to_text_generation(payload)
-        validated_model = await self._validate_model_has_instance(
+        task_params = ollama_request_to_text_generation(request)
+        resolved_model = await self._validate_model_has_instance(
             ModelId(task_params.model)
         )
-        task_params = task_params.model_copy(update={"model": validated_model})
+        task_params = task_params.model_copy(update={"model": resolved_model})
 
         command = await self._send_text_generation_with_images(task_params)
 
-        if payload.stream:
+        if request.stream:
             return StreamingResponse(
                 generate_ollama_chat_stream(
                     command.command_id,
@@ -1645,20 +1658,18 @@ class API:
             )
 
     async def ollama_generate(
-        self, request: Request
+        self, request: OllamaGenerateRequest
     ) -> OllamaGenerateResponse | StreamingResponse:
         """Ollama Generate API — accepts JSON regardless of Content-Type."""
-        body = await request.body()
-        payload = OllamaGenerateRequest.model_validate_json(body)
-        task_params = ollama_generate_request_to_text_generation(payload)
-        validated_model = await self._validate_model_has_instance(
+        task_params = ollama_generate_request_to_text_generation(request)
+        resolved_model = await self._validate_model_has_instance(
             ModelId(task_params.model)
         )
-        task_params = task_params.model_copy(update={"model": validated_model})
+        task_params = task_params.model_copy(update={"model": resolved_model})
 
         command = await self._send_text_generation_with_images(task_params)
 
-        if payload.stream:
+        if request.stream:
             return StreamingResponse(
                 generate_ollama_generate_stream(
                     command.command_id,
@@ -1713,11 +1724,9 @@ class API:
             ]
         )
 
-    async def ollama_show(self, request: Request) -> OllamaShowResponse:
+    async def ollama_show(self, request: OllamaShowRequest) -> OllamaShowResponse:
         """Returns model information in Ollama show format."""
-        body = await request.body()
-        payload = OllamaShowRequest.model_validate_json(body)
-        model_name = payload.name or payload.model
+        model_name = request.name or request.model
         if not model_name:
             raise HTTPException(status_code=400, detail="name or model is required")
         try:
@@ -1778,7 +1787,7 @@ class API:
         """Calculate total available memory across all nodes in bytes."""
         total_available = Memory()
 
-        for memory in self.state.node_memory.values():
+        for memory in self.state.with_aggregator(self.aggregator).node_memory.values():
             total_available += memory.ram_available
 
         return total_available
@@ -1827,11 +1836,8 @@ class API:
                 status_code=400, detail=f"Failed to fetch model: {exc}"
             ) from exc
 
-        await self.command_sender.send(
-            ForwarderCommand(
-                origin=self._system_id,
-                command=AddCustomModelCard(model_card=card),
-            )
+        await self.storage.put(
+            f"custom_model_cards/{card.model_id.normalize()}", card.model_dump_json()
         )
 
         # Immediately update the local cache so the subsequent GET /models
@@ -1856,12 +1862,7 @@ class API:
         if card is None or not card.is_custom:
             raise HTTPException(status_code=404, detail="Custom model card not found")
 
-        await self.command_sender.send(
-            ForwarderCommand(
-                origin=self._system_id,
-                command=DeleteCustomModelCard(model_id=model_id),
-            )
-        )
+        await self.storage.delete(f"custom_model_cards/{card.model_id.normalize()}")
 
         return JSONResponse(
             {"message": "Model card deleted", "model_id": str(model_id)}

src/exo/api/types/api.py

@@ -329,7 +329,6 @@ class InstanceLinkBody(BaseModel):
 
 class InstanceLinkResponse(BaseModel):
     message: str
-    command_id: CommandId
 
 
 ImageSize = Literal[

src/exo/main.py

@@ -10,18 +10,17 @@ from typing import Self
 import anyio
 from anyio.lowlevel import checkpoint as anyio_checkpoint
 from daemon import DaemonContext  # pyright: ignore[reportMissingTypeStubs]
-from exo_rs import Pidfile, PidfileError
+from exo_rs import Pidfile, PidfileError, SessionHandle
 from loguru import logger
 from pydantic import PositiveInt
 
 import exo.routing.topics as topics
-from exo import __version__
 from exo.api.main import API
 from exo.download.coordinator import DownloadCoordinator
 from exo.download.impl_shard_downloader import exo_shard_downloader
 from exo.master.main import Master
 from exo.routing.event_router import EventRouter
-from exo.routing.router import Router, get_node_zid
+from exo.routing.router import Router
 from exo.shared.constants import EXO_DEFAULT_MODELS_DIR, EXO_LOG, EXO_PID_FILE
 from exo.shared.election import Election, ElectionResult
 from exo.shared.logging import logger_cleanup, logger_setup
@@ -47,18 +46,17 @@ class Node:
     node_id: NodeId
     offline: bool
     _api_port: int
+    _sh: SessionHandle
     _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
 
     @classmethod
     async def create(cls, args: "Args") -> Self:
-        node_id = get_node_zid()
+        identity = os.urandom(16).hex().lstrip("0")
+        node_id = NodeId(identity)
         session_id = SessionId(master_node_id=node_id, election_clock=0)
-        router = Router.create(
-            node_id,
-            namespace=args.namespace,
-            listen_port=args.zenoh_port,
-            discovery_service_port=args.discovery_port,
-        )
+        session_handle, _nh = SessionHandle.new(identity, args.zenoh_port, 52413)
+        router = Router(_nh)
+
         await router.register_topic(topics.GLOBAL_EVENTS)
         await router.register_topic(topics.LOCAL_EVENTS)
         await router.register_topic(topics.COMMANDS)
@@ -97,6 +95,7 @@ class Node:
                 command_sender=router.sender(topics.COMMANDS),
                 download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
                 election_receiver=router.receiver(topics.ELECTION_MESSAGES),
+                session_handle=session_handle,
             )
         else:
             api = None
@@ -108,6 +107,7 @@ class Node:
                 event_sender=event_router.sender(),
                 command_sender=router.sender(topics.COMMANDS),
                 download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
+                session_handle=session_handle,
                 api_port=args.api_port,
             )
         else:
@@ -122,6 +122,8 @@ class Node:
             local_event_receiver=router.receiver(topics.LOCAL_EVENTS),
             command_receiver=router.receiver(topics.COMMANDS),
             download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
+            aggregator=session_handle.last_value_aggregator("metrics"),
+            storage=session_handle.storage_interface(),
         )
 
         er_send, er_recv = channel[ElectionResult]()
@@ -150,6 +152,7 @@ class Node:
             node_id,
             args.offline,
             args.api_port,
+            session_handle,
         )
 
     async def run(self):
@@ -225,6 +228,8 @@ class Node:
                         download_command_sender=self.router.sender(
                             topics.DOWNLOAD_COMMANDS
                         ),
+                        aggregator=self._sh.last_value_aggregator("metrics"),
+                        storage=self._sh.storage_interface(),
                     )
                     self._tg.start_soon(self.master.run)
                 elif (
@@ -264,6 +269,7 @@ class Node:
                             download_command_sender=self.router.sender(
                                 topics.DOWNLOAD_COMMANDS
                             ),
+                            session_handle=self._sh,
                             api_port=self._api_port,
                         )
                         self._tg.start_soon(self.worker.run)
@@ -378,21 +384,18 @@ def main_inner(args: "Args"):
 
 
 class Args(FrozenModel):
-    verbosity: int = 0
-    force_master: bool = False
-    spawn_api: bool = False
-    api_port: PositiveInt = 52415
-    tb_only: bool = False
+    verbosity: int
+    force_master: bool
+    spawn_api: bool
+    api_port: PositiveInt
     no_worker: bool = False
     no_downloads: bool = False
-    offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true"
-    no_batch: bool = False
+    offline: bool
+    no_batch: bool
     fast_synch: bool | None = None  # None = auto, True = force on, False = force off
-    legacy_daemon: bool = False
+    legacy_daemon: bool
     bootstrap_peers: list[str] = []
-    namespace: str
     zenoh_port: int
-    discovery_port: int
 
     @classmethod
     def parse(cls) -> Self:
@@ -464,26 +467,12 @@ class Args(FrozenModel):
             dest="bootstrap_peers",
             help="Comma-separated libp2p multiaddrs to dial on startup (env: EXO_BOOTSTRAP_PEERS)",
         )
-        parser.add_argument(
-            "--namespace",
-            type=str,
-            default=__version__,
-            dest="namespace",
-            help="Discovery namespace, nodes with different namespaces will not connect.",
-        )
         parser.add_argument(
             "--zenoh-port",
             type=int,
             default=52414,
             dest="zenoh_port",
-            help="Fixed TCP port for zenoh to listen.",
-        )
-        parser.add_argument(
-            "--discovery-port",
-            type=int,
-            default=52413,
-            dest="discovery_port",
-            help="Fixed UDP port for the discovery service.",
+            help="Fixed port for zenoh to listen on.",
         )
         fast_synch_group = parser.add_mutually_exclusive_group()
         fast_synch_group.add_argument(

src/exo/master/main.py

@@ -1,7 +1,9 @@
 from datetime import datetime, timedelta, timezone
 
 import anyio
+from exo_rs import LVAggregator, Storage
 from loguru import logger
+from pydantic import ValidationError
 
 from exo.master.placement import (
     add_instance_to_placements,
@@ -18,11 +20,8 @@ from exo.routing.event_router import (
 from exo.shared.apply import apply
 from exo.shared.constants import EXO_EVENT_LOG_DIR, EXO_TRACING_ENABLED
 from exo.shared.types.commands import (
-    AddCustomModelCard,
     CreateInstance,
-    DeleteCustomModelCard,
     DeleteInstance,
-    DeleteInstanceLink,
     ForwarderCommand,
     ForwarderDownloadCommand,
     ImageEdits,
@@ -30,7 +29,6 @@ from exo.shared.types.commands import (
     PlaceInstance,
     RequestEventLog,
     SendInputChunk,
-    SetInstanceLink,
     TaskCancelled,
     TaskFinished,
     TestCommand,
@@ -38,15 +36,11 @@ from exo.shared.types.commands import (
 )
 from exo.shared.types.common import CommandId, NodeId, SessionId, SystemId
 from exo.shared.types.events import (
-    CustomModelCardAdded,
-    CustomModelCardDeleted,
     Event,
     GlobalForwarderEvent,
     IndexedEvent,
     InputChunkReceived,
     InstanceDeleted,
-    InstanceLinkCreated,
-    InstanceLinkDeleted,
     LocalForwarderEvent,
     NodeGatheredInfo,
     NodeTimedOut,
@@ -79,16 +73,18 @@ from exo.utils.event_buffer import MultiSourceBuffer
 from exo.utils.task_group import TaskGroup
 
 
-def _prefill_endpoint_for(state: State, decode_instance_id: InstanceId) -> str | None:
+def _prefill_endpoint_for(
+    state: State, instance_links: list[InstanceLink], decode_instance_id: InstanceId
+) -> str | None:
     decode = state.instances.get(decode_instance_id)
     if decode is None:
         return None
-    decode_node = next(iter(decode.shard_assignments.node_to_runner.keys()), None)
-    if decode_node is None:
-        return None
+    decode_node = decode.shard_assignments.shards[
+        decode.shard_assignments.primary_output_node
+    ].node_id
 
     sources: set[InstanceId] = set()
-    for link in state.instance_links.values():
+    for link in instance_links:
         if decode_instance_id in link.decode_instances:
             sources.update(link.prefill_instances)
     sources.discard(decode_instance_id)
@@ -106,7 +102,7 @@ def _prefill_endpoint_for(state: State, decode_instance_id: InstanceId) -> str |
         instance = state.instances.get(src_id)
         if instance is None:
             continue
-        for node_id, runner_id in instance.shard_assignments.node_to_runner.items():
+        for node_id, runner_id, _ in instance.shard_assignments.shards:
             port = state.prefill_server_ports.get(runner_id)
             if port is None:
                 continue
@@ -130,6 +126,8 @@ class Master:
         local_event_receiver: Receiver[LocalForwarderEvent],
         global_event_sender: Sender[GlobalForwarderEvent],
         download_command_sender: Sender[ForwarderDownloadCommand],
+        aggregator: LVAggregator,
+        storage: Storage,
     ):
         self.node_id = node_id
         self.session_id = session_id
@@ -145,7 +143,7 @@ class Master:
         self._multi_buffer = MultiSourceBuffer[SystemId, Event]()
         self._event_log = DiskEventLog(EXO_EVENT_LOG_DIR / "master")
         self._pending_traces: dict[TaskId, dict[int, list[TraceEventData]]] = {}
-        self._expected_ranks: dict[TaskId, set[int]] = {}
+        self._world_sizes: dict[TaskId, int] = {}
 
     async def run(self):
         logger.info("Starting Master")
@@ -182,10 +180,21 @@ class Master:
                             pass
                         case TextGeneration():
                             # set-difference => prefill-only nodes
+                            instance_links: list[InstanceLink] = []
                             prefill_only: set[InstanceId] = set()
-                            for link in self.state.instance_links.values():
+                            for _, link in (
+                                await self.storage.dump("instance_links/")
+                            ).items():
+                                try:
+                                    instance_links.append(
+                                        InstanceLink.model_validate_json(link)
+                                    )
+                                except ValidationError:
+                                    continue
+
+                            for link in instance_links:
                                 prefill_only.update(link.prefill_instances)
-                            for link in self.state.instance_links.values():
+                            for link in instance_links:
                                 prefill_only.difference_update(link.decode_instances)
 
                             for instance in self.state.instances.values():
@@ -225,7 +234,9 @@ class Master:
                             params = command.task_params.model_copy(
                                 update={
                                     "prefill_endpoint": _prefill_endpoint_for(
-                                        self.state, decode_instance_id
+                                        self.state.with_aggregator(self.aggregator),
+                                        instance_links,
+                                        decode_instance_id,
                                     ),
                                 }
                             )
@@ -293,11 +304,9 @@ class Master:
                                     selected_instance_id
                                 )
                                 if selected_instance:
-                                    ranks = set(
-                                        shard.device_rank
-                                        for shard in selected_instance.shard_assignments.runner_to_shard.values()
+                                    self._world_sizes[task_id] = len(
+                                        selected_instance.shard_assignments.shards
                                     )
-                                    self._expected_ranks[task_id] = ranks
                         case ImageEdits():
                             for instance in self.state.instances.values():
                                 if (
@@ -349,11 +358,9 @@ class Master:
                                     selected_instance_id
                                 )
                                 if selected_instance:
-                                    ranks = set(
-                                        shard.device_rank
-                                        for shard in selected_instance.shard_assignments.runner_to_shard.values()
+                                    self._world_sizes[task_id] = len(
+                                        selected_instance.shard_assignments.shards
                                     )
-                                    self._expected_ranks[task_id] = ranks
                         case DeleteInstance():
                             placement = delete_instance(command, self.state.instances)
                             transition_events = get_transition_events(
@@ -369,15 +376,16 @@ class Master:
                                 )
                             generated_events.extend(transition_events)
                         case PlaceInstance():
+                            state = self.state.with_aggregator(self.aggregator)
                             placement = place_instance(
                                 command,
-                                self.state.topology,
-                                self.state.instances,
-                                self.state.node_memory,
-                                self.state.node_network,
-                                self.state.node_backends,
-                                download_status=self.state.downloads,
-                                node_rdma_ctl=self.state.node_rdma_ctl,
+                                state.topology,
+                                state.instances,
+                                state.node_memory,
+                                state.node_network,
+                                state.node_backends,
+                                download_status=state.downloads,
+                                node_rdma_ctl=state.node_rdma_ctl,
                             )
                             transition_events = get_transition_events(
                                 self.state.instances, placement, self.state.tasks
@@ -428,29 +436,6 @@ class Master:
                                     f"Finished command {command.finished_command_id} finished"
                                 )
 
-                        case AddCustomModelCard():
-                            generated_events.append(
-                                CustomModelCardAdded(model_card=command.model_card)
-                            )
-                        case DeleteCustomModelCard():
-                            generated_events.append(
-                                CustomModelCardDeleted(model_id=command.model_id)
-                            )
-                        case SetInstanceLink():
-                            link = InstanceLink(
-                                link_id=command.link_id,
-                                prefill_instances=list(
-                                    dict.fromkeys(command.prefill_instances)
-                                ),
-                                decode_instances=list(
-                                    dict.fromkeys(command.decode_instances)
-                                ),
-                            )
-                            generated_events.append(InstanceLinkCreated(link=link))
-                        case DeleteInstanceLink():
-                            generated_events.append(
-                                InstanceLinkDeleted(link_id=command.link_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
@@ -471,9 +456,11 @@ class Master:
     async def _plan(self) -> None:
         while True:
             # kill broken instances
-            connected_node_ids = set(self.state.topology.list_nodes())
+            connected_node_ids = set(
+                self.state.with_aggregator(self.aggregator).topology.list_nodes()
+            )
             for instance_id, instance in self.state.instances.items():
-                for node_id in instance.shard_assignments.node_to_runner:
+                for node_id, _, _ in instance.shard_assignments.shards:
                     if node_id not in connected_node_ids:
                         await self.event_sender.send(
                             InstanceDeleted(instance_id=instance_id)
@@ -481,7 +468,9 @@ class Master:
                         break
 
             # time out dead nodes
-            for node_id, time in self.state.last_seen.items():
+            for node_id, time in self.state.with_aggregator(
+                self.aggregator
+            ).last_seen.items():
                 now = datetime.now(tz=timezone.utc)
                 if now - time > timedelta(seconds=30):
                     logger.info(f"Manually removing node {node_id} due to inactivity")
@@ -540,9 +529,8 @@ class Master:
         self._pending_traces[task_id][event.rank] = event.traces
 
         if (
-            task_id in self._expected_ranks
-            and set(self._pending_traces[task_id].keys())
-            >= self._expected_ranks[task_id]
+            task_id in self._world_sizes
+            and len(self._pending_traces[task_id]) >= self._world_sizes[task_id]
         ):
             await self._merge_and_save_traces(task_id)
 
@@ -556,5 +544,5 @@ class Master:
         )
 
         del self._pending_traces[task_id]
-        if task_id in self._expected_ranks:
-            del self._expected_ranks[task_id]
+        if task_id in self._world_sizes:
+            del self._world_sizes[task_id]

src/exo/master/placement.py

@@ -262,20 +262,7 @@ def place_instance(
 
     match command.instance_meta:
         case InstanceMeta.MlxJaccl:
-            # TODO(evan): shard assignments should contain information about ranks, this is ugly
-            def get_device_rank(node_id: NodeId) -> int:
-                runner_id = shard_assignments.node_to_runner[node_id]
-                shard_metadata = shard_assignments.runner_to_shard.get(runner_id)
-                assert shard_metadata is not None
-                return shard_metadata.device_rank
-
-            zero_node_ids = [
-                node_id
-                for node_id in selected_cycle.node_ids
-                if get_device_rank(node_id) == 0
-            ]
-            assert len(zero_node_ids) == 1
-            coordinator_node_id = zero_node_ids[0]
+            coordinator_node_id = shard_assignments.shards[0].node_id
 
             mlx_jaccl_devices = get_mlx_jaccl_devices_matrix(
                 [node_id for node_id in selected_cycle],
@@ -376,10 +363,10 @@ def cancel_unnecessary_downloads(
     active_models = set(
         (
             node_id,
-            instance.shard_assignments.runner_to_shard[runner_id].model_card.model_id,
+            instance.shard_assignments.model_id,
         )
         for instance in instances.values()
-        for node_id, runner_id in instance.shard_assignments.node_to_runner.items()
+        for node_id, _, _ in instance.shard_assignments.shards
     )
     for pair in currently_downloading:
         if pair not in active_models:

src/exo/master/placement_utils.py

@@ -8,12 +8,11 @@ from exo.shared.types.common import Host, NodeId
 from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
 from exo.shared.types.topology import Cycle, RDMAConnection, SocketConnection
-from exo.shared.types.worker.runners import RunnerId, ShardAssignments
+from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardWithId
 from exo.shared.types.worker.shards import (
     CfgShardMetadata,
     PipelineShardMetadata,
     Sharding,
-    ShardMetadata,
     TensorShardMetadata,
 )
 
@@ -152,27 +151,27 @@ def _get_shard_assignments_for_cfg_parallel(
     _validate_cycle(cycle)
 
     world_size = len(cycle)
-    cfg_world_size = 2
-    pipeline_world_size = world_size // cfg_world_size
+    pipeline_world_size = world_size // 2
 
     # Allocate layers for one pipeline group (both groups run the same layers)
     pipeline_node_ids = cycle.node_ids[:pipeline_world_size]
     pipeline_memory = _compute_total_memory(pipeline_node_ids, node_memory)
+
+    # nb: only validates the forward path...
     layer_allocations = _allocate_and_validate_layers(
         pipeline_node_ids, node_memory, pipeline_memory, model_card
     )
 
     # Ring topology: group 0 ascending [0,1,2,...], group 1 descending [...,2,1,0]
     # This places both last stages as neighbors for CFG exchange.
-    position_to_cfg_pipeline = [(0, r) for r in range(pipeline_world_size)] + [
-        (1, r) for r in reversed(range(pipeline_world_size))
-    ]
+    position_to_cfg_pipeline = list(range(pipeline_world_size)) + list(
+        reversed(range(pipeline_world_size))
+    )
 
-    runner_to_shard: dict[RunnerId, ShardMetadata] = {}
-    node_to_runner: dict[NodeId, RunnerId] = {}
+    shards: list[ShardWithId] = []
 
     for device_rank, node_id in enumerate(cycle.node_ids):
-        cfg_rank, pipeline_rank = position_to_cfg_pipeline[device_rank]
+        pipeline_rank = position_to_cfg_pipeline[device_rank]
         layers_before = sum(layer_allocations[:pipeline_rank])
         node_layers = layer_allocations[pipeline_rank]
 
@@ -183,20 +182,15 @@ def _get_shard_assignments_for_cfg_parallel(
             start_layer=layers_before,
             end_layer=layers_before + node_layers,
             n_layers=model_card.n_layers,
-            cfg_rank=cfg_rank,
-            cfg_world_size=cfg_world_size,
-            pipeline_rank=pipeline_rank,
-            pipeline_world_size=pipeline_world_size,
         )
 
         runner_id = RunnerId()
-        runner_to_shard[runner_id] = shard
-        node_to_runner[node_id] = runner_id
+        shards.append(ShardWithId(node_id, runner_id, shard))
 
     return ShardAssignments(
         model_id=model_card.model_id,
-        runner_to_shard=runner_to_shard,
-        node_to_runner=node_to_runner,
+        shards=shards,
+        primary_output_node=pipeline_world_size - 1,
     )
 
 
@@ -208,13 +202,13 @@ def _get_shard_assignments_for_pure_pipeline(
     """Create shard assignments for pure pipeline execution."""
     _validate_cycle(cycle)
     total_memory = _compute_total_memory(cycle.node_ids, node_memory)
+    world_size = len(cycle)
 
     layer_allocations = _allocate_and_validate_layers(
         cycle.node_ids, node_memory, total_memory, model_card
     )
 
-    runner_to_shard: dict[RunnerId, ShardMetadata] = {}
-    node_to_runner: dict[NodeId, RunnerId] = {}
+    shards: list[ShardWithId] = []
 
     for pipeline_rank, node_id in enumerate(cycle.node_ids):
         layers_before = sum(layer_allocations[:pipeline_rank])
@@ -223,20 +217,17 @@ def _get_shard_assignments_for_pure_pipeline(
         shard = PipelineShardMetadata(
             model_card=model_card,
             device_rank=pipeline_rank,
-            world_size=len(cycle),
+            world_size=world_size,
             start_layer=layers_before,
             end_layer=layers_before + node_layers,
             n_layers=model_card.n_layers,
         )
 
         runner_id = RunnerId()
-        runner_to_shard[runner_id] = shard
-        node_to_runner[node_id] = runner_id
+        shards.append(ShardWithId(node_id, runner_id, shard))
 
     return ShardAssignments(
-        model_id=model_card.model_id,
-        runner_to_shard=runner_to_shard,
-        node_to_runner=node_to_runner,
+        model_id=model_card.model_id, shards=shards, primary_output_node=world_size - 1
     )
 
 
@@ -246,8 +237,7 @@ def get_shard_assignments_for_tensor_parallel(
 ):
     total_layers = model_card.n_layers
     world_size = len(cycle)
-    runner_to_shard: dict[RunnerId, ShardMetadata] = {}
-    node_to_runner: dict[NodeId, RunnerId] = {}
+    shards: list[ShardWithId] = []
 
     for i, node_id in enumerate(cycle):
         shard = TensorShardMetadata(
@@ -260,14 +250,10 @@ def get_shard_assignments_for_tensor_parallel(
         )
 
         runner_id = RunnerId()
-
-        runner_to_shard[runner_id] = shard
-        node_to_runner[node_id] = runner_id
+        shards.append(ShardWithId(node_id, runner_id, shard))
 
     shard_assignments = ShardAssignments(
-        model_id=model_card.model_id,
-        runner_to_shard=runner_to_shard,
-        node_to_runner=node_to_runner,
+        model_id=model_card.model_id, shards=shards, primary_output_node=world_size - 1
     )
 
     return shard_assignments

src/exo/master/tests/test_master.py

@@ -6,7 +6,6 @@ import pytest
 from loguru import logger
 
 from exo.master.main import Master
-from exo.routing.router import get_node_zid
 from exo.shared.models.model_cards import ModelCard, ModelTask
 from exo.shared.types.backends import Backend
 from exo.shared.types.commands import (
@@ -16,7 +15,7 @@ from exo.shared.types.commands import (
     PlaceInstance,
     TextGeneration,
 )
-from exo.shared.types.common import ModelId, SessionId, SystemId
+from exo.shared.types.common import ModelId, NodeId, SessionId, SystemId
 from exo.shared.types.events import (
     Event,
     GlobalForwarderEvent,
@@ -42,14 +41,34 @@ from exo.shared.types.worker.instances import (
     MlxRingInstance,
     ShardAssignments,
 )
+from exo.shared.types.worker.runners import ShardWithId
 from exo.shared.types.worker.shards import PipelineShardMetadata, Sharding
 from exo.utils.channels import channel
 from exo.utils.info_gatherer.info_gatherer import NodeBackends
 
 
+class MockAggregator:
+    def dump(self) -> dict[str, str]:
+        return {}
+
+
+class MockStorage:
+    async def get(self, _: str) -> None:
+        return None
+
+    async def put(self, _1: str, _2: str) -> None:
+        return None
+
+    async def delete(self, _: str) -> None:
+        return None
+
+    async def dump(self, _: str) -> dict[str, str]:
+        return {}
+
+
 @pytest.mark.asyncio
 async def test_master():
-    node_id = get_node_zid()
+    node_id = NodeId("yoooo")
     session_id = SessionId(master_node_id=node_id, election_clock=0)
 
     ge_sender, global_event_receiver = channel[GlobalForwarderEvent]()
@@ -94,6 +113,8 @@ async def test_master():
         local_event_receiver=le_receiver,
         command_receiver=co_receiver,
         download_command_sender=fcds,
+        aggregator=MockAggregator(),  # pyright: ignore[reportArgumentType]
+        storage=MockStorage(),  # pyright: ignore[reportArgumentType]
     )
     logger.info("run the master")
     async with anyio.create_task_group() as tg:
@@ -205,29 +226,33 @@ async def test_master():
         assert isinstance(events[2].event, InstanceCreated)
         created_instance = events[2].event.instance
         assert isinstance(created_instance, MlxRingInstance)
-        runner_id = list(created_instance.shard_assignments.runner_to_shard.keys())[0]
+        runner_id = created_instance.shard_assignments.shards[0].runner_id
         # Validate the shard assignments
         expected_shard_assignments = ShardAssignments(
             model_id=ModelId("llama-3.2-1b"),
-            runner_to_shard={
-                (runner_id): PipelineShardMetadata(
-                    start_layer=0,
-                    end_layer=16,
-                    n_layers=16,
-                    model_card=ModelCard(
-                        model_id=ModelId("llama-3.2-1b"),
+            shards=[
+                ShardWithId(
+                    node_id,
+                    runner_id,
+                    PipelineShardMetadata(
+                        start_layer=0,
+                        end_layer=16,
                         n_layers=16,
-                        storage_size=Memory.from_bytes(678948),
-                        hidden_size=7168,
-                        supports_tensor=True,
-                        tasks=[ModelTask.TextGeneration],
-                        backends=[Backend.MlxMetal],
+                        model_card=ModelCard(
+                            model_id=ModelId("llama-3.2-1b"),
+                            n_layers=16,
+                            storage_size=Memory.from_bytes(678948),
+                            hidden_size=7168,
+                            supports_tensor=True,
+                            tasks=[ModelTask.TextGeneration],
+                            backends=[Backend.MlxMetal],
+                        ),
+                        device_rank=0,
+                        world_size=1,
                     ),
-                    device_rank=0,
-                    world_size=1,
                 )
-            },
-            node_to_runner={node_id: runner_id},
+            ],
+            primary_output_node=0,
         )
         assert created_instance.shard_assignments == expected_shard_assignments
         # For single-node, hosts_by_node should have one entry with self-binding

src/exo/master/tests/test_placement.py

@@ -49,16 +49,36 @@ from exo.shared.types.worker.instances import (
     MlxJacclInstance,
     MlxRingInstance,
 )
-from exo.shared.types.worker.runners import ShardAssignments
+from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardWithId
 from exo.shared.types.worker.shards import PipelineShardMetadata, Sharding
 
 
+class MockShard:
+    def is_primary_output(self) -> bool:
+        return True
+
+
 @pytest.fixture
-def instance() -> Instance:
+def instance(model_card: ModelCard) -> Instance:
     return MlxRingInstance(
         instance_id=InstanceId(),
         shard_assignments=ShardAssignments(
-            model_id=ModelId("test-model"), runner_to_shard={}, node_to_runner={}
+            model_id=ModelId("test-model"),
+            shards=[
+                ShardWithId(
+                    NodeId(),
+                    RunnerId(),
+                    PipelineShardMetadata(
+                        model_card=model_card,
+                        device_rank=0,
+                        world_size=1,
+                        start_layer=0,
+                        end_layer=model_card.n_layers,
+                        n_layers=model_card.n_layers,
+                    ),
+                )
+            ],
+            primary_output_node=0,
         ),
         hosts_by_node={},
         ephemeral_port=50000,
@@ -123,6 +143,11 @@ def test_get_instance_placements_create_instance(
     node_id_a = NodeId()
     node_id_b = NodeId()
     node_id_c = NodeId()
+    node_to_layers = {
+        node_id_a: expected_layers[0],
+        node_id_b: expected_layers[1],
+        node_id_c: expected_layers[2],
+    }
 
     # fully connected (directed) between the 3 nodes
     conn_a_b = Connection(
@@ -175,22 +200,11 @@ def test_get_instance_placements_create_instance(
     instance = placements[instance_id]
     assert instance.shard_assignments.model_id == model_card.model_id
 
-    runner_id_a = instance.shard_assignments.node_to_runner[node_id_a]
-    runner_id_b = instance.shard_assignments.node_to_runner[node_id_b]
-    runner_id_c = instance.shard_assignments.node_to_runner[node_id_c]
+    for nid, _, shard in (shards := instance.shard_assignments.shards):
+        assert shard.end_layer - shard.start_layer == node_to_layers[nid]
 
-    shard_a = instance.shard_assignments.runner_to_shard[runner_id_a]
-    shard_b = instance.shard_assignments.runner_to_shard[runner_id_b]
-    shard_c = instance.shard_assignments.runner_to_shard[runner_id_c]
-
-    assert shard_a.end_layer - shard_a.start_layer == expected_layers[0]
-    assert shard_b.end_layer - shard_b.start_layer == expected_layers[1]
-    assert shard_c.end_layer - shard_c.start_layer == expected_layers[2]
-
-    shards = [shard_a, shard_b, shard_c]
-    shards_sorted = sorted(shards, key=lambda s: s.start_layer)
-    assert shards_sorted[0].start_layer == 0
-    assert shards_sorted[-1].end_layer == total_layers
+    assert shards[0].shard.start_layer == 0
+    assert shards[-1].shard.end_layer == total_layers
 
 
 def test_get_instance_placements_one_node_exact_fit() -> None:
@@ -218,9 +232,7 @@ def test_get_instance_placements_one_node_exact_fit() -> None:
     instance_id = list(placements.keys())[0]
     instance = placements[instance_id]
     assert instance.shard_assignments.model_id == "test-model"
-    assert len(instance.shard_assignments.node_to_runner) == 1
-    assert len(instance.shard_assignments.runner_to_shard) == 1
-    assert len(instance.shard_assignments.runner_to_shard) == 1
+    assert len(instance.shard_assignments.shards) == 1
 
 
 def test_get_instance_placements_one_node_fits_with_extra_memory() -> None:
@@ -248,9 +260,7 @@ def test_get_instance_placements_one_node_fits_with_extra_memory() -> None:
     instance_id = list(placements.keys())[0]
     instance = placements[instance_id]
     assert instance.shard_assignments.model_id == "test-model"
-    assert len(instance.shard_assignments.node_to_runner) == 1
-    assert len(instance.shard_assignments.runner_to_shard) == 1
-    assert len(instance.shard_assignments.runner_to_shard) == 1
+    assert len(instance.shard_assignments.shards) == 1
 
 
 def test_get_instance_placements_one_node_not_fit() -> None:
@@ -381,7 +391,7 @@ def test_placement_selects_leaf_nodes(
     assert len(placements) == 1
     instance = list(placements.values())[0]
 
-    assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
+    assigned_nodes = set(map(lambda it: it.node_id, instance.shard_assignments.shards))
     assert assigned_nodes == set((node_id_a, node_id_b)) or assigned_nodes == set(
         (
             node_id_c,
@@ -498,8 +508,8 @@ def test_tensor_rdma_backend_connectivity_matrix(
     for i in range(3):
         assert matrix[i][i] is None
 
-    assigned_nodes = list(instance.shard_assignments.node_to_runner.keys())
-    node_to_idx = {node_id: idx for idx, node_id in enumerate(assigned_nodes)}
+    assigned_nodes = list(instance.shard_assignments.shards)
+    node_to_idx = {node_id: idx for idx, (node_id, _, _) in enumerate(assigned_nodes)}
 
     idx_a = node_to_idx[node_a]
     idx_b = node_to_idx[node_b]
@@ -511,7 +521,7 @@ def test_tensor_rdma_backend_connectivity_matrix(
 
     # Verify coordinators are set for all nodes
     assert len(instance.jaccl_coordinators) == 3
-    for node_id in assigned_nodes:
+    for node_id, _, _ in assigned_nodes:
         assert node_id in instance.jaccl_coordinators
         coordinator = instance.jaccl_coordinators[node_id]
         assert ":" in coordinator
@@ -825,7 +835,7 @@ def test_placement_prefers_cycle_with_downloaded_model(
 
     assert len(placements) == 1
     instance = list(placements.values())[0]
-    assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
+    assigned_nodes = set(map(lambda it: it.node_id, instance.shard_assignments.shards))
     assert assigned_nodes == {node_b}
 
 
@@ -903,7 +913,7 @@ def test_placement_prefers_cycle_with_higher_download_progress(
 
     assert len(placements) == 1
     instance = list(placements.values())[0]
-    assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
+    assigned_nodes = set(map(lambda it: it.node_id, instance.shard_assignments.shards))
     assert assigned_nodes == {node_b}
 
 
@@ -957,7 +967,7 @@ def test_placement_does_not_prefer_cycle_with_failed_download(
 
     assert len(placements) == 1
     instance = list(placements.values())[0]
-    assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
+    assigned_nodes = set(map(lambda it: it.node_id, instance.shard_assignments.shards))
     # node_a should win on RAM tiebreaker since failed download scores 0.0
     assert assigned_nodes == {node_a}
 

src/exo/master/tests/test_placement_utils.py

@@ -204,6 +204,11 @@ def test_get_shard_assignments(
     node_a_id = NodeId()
     node_b_id = NodeId()
     node_c_id = NodeId()
+    layers_by_node = {
+        node_a_id: expected_layers[0],
+        node_b_id: expected_layers[1],
+        node_c_id: expected_layers[2],
+    }
 
     # create connections (A -> B -> C -> A forms a 3-cycle, plus B -> A also exists)
     connection1 = Connection(
@@ -258,25 +263,8 @@ def test_get_shard_assignments(
     )
 
     # assert
-    runner_id_a = shard_assignments.node_to_runner[node_a_id]
-    runner_id_b = shard_assignments.node_to_runner[node_b_id]
-    runner_id_c = shard_assignments.node_to_runner[node_c_id]
-
-    assert (
-        shard_assignments.runner_to_shard[runner_id_a].end_layer
-        - shard_assignments.runner_to_shard[runner_id_a].start_layer
-        == expected_layers[0]
-    )
-    assert (
-        shard_assignments.runner_to_shard[runner_id_b].end_layer
-        - shard_assignments.runner_to_shard[runner_id_b].start_layer
-        == expected_layers[1]
-    )
-    assert (
-        shard_assignments.runner_to_shard[runner_id_c].end_layer
-        - shard_assignments.runner_to_shard[runner_id_c].start_layer
-        == expected_layers[2]
-    )
+    for nid, _, shard in shard_assignments.shards:
+        assert shard.end_layer - shard.start_layer == layers_by_node[nid]
 
 
 def test_get_mlx_jaccl_coordinators():
@@ -543,11 +531,11 @@ class TestCfgParallelPlacement:
             model_card, cycle, node_memory
         )
 
-        shards = list(assignments.runner_to_shard.values())
+        shards = list(assignments.shards)
         assert len(shards) == 2
 
         # CFG models should get CfgShardMetadata
-        for shard in shards:
+        for _, _, shard in shards:
             assert isinstance(shard, CfgShardMetadata)
             # Both nodes should have all layers (no pipeline split)
             assert shard.start_layer == 0
@@ -558,7 +546,7 @@ class TestCfgParallelPlacement:
             assert shard.pipeline_rank == 0
 
         cfg_ranks = sorted(
-            s.cfg_rank for s in shards if isinstance(s, CfgShardMetadata)
+            s.shard.cfg_rank for s in shards if isinstance(s.shard, CfgShardMetadata)
         )
         assert cfg_ranks == [0, 1]
 
@@ -587,11 +575,11 @@ class TestCfgParallelPlacement:
             model_card, cycle, node_memory
         )
 
-        shards = list(assignments.runner_to_shard.values())
+        shards = assignments.shards
         assert len(shards) == 4
 
         # CFG models should get CfgShardMetadata
-        for shard in shards:
+        for _, _, shard in shards:
             assert isinstance(shard, CfgShardMetadata)
             assert shard.cfg_world_size == 2
             assert shard.pipeline_world_size == 2
@@ -599,10 +587,14 @@ class TestCfgParallelPlacement:
 
         # Check we have 2 nodes in each CFG group
         cfg_0_shards = [
-            s for s in shards if isinstance(s, CfgShardMetadata) and s.cfg_rank == 0
+            s.shard
+            for s in shards
+            if isinstance(s.shard, CfgShardMetadata) and s.shard.cfg_rank == 0
         ]
         cfg_1_shards = [
-            s for s in shards if isinstance(s, CfgShardMetadata) and s.cfg_rank == 1
+            s.shard
+            for s in shards
+            if isinstance(s.shard, CfgShardMetadata) and s.shard.cfg_rank == 1
         ]
         assert len(cfg_0_shards) == 2
         assert len(cfg_1_shards) == 2
@@ -637,11 +629,11 @@ class TestCfgParallelPlacement:
             model_card, cycle, node_memory
         )
 
-        shards = list(assignments.runner_to_shard.values())
+        shards = list(assignments.shards)
         assert len(shards) == 3
 
         # Odd node count with CFG model falls back to PipelineShardMetadata (sequential CFG)
-        for shard in shards:
+        for _, _, shard in shards:
             assert isinstance(shard, PipelineShardMetadata)
 
     def test_two_nodes_non_cfg_model_uses_pipeline(self):
@@ -673,18 +665,18 @@ class TestCfgParallelPlacement:
             model_card, cycle, node_memory
         )
 
-        shards = list(assignments.runner_to_shard.values())
+        shards = list(assignments.shards)
         assert len(shards) == 2
 
         # Non-CFG models should get PipelineShardMetadata
-        for shard in shards:
+        for _, _, shard in shards:
             assert isinstance(shard, PipelineShardMetadata)
 
         # Should have actual layer sharding (pipeline)
         layer_ranges = sorted(
-            (s.start_layer, s.end_layer)
+            (s.shard.start_layer, s.shard.end_layer)
             for s in shards
-            if isinstance(s, PipelineShardMetadata)
+            if isinstance(s.shard, PipelineShardMetadata)
         )
         # First shard starts at 0, last shard ends at 57
         assert layer_ranges[0][0] == 0

src/exo/routing/connection_message.py

@@ -1,4 +1,4 @@
-from exo_rs import FromSwarm
+from exo_rs import PyFromSwarm
 
 from exo.utils.pydantic_ext import FrozenModel
 
@@ -9,5 +9,5 @@ class ConnectionMessage(FrozenModel):
     connected: bool
 
     @classmethod
-    def from_update(cls, update: FromSwarm.Connection) -> "ConnectionMessage":
+    def from_update(cls, update: PyFromSwarm.Connection) -> "ConnectionMessage":
         return cls(connected=update.connected)

src/exo/routing/router.py

@@ -12,8 +12,8 @@ from anyio import (
     sleep_forever,
 )
 from exo_rs import (
-    FromSwarm,
     NetworkingHandle,
+    PyFromSwarm,
 )
 from loguru import logger
 
@@ -101,14 +101,11 @@ class Router:
     def create(
         cls,
         identity: str,
-        namespace: str,
         listen_port: int,
         discovery_service_port: int,
     ) -> "Router":
         return cls(
-            handle=NetworkingHandle.new(
-                identity, namespace, listen_port, discovery_service_port
-            )
+            handle=NetworkingHandle.new(identity, listen_port, discovery_service_port)
         )
 
     def __init__(self, handle: NetworkingHandle):
@@ -189,7 +186,7 @@ class Router:
                 from_swarm = await self._net.recv()
                 logger.debug(from_swarm)
                 match from_swarm:
-                    case FromSwarm.Message(topic, data):
+                    case PyFromSwarm.Message(topic, data):
                         logger.trace(f"Received message on {topic} with payload {data}")
                         if topic not in self.topic_routers:
                             logger.warning(
@@ -198,7 +195,7 @@ class Router:
                             continue
                         router = self.topic_routers[topic]
                         await router.publish_bytes(data)
-                    case FromSwarm.Connection():
+                    case PyFromSwarm.Connection():
                         message = ConnectionMessage.from_update(from_swarm)
                         logger.trace(
                             f"Received message on connection_messages with payload {message}"
@@ -237,7 +234,7 @@ def get_node_zid(
     Obtain the :class:`PeerId` by from it.
     """
     # TODO(evan): bring back node id persistence once we figure out how to deal with duplicates
-    return NodeId(os.urandom(16).hex().lstrip("0"))
+    return NodeId(os.urandom(16).hex())
 
     """
     def lock_path(path: str | bytes | PathLike[str] | PathLike[bytes]) -> Path:

src/exo/shared/apply.py

@@ -4,19 +4,14 @@ from datetime import datetime
 
 from loguru import logger
 
-from exo.shared.models.model_cards import ModelCard
-from exo.shared.types.common import ModelId, NodeId
+from exo.shared.types.common import NodeId
 from exo.shared.types.events import (
     ChunkGenerated,
-    CustomModelCardAdded,
-    CustomModelCardDeleted,
     Event,
     IndexedEvent,
     InputChunkReceived,
     InstanceCreated,
     InstanceDeleted,
-    InstanceLinkCreated,
-    InstanceLinkDeleted,
     NodeDownloadProgress,
     NodeGatheredInfo,
     NodeTimedOut,
@@ -32,7 +27,6 @@ from exo.shared.types.events import (
     TracesCollected,
     TracesMerged,
 )
-from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.profiling import (
     NodeIdentity,
     NodeNetworkInfo,
@@ -42,7 +36,6 @@ from exo.shared.types.profiling import (
 )
 from exo.shared.types.state import State
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
-from exo.shared.types.topology import Connection, RDMAConnection
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
 from exo.shared.types.worker.runners import (
@@ -67,18 +60,6 @@ from exo.utils.info_gatherer.info_gatherer import (
 )
 
 
-def _is_rdma_ctl_enabled(
-    node_id: NodeId, node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus]
-) -> bool:
-    """A node is RDMA-capable only if rdma_ctl status has been observed as enabled.
-
-    Missing entries default to ``False`` — if we have not yet observed (or the node
-    cannot run) ``rdma_ctl``, it must not participate in an RDMA-backed instance.
-    """
-    status = node_rdma_ctl.get(node_id)
-    return status is not None and status.enabled
-
-
 def event_apply(event: Event, state: State) -> State:
     """Apply an event to state."""
     match event:
@@ -91,10 +72,6 @@ def event_apply(event: Event, state: State) -> State:
             | TracesMerged()
         ):  # Pass-through events that don't modify state
             return state
-        case CustomModelCardAdded():
-            return apply_custom_model_card_added(event, state)
-        case CustomModelCardDeleted():
-            return apply_custom_model_card_deleted(event, state)
         case InstanceCreated():
             return apply_instance_created(event, state)
         case InstanceDeleted():
@@ -119,10 +96,6 @@ def event_apply(event: Event, state: State) -> State:
             return apply_topology_edge_created(event, state)
         case TopologyEdgeDeleted():
             return apply_topology_edge_deleted(event, state)
-        case InstanceLinkCreated():
-            return apply_instance_link_created(event, state)
-        case InstanceLinkDeleted():
-            return apply_instance_link_deleted(event, state)
 
 
 def apply(state: State, event: IndexedEvent) -> State:
@@ -222,38 +195,7 @@ def apply_instance_deleted(event: InstanceDeleted, state: State) -> State:
     new_instances: Mapping[InstanceId, Instance] = {
         iid: inst for iid, inst in state.instances.items() if iid != event.instance_id
     }
-    new_links: dict[InstanceLinkId, InstanceLink] = {}
-    for link_id, link in state.instance_links.items():
-        prefill = [i for i in link.prefill_instances if i != event.instance_id]
-        decode = [i for i in link.decode_instances if i != event.instance_id]
-        if not prefill or not decode:
-            continue
-        if prefill == list(link.prefill_instances) and decode == list(
-            link.decode_instances
-        ):
-            new_links[link_id] = link
-        else:
-            new_links[link_id] = link.model_copy(
-                update={"prefill_instances": prefill, "decode_instances": decode}
-            )
-    return state.model_copy(
-        update={"instances": new_instances, "instance_links": new_links}
-    )
-
-
-def apply_instance_link_created(event: InstanceLinkCreated, state: State) -> State:
-    new_links: Mapping[InstanceLinkId, InstanceLink] = {
-        **state.instance_links,
-        event.link.link_id: event.link,
-    }
-    return state.model_copy(update={"instance_links": new_links})
-
-
-def apply_instance_link_deleted(event: InstanceLinkDeleted, state: State) -> State:
-    new_links: Mapping[InstanceLinkId, InstanceLink] = {
-        lid: link for lid, link in state.instance_links.items() if lid != event.link_id
-    }
-    return state.model_copy(update={"instance_links": new_links})
+    return state.model_copy(update={"instances": new_instances})
 
 
 def apply_runner_status_updated(event: RunnerStatusUpdated, state: State) -> State:
@@ -408,59 +350,26 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
                 event.node_id: NodeThunderboltInfo(interfaces=info.idents),
             }
         case MacThunderboltConnections():
-            conn_map = {
-                tb_ident.domain_uuid: (nid, tb_ident.rdma_interface)
-                for nid in state.node_thunderbolt
-                for tb_ident in state.node_thunderbolt[nid].interfaces
+            update["node_thunderbolt_connections"] = {
+                **state.node_thunderbolt_connections,
+                event.node_id: info,
             }
-            source_is_rdma_enabled = _is_rdma_ctl_enabled(
-                event.node_id, state.node_rdma_ctl
-            )
-            as_rdma_conns = [
-                Connection(
-                    source=event.node_id,
-                    sink=conn_map[tb_conn.sink_uuid][0],
-                    edge=RDMAConnection(
-                        source_rdma_iface=conn_map[tb_conn.source_uuid][1],
-                        sink_rdma_iface=conn_map[tb_conn.sink_uuid][1],
-                    ),
-                )
-                for tb_conn in info.conns
-                if tb_conn.source_uuid in conn_map
-                if tb_conn.sink_uuid in conn_map
-                if source_is_rdma_enabled
-                and _is_rdma_ctl_enabled(
-                    conn_map[tb_conn.sink_uuid][0], state.node_rdma_ctl
-                )
-            ]
-            topology.replace_all_out_rdma_connections(event.node_id, as_rdma_conns)
         case ThunderboltBridgeInfo():
             new_tb_bridge: dict[NodeId, ThunderboltBridgeStatus] = {
                 **state.node_thunderbolt_bridge,
                 event.node_id: info.status,
             }
             update["node_thunderbolt_bridge"] = new_tb_bridge
-            # Only recompute cycles if the enabled status changed
-            old_status = state.node_thunderbolt_bridge.get(event.node_id)
-            old_enabled = old_status.enabled if old_status else False
-            new_enabled = info.status.enabled
-            if old_enabled != new_enabled:
-                update["thunderbolt_bridge_cycles"] = (
-                    topology.get_thunderbolt_bridge_cycles(
-                        new_tb_bridge, state.node_network
-                    )
+            update["thunderbolt_bridge_cycles"] = (
+                topology.get_thunderbolt_bridge_cycles(
+                    new_tb_bridge, state.node_network
                 )
+            )
         case RdmaCtlStatus():
             update["node_rdma_ctl"] = {
                 **state.node_rdma_ctl,
                 event.node_id: NodeRdmaCtlStatus(enabled=info.enabled),
             }
-            # If RDMA just got disabled on this node, drop any RDMA edges touching it
-            # so placement / topology consumers cannot pick a disabled node for an
-            # RDMA-backed instance. (Edges will repopulate on the next
-            # MacThunderboltConnections poll once both endpoints are enabled again.)
-            if not info.enabled:
-                topology.remove_all_rdma_connections_touching(event.node_id)
         case NodeBackends():
             update["node_backends"] = {
                 **state.node_backends,
@@ -471,32 +380,45 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
 
 
 def apply_topology_edge_created(event: TopologyEdgeCreated, state: State) -> State:
-    topology = copy.deepcopy(state.topology)
-    topology.add_connection(event.conn)
-    return state.model_copy(update={"topology": topology})
+    source_connections = state.node_socket_connections.get(event.conn.source, {})
+    sink_connections = source_connections.get(event.conn.sink, [])
+
+    update = {
+        "node_socket_connections": {
+            **state.node_socket_connections,
+            event.conn.source: {
+                **source_connections,
+                event.conn.sink: sink_connections
+                if event.conn.edge in sink_connections
+                else [*sink_connections, event.conn.edge],
+            },
+        }
+    }
+    return state.model_copy(update=update)
 
 
 def apply_topology_edge_deleted(event: TopologyEdgeDeleted, state: State) -> State:
-    topology = copy.deepcopy(state.topology)
-    topology.remove_connection(event.conn)
-    # TODO: Clean up removing the reverse connection
-    return state.model_copy(update={"topology": topology})
-
-
-def apply_custom_model_card_added(event: CustomModelCardAdded, state: State) -> State:
-    new_cards: Mapping[ModelId, ModelCard] = {
-        **state.custom_model_cards,
-        event.model_card.model_id: event.model_card,
+    inner_update = {
+        sink: final_edges
+        for sink, edges in state.node_socket_connections.get(
+            event.conn.source, {}
+        ).items()
+        if (
+            final_edges := [
+                edge
+                for edge in edges
+                if (edge != event.conn.edge or sink != event.conn.sink)
+            ]
+        )
     }
-    return state.model_copy(update={"custom_model_cards": new_cards})
-
-
-def apply_custom_model_card_deleted(
-    event: CustomModelCardDeleted, state: State
-) -> State:
-    new_cards: Mapping[ModelId, ModelCard] = {
-        model_id: card
-        for model_id, card in state.custom_model_cards.items()
-        if model_id != event.model_id
+    update = {
+        "node_socket_connections": {
+            source: maps
+            for source, maps in {
+                **state.node_socket_connections,
+                event.conn.source: inner_update,
+            }.items()
+            if maps
+        }
     }
-    return state.model_copy(update={"custom_model_cards": new_cards})
+    return state.model_copy(update=update)

src/exo/shared/models/model_cards.py

@@ -55,7 +55,7 @@ class _CardCache:
         except OSError as e:
             logger.warning(f"failed to save custom model card ({e.strerror})")
 
-    async def pop(self, model_id: ModelId) -> "ModelCard | None":
+    async def delete(self, model_id: ModelId) -> "ModelCard | None":
         """Delete a user-added custom model card. Returns True if deleted."""
         card_path = _custom_cards_dir / (ModelId(model_id).normalize() + ".toml")
         try:

src/exo/shared/tests/test_apply/test_apply_custom_model_cards.py

@@ -1,46 +0,0 @@
-from exo.shared.apply import apply
-from exo.shared.models.model_cards import ModelCard, ModelTask
-from exo.shared.types.backends import Backend
-from exo.shared.types.common import ModelId
-from exo.shared.types.events import (
-    CustomModelCardAdded,
-    CustomModelCardDeleted,
-    IndexedEvent,
-)
-from exo.shared.types.memory import Memory
-from exo.shared.types.state import State
-
-
-def _model_card(model_id: ModelId) -> ModelCard:
-    return ModelCard(
-        model_id=model_id,
-        n_layers=1,
-        storage_size=Memory.from_bytes(1),
-        hidden_size=1,
-        supports_tensor=True,
-        tasks=[ModelTask.TextGeneration],
-        backends=[Backend.MlxMetal],
-    )
-
-
-def test_custom_model_card_added_is_reduced_into_state() -> None:
-    card = _model_card(ModelId("custom/model"))
-
-    state = apply(
-        State(),
-        IndexedEvent(idx=0, event=CustomModelCardAdded(model_card=card)),
-    )
-
-    assert state.custom_model_cards == {card.model_id: card}
-
-
-def test_custom_model_card_deleted_removes_card_from_state() -> None:
-    card = _model_card(ModelId("custom/model"))
-    state = State(custom_model_cards={card.model_id: card}, last_event_applied_idx=0)
-
-    state = apply(
-        state,
-        IndexedEvent(idx=1, event=CustomModelCardDeleted(model_id=card.model_id)),
-    )
-
-    assert state.custom_model_cards == {}

src/exo/shared/tests/test_apply/test_apply_instance_link.py

@@ -1,72 +0,0 @@
-from exo.shared.apply import (
-    apply_instance_deleted,
-    apply_instance_link_created,
-    apply_instance_link_deleted,
-)
-from exo.shared.types.events import (
-    InstanceDeleted,
-    InstanceLinkCreated,
-    InstanceLinkDeleted,
-)
-from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
-from exo.shared.types.state import State
-from exo.shared.types.worker.instances import InstanceId
-
-
-def _link(
-    prefill: list[InstanceId],
-    decode: list[InstanceId],
-    link_id: InstanceLinkId | None = None,
-) -> InstanceLink:
-    return InstanceLink(
-        link_id=link_id or InstanceLinkId(),
-        prefill_instances=prefill,
-        decode_instances=decode,
-    )
-
-
-def test_create_link() -> None:
-    state = State()
-    link = _link([InstanceId("a")], [InstanceId("b")])
-    new_state = apply_instance_link_created(InstanceLinkCreated(link=link), state)
-    assert new_state.instance_links == {link.link_id: link}
-
-
-def test_update_replaces_existing_link() -> None:
-    a, b, c = InstanceId("a"), InstanceId("b"), InstanceId("c")
-    link = _link([a], [b])
-    state = State(instance_links={link.link_id: link})
-
-    updated = link.model_copy(update={"decode_instances": [b, c]})
-    new_state = apply_instance_link_created(InstanceLinkCreated(link=updated), state)
-    assert set(new_state.instance_links[link.link_id].decode_instances) == {b, c}
-
-
-def test_delete_link() -> None:
-    link = _link([InstanceId("a")], [InstanceId("b")])
-    state = State(instance_links={link.link_id: link})
-
-    new_state = apply_instance_link_deleted(
-        InstanceLinkDeleted(link_id=link.link_id), state
-    )
-    assert new_state.instance_links == {}
-
-
-def test_instance_deleted_strips_from_links() -> None:
-    a, b, c = InstanceId("a"), InstanceId("b"), InstanceId("c")
-    link = _link([a, c], [b])
-    state = State(instance_links={link.link_id: link})
-
-    new_state = apply_instance_deleted(InstanceDeleted(instance_id=a), state)
-    remaining = new_state.instance_links[link.link_id]
-    assert remaining.prefill_instances == [c]
-    assert remaining.decode_instances == [b]
-
-
-def test_instance_deleted_drops_link_when_role_empties() -> None:
-    a, b = InstanceId("a"), InstanceId("b")
-    link = _link([a], [b])
-    state = State(instance_links={link.link_id: link})
-
-    new_state = apply_instance_deleted(InstanceDeleted(instance_id=a), state)
-    assert link.link_id not in new_state.instance_links

src/exo/shared/tests/test_apply/test_apply_rdma_gating.py

@@ -217,7 +217,7 @@ def test_topology_remove_all_rdma_connections_touching_keeps_socket_edges():
         )
     )
     socket_edge = SocketConnection(
-        sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000")
+        sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000"),
     )
     topology.add_connection(Connection(source=node_a, sink=node_b, edge=socket_edge))
 

src/exo/shared/tests/test_state_serialization.py

@@ -1,35 +0,0 @@
-from exo.shared.types.common import NodeId
-from exo.shared.types.multiaddr import Multiaddr
-from exo.shared.types.state import State
-from exo.shared.types.topology import Connection, SocketConnection
-
-
-def test_state_serialization_roundtrip() -> None:
-    """Verify that State → JSON → State round-trip preserves topology."""
-
-    # --- build a simple state ------------------------------------------------
-    node_a = NodeId("node-a")
-    node_b = NodeId("node-b")
-
-    connection = Connection(
-        source=node_a,
-        sink=node_b,
-        edge=SocketConnection(
-            sink_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/10001"),
-        ),
-    )
-
-    state = State()
-    state.topology.add_connection(connection)
-
-    json_repr = state.model_dump_json()
-    restored_state = State.model_validate_json(json_repr)
-
-    assert (
-        state.topology.to_snapshot().nodes
-        == restored_state.topology.to_snapshot().nodes
-    )
-    assert set(state.topology.to_snapshot().connections) == set(
-        restored_state.topology.to_snapshot().connections
-    )
-    assert restored_state.model_dump_json() == json_repr

src/exo/shared/topology.py

@@ -1,10 +1,8 @@
-import contextlib
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field
 from typing import Iterable
 
 import rustworkx as rx
-from pydantic import BaseModel, ConfigDict
 
 from exo.shared.types.common import NodeId
 from exo.shared.types.profiling import (
@@ -20,15 +18,6 @@ from exo.shared.types.topology import (
 )
 
 
-class TopologySnapshot(BaseModel):
-    nodes: Sequence[NodeId]
-    connections: Mapping[
-        NodeId, Mapping[NodeId, Sequence[SocketConnection | RDMAConnection]]
-    ]
-
-    model_config = ConfigDict(frozen=True, extra="forbid")
-
-
 @dataclass
 class Topology:
     _graph: rx.PyDiGraph[NodeId, SocketConnection | RDMAConnection] = field(
@@ -36,28 +25,6 @@ class Topology:
     )
     _vertex_indices: dict[NodeId, int] = field(init=False, default_factory=dict)
 
-    def to_snapshot(self) -> TopologySnapshot:
-        return TopologySnapshot(
-            nodes=list(self.list_nodes()), connections=self.map_connections()
-        )
-
-    @classmethod
-    def from_snapshot(cls, snapshot: TopologySnapshot) -> "Topology":
-        topology = cls()
-
-        for node_id in snapshot.nodes:
-            with contextlib.suppress(ValueError):
-                topology.add_node(node_id)
-
-        for source in snapshot.connections:
-            for sink in snapshot.connections[source]:
-                for edge in snapshot.connections[source][sink]:
-                    topology.add_connection(
-                        Connection(source=source, sink=sink, edge=edge)
-                    )
-
-        return topology
-
     def add_node(self, node_id: NodeId) -> None:
         if node_id in self._vertex_indices:
             return

src/exo/shared/types/commands.py

@@ -7,7 +7,6 @@ from exo.api.types import (
 from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.chunks import InputImageChunk
 from exo.shared.types.common import CommandId, NodeId, SystemId
-from exo.shared.types.instance_link import InstanceLinkId
 from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding, ShardMetadata
@@ -82,24 +81,6 @@ class CancelDownload(BaseCommand):
     model_id: ModelId
 
 
-class AddCustomModelCard(BaseCommand):
-    model_card: ModelCard
-
-
-class DeleteCustomModelCard(BaseCommand):
-    model_id: ModelId
-
-
-class SetInstanceLink(BaseCommand):
-    link_id: InstanceLinkId
-    prefill_instances: list[InstanceId]
-    decode_instances: list[InstanceId]
-
-
-class DeleteInstanceLink(BaseCommand):
-    link_id: InstanceLinkId
-
-
 DownloadCommand = StartDownload | DeleteDownload | CancelDownload
 
 
@@ -115,10 +96,6 @@ Command = (
     | TaskCancelled
     | TaskFinished
     | SendInputChunk
-    | AddCustomModelCard
-    | DeleteCustomModelCard
-    | SetInstanceLink
-    | DeleteInstanceLink
 )
 
 

src/exo/shared/types/events.py

@@ -3,11 +3,9 @@ from typing import final
 
 from pydantic import Field
 
-from exo.shared.models.model_cards import ModelCard
 from exo.shared.topology import Connection
 from exo.shared.types.chunks import Chunk, InputImageChunk
-from exo.shared.types.common import CommandId, Id, ModelId, NodeId, SessionId, SystemId
-from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
+from exo.shared.types.common import CommandId, Id, NodeId, SessionId, SystemId
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
@@ -108,14 +106,6 @@ class TopologyEdgeDeleted(BaseEvent):
     conn: Connection
 
 
-class CustomModelCardAdded(BaseEvent):
-    model_card: ModelCard
-
-
-class CustomModelCardDeleted(BaseEvent):
-    model_id: ModelId
-
-
 @final
 class TraceEventData(FrozenModel):
     name: str
@@ -138,14 +128,6 @@ class TracesMerged(BaseEvent):
     traces: list[TraceEventData]
 
 
-class InstanceLinkCreated(BaseEvent):
-    link: InstanceLink
-
-
-class InstanceLinkDeleted(BaseEvent):
-    link_id: InstanceLinkId
-
-
 Event = (
     TestEvent
     | TaskCreated
@@ -165,10 +147,6 @@ Event = (
     | TopologyEdgeDeleted
     | TracesCollected
     | TracesMerged
-    | CustomModelCardAdded
-    | CustomModelCardDeleted
-    | InstanceLinkCreated
-    | InstanceLinkDeleted
 )
 
 

src/exo/shared/types/profiling.py

@@ -16,6 +16,10 @@ class MemoryUsage(FrozenModel):
     swap_total: Memory
     swap_available: Memory
 
+    @classmethod
+    def tag(cls) -> str:
+        return cls.__name__
+
     @classmethod
     def from_bytes(
         cls, *, ram_total: int, ram_available: int, swap_total: int, swap_available: int

src/exo/shared/types/state.py

@@ -1,15 +1,15 @@
 from collections.abc import Mapping, Sequence
 from datetime import datetime
-from typing import Any, cast
+from typing import Any
 
-from pydantic import ConfigDict, Field, field_serializer, field_validator
+from exo_rs import LVAggregator
+from pydantic import ConfigDict, Field, model_serializer
 from pydantic.alias_generators import to_camel
+from pydantic_core.core_schema import SerializerFunctionWrapHandler
 
-from exo.shared.models.model_cards import ModelCard
-from exo.shared.topology import Topology, TopologySnapshot
+from exo.shared.topology import Topology
 from exo.shared.types.backends import Backend
-from exo.shared.types.common import ModelId, NodeId
-from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
+from exo.shared.types.common import NodeId
 from exo.shared.types.profiling import (
     DiskUsage,
     MemoryUsage,
@@ -21,9 +21,15 @@ from exo.shared.types.profiling import (
     ThunderboltBridgeStatus,
 )
 from exo.shared.types.tasks import Task, TaskId
+from exo.shared.types.topology import (
+    Connection,
+    RDMAConnection,
+    SocketConnection,
+)
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
 from exo.shared.types.worker.runners import RunnerId, RunnerStatus
+from exo.utils.info_gatherer.info_gatherer import MacThunderboltConnections
 from exo.utils.pydantic_ext import FrozenModel
 
 
@@ -39,7 +45,6 @@ class State(FrozenModel):
         alias_generator=to_camel,
         validate_by_name=True,
         extra="forbid",
-        # I want to reenable this ASAP, but it's causing an issue with TaskStatus
         strict=True,
         arbitrary_types_allowed=True,
     )
@@ -48,7 +53,6 @@ class State(FrozenModel):
     downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {}
     tasks: Mapping[TaskId, Task] = {}
     last_seen: Mapping[NodeId, datetime] = {}
-    topology: Topology = Field(default_factory=Topology)
     last_event_applied_idx: int = Field(default=-1, ge=-1)
 
     # Granular node state mappings (update independently at different frequencies)
@@ -61,34 +65,94 @@ class State(FrozenModel):
     node_thunderbolt_bridge: Mapping[NodeId, ThunderboltBridgeStatus] = {}
     node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus] = {}
     node_backends: Mapping[NodeId, list[Backend]] = {}
+    node_socket_connections: Mapping[
+        NodeId, Mapping[NodeId, Sequence[SocketConnection]]
+    ] = {}
+    node_thunderbolt_connections: Mapping[NodeId, MacThunderboltConnections] = {}
 
     # Detected cycles where all nodes have Thunderbolt bridge enabled (>2 nodes)
     thunderbolt_bridge_cycles: Sequence[Sequence[NodeId]] = []
 
-    instance_links: Mapping[InstanceLinkId, InstanceLink] = {}
     prefill_server_ports: Mapping[RunnerId, int] = {}
 
-    # User-added model cards. Workers can reconcile their on-disk custom card cache
-    custom_model_cards: Mapping[ModelId, ModelCard] = {}
+    @model_serializer(mode="wrap")
+    def _serialize(self, handler: SerializerFunctionWrapHandler) -> dict[str, Any]:
+        data = handler(self)  # pyright: ignore[reportAny]
+        data["topology"] = {
+            "nodes": list(self.node_identities.keys()),
+            "connections": self.topology.map_connections(),
+        }
+        return data  # pyright: ignore[reportAny]
 
-    @field_serializer("topology", mode="plain")
-    def _encode_topology(self, value: Topology) -> TopologySnapshot:
-        return value.to_snapshot()
+    @property
+    def topology(self) -> Topology:
+        topology = Topology()
+        thunderbolt_by_uuid = {
+            ident.domain_uuid: (node_id, ident.rdma_interface)
+            for node_id, info in self.node_thunderbolt.items()
+            for ident in info.interfaces
+        }
+        for node_id in self.node_identities:
+            topology.add_node(node_id)
 
-    @field_validator("topology", mode="before")
-    @classmethod
-    def _deserialize_topology(cls, value: object) -> Topology:  # noqa: D401 – Pydantic validator signature
-        """Convert an incoming *value* into a :class:`Topology` instance.
+        for source, data in self.node_socket_connections.items():
+            for sink, conns in data.items():
+                for conn in conns:
+                    topology.add_connection(
+                        Connection(source=source, sink=sink, edge=conn)
+                    )
 
-        Accepts either an already constructed :class:`Topology` or a mapping
-        representing :class:`~shared.topology.TopologySnapshot`.
-        """
+        for source, connections in self.node_thunderbolt_connections.items():
+            if not self.node_rdma_ctl.get(
+                source, NodeRdmaCtlStatus(enabled=False)
+            ).enabled:
+                continue
+            for connection in connections.conns:
+                if (
+                    source_iface := thunderbolt_by_uuid.get(connection.source_uuid)
+                ) is None or (
+                    sink_iface := thunderbolt_by_uuid.get(connection.sink_uuid)
+                ) is None:
+                    continue
+                if not self.node_rdma_ctl.get(
+                    sink_iface[0], NodeRdmaCtlStatus(enabled=False)
+                ).enabled:
+                    continue
+                assert source_iface[0] == source, "registered invalid source uuid"
+                topology.add_connection(
+                    Connection(
+                        source=source_iface[0],
+                        sink=sink_iface[0],
+                        edge=RDMAConnection(
+                            source_rdma_iface=source_iface[1],
+                            sink_rdma_iface=sink_iface[1],
+                        ),
+                    )
+                )
 
-        if isinstance(value, Topology):
-            return value
+        return topology
 
-        if isinstance(value, Mapping):  # likely a snapshot-dict coming from JSON
-            snapshot = TopologySnapshot(**cast(dict[str, Any], value))  # type: ignore[arg-type]
-            return Topology.from_snapshot(snapshot)
+    def with_aggregator(self, aggregator: LVAggregator) -> "State":
+        from datetime import datetime, timezone
 
-        raise TypeError("Invalid representation for Topology field in State")
+        from pydantic import TypeAdapter
+
+        from exo.shared.apply import event_apply
+        from exo.shared.types.events import NodeGatheredInfo
+        from exo.utils.info_gatherer.info_gatherer import GatheredInfo
+
+        state = self.model_copy()
+        for key, value in aggregator.dump().items():
+            try:
+                data = TypeAdapter[GatheredInfo](GatheredInfo).validate_json(value)
+                node_id = NodeId(key.split("/")[0])
+                event = NodeGatheredInfo(
+                    node_id=node_id, when=str(datetime.now(tz=timezone.utc)), info=data
+                )
+                state = event_apply(event, state)
+            except Exception as e:
+                print(
+                    f"\n{'=' * 10}key: {key} with exception {str(e)}\nvalue: {value}{'=' * 10}\n"
+                )
+
+        return state

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

@@ -1,3 +1,4 @@
+from collections.abc import Iterable
 from enum import Enum
 
 from pydantic import model_validator
@@ -22,7 +23,14 @@ class BaseInstance(TaggedModel):
     shard_assignments: ShardAssignments
 
     def shard(self, runner_id: RunnerId) -> ShardMetadata | None:
-        return self.shard_assignments.runner_to_shard.get(runner_id, None)
+        for _, rid, shard in self.shard_assignments.shards:
+            if rid == runner_id:
+                return shard
+
+    def runners_for(self, node_id: NodeId) -> Iterable[RunnerId]:
+        for nid, rid, _ in self.shard_assignments.shards:
+            if nid == node_id:
+                yield rid
 
 
 class MlxRingInstance(BaseInstance):
@@ -59,8 +67,9 @@ class BoundInstance(FrozenModel):
 
     @model_validator(mode="after")
     def validate_shard_exists(self) -> "BoundInstance":
-        assert (
-            self.bound_runner_id in self.instance.shard_assignments.runner_to_shard
+        assert any(
+            rid == self.bound_runner_id
+            for (_, rid, _) in self.instance.shard_assignments.shards
         ), (
             "Bound Instance must be constructed with a runner_id that is in the instances assigned shards"
         )

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

@@ -1,4 +1,5 @@
-from collections.abc import Mapping
+from collections.abc import Sequence
+from typing import NamedTuple
 
 from pydantic import model_validator
 
@@ -83,16 +84,26 @@ RunnerStatus = (
 )
 
 
+class ShardWithId(NamedTuple):
+    node_id: NodeId
+    runner_id: RunnerId
+    shard: ShardMetadata
+
+
 class ShardAssignments(FrozenModel):
     model_id: ModelId
-    runner_to_shard: Mapping[RunnerId, ShardMetadata]
-    node_to_runner: Mapping[NodeId, RunnerId]
+    shards: Sequence[ShardWithId]
+    # this node needs to be connected to the API node for the stream to be considered ready
+    # (this is a device rank)
+    primary_output_node: int
 
     @model_validator(mode="after")
     def validate_runners_exist(self) -> "ShardAssignments":
-        for runner_id in self.node_to_runner.values():
-            if runner_id not in self.runner_to_shard:
-                raise ValueError(
-                    f"Runner {runner_id} in node_to_runner does not exist in runner_to_shard"
-                )
+        for position, shard in enumerate(self.shards):
+            if shard.shard.device_rank != position:
+                raise ValueError("shard position does not correspond to device rank")
+
+        if not self.shards[self.primary_output_node].shard.is_primary_output():
+            raise ValueError("primary output node does not correspond to primary shard")
+
         return self

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

@@ -15,18 +15,14 @@ class Sharding(str, Enum):
 class BaseShardMetadata(TaggedModel):
     """
     Defines a specific shard of the model that is ready to be run on a device.
-    Replaces previous `Shard` object.
+    Layers are represented as a half-open interval [start_layer, end_layer),
+    where start_layer is inclusive and end_layer is exclusive.
     """
 
     model_card: ModelCard
     device_rank: int
     world_size: int
 
-    # Error handling; equivalent to monkey-patch, but we can't monkey-patch runner.py
-    # This is kinda annoying because it allocates memory in the ShardMetadata object. Can be rethought after Shanghai.
-    immediate_exception: bool = False
-    should_timeout: float | None = None
-
     start_layer: int = Field(ge=0)
     end_layer: int = Field(ge=0)
     n_layers: int = Field(ge=0)
@@ -51,27 +47,56 @@ class BaseShardMetadata(TaggedModel):
             )
         )
 
+    def is_primary_output(self) -> bool:
+        return self.device_rank == self.world_size - 1
+
 
 @final
 class PipelineShardMetadata(BaseShardMetadata):
-    """
-    Pipeline parallelism shard meta.
-
-    Layers are represented as a half-open interval [start_layer, end_layer),
-    where start_layer is inclusive and end_layer is exclusive.
-    """
+    pass
 
 
 @final
 class CfgShardMetadata(BaseShardMetadata):
-    """Shard metadata for CFG-parallel image generation models."""
+    # example
+    # world_size 6
+    # rank prank crank
+    #  0     0     0
+    #  1     1     0
+    #  2     2     0
+    #  3     2     1
+    #  4     1     1
+    #  5     0     1
 
-    cfg_rank: int  # 0 = positive branch, 1 = negative branch
-    cfg_world_size: int = 2
+    @property
+    def cfg_rank(self) -> int:
+        # 0 = positive branch, 1 = negative branch
+        return 0 if self.device_rank < self.world_size // 2 else 1
 
-    # Pipeline-relative coordinates (computed at placement time)
-    pipeline_rank: int  # rank within the pipeline group (0, 1, 2, ...)
-    pipeline_world_size: int  # number of nodes per pipeline group
+    @property
+    def cfg_world_size(self) -> int:
+        return 2
+
+    @property
+    def pipeline_rank(self) -> int:
+        return (
+            self.device_rank
+            if self.cfg_rank == 0
+            else (self.world_size - self.device_rank - 1)
+        )
+
+    @property
+    def pipeline_world_size(self) -> int:
+        return self.world_size // 2
+
+    def is_primary_output(self) -> bool:
+        """
+        For CFG models: the last pipeline stage in CFG group 0 (positive prompt).
+        For non-CFG models: the last pipeline stage.
+        """
+        assert self.pipeline_world_size == self.world_size // 2
+        assert self.world_size % 2 == 0
+        return self.device_rank == (self.world_size // 2) - 1
 
 
 @final

src/exo/utils/__init__.py

@@ -1,4 +1,4 @@
-from typing import Any, Type
+from typing import Any, Callable, Iterable, Iterator, Type, TypeGuard
 
 from .phantom import PhantomData
 
@@ -19,3 +19,11 @@ def todo[T](
     _phantom: PhantomData[T] = None,
 ) -> T:
     raise NotImplementedError(msg)
+
+
+def not_none[T](t: T | None) -> TypeGuard[T]:
+    return t is not None
+
+
+def fmap[T, U](f: Callable[[T], U | None], s: Iterable[T]) -> Iterator[U]:
+    return filter(not_none, map(f, s))

src/exo/utils/info_gatherer/info_gatherer.py

@@ -10,11 +10,13 @@ from typing import Self, cast
 import anyio
 from anyio import fail_after, open_process, to_thread
 from anyio.streams.buffered import BufferedByteReceiveStream
+from exo_rs import LVPublisher, SessionHandle
 from loguru import logger
 from pydantic import ValidationError
 
 from exo.shared.constants import EXO_CONFIG_FILE, EXO_DEFAULT_MODELS_DIR
 from exo.shared.types.backends import Backend
+from exo.shared.types.common import NodeId
 from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import (
     DiskUsage,
@@ -27,7 +29,6 @@ from exo.shared.types.thunderbolt import (
     ThunderboltConnectivity,
     ThunderboltIdentifier,
 )
-from exo.utils.channels import Sender
 from exo.utils.pydantic_ext import TaggedModel
 from exo.utils.task_group import TaskGroup
 
@@ -401,10 +402,42 @@ GatheredInfo = (
 
 @dataclass
 class InfoGatherer:
-    info_sender: Sender[GatheredInfo]
+    session_handle: SessionHandle
+    node_id: NodeId
+    info_senders: dict[str, LVPublisher] = field(init=False, default_factory=dict)
     _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
     _psutil_enabled: bool = field(init=False, default=False)
 
+    async def send(self, info: GatheredInfo):
+        if (tag := info.tag()) not in self.info_senders:
+            self.info_senders[tag] = self.session_handle.last_value_publisher(
+                f"metrics/{self.node_id}/{tag}"
+            )
+        await self.info_senders[tag].put(info.model_dump_json())
+
+    async def run(self):
+        async with self._tg as tg:
+            if IS_DARWIN:
+                tg.start_soon(self._monitor_macmon, 1)
+                tg.start_soon(self._monitor_system_profiler_thunderbolt_data, 5)
+                tg.start_soon(self._monitor_thunderbolt_bridge_status, 10)
+                tg.start_soon(self._monitor_rdma_ctl_status, 10)
+            if not IS_DARWIN:
+                tg.start_soon(self._monitor_memory_usage, 1)
+            tg.start_soon(self._watch_system_info, 10)
+            tg.start_soon(self._monitor_misc, 60)
+            tg.start_soon(self._monitor_static_info, 60)
+            tg.start_soon(self._monitor_disk_usage, 30)
+
+            nc = await NodeConfig.gather()
+            if nc is not None:
+                await self.send(nc)
+
+            await self.send(await NodeBackends.gather())
+
+    def shutdown(self):
+        self._tg.cancel_tasks()
+
     async def _can_read_macmon_metrics(self, macmon_path: str) -> bool:
         try:
             with fail_after(5):
@@ -441,34 +474,11 @@ class InfoGatherer:
 
         return True
 
-    async def run(self):
-        async with self._tg as tg:
-            if IS_DARWIN:
-                tg.start_soon(self._monitor_macmon, 1)
-                tg.start_soon(self._monitor_system_profiler_thunderbolt_data, 5)
-                tg.start_soon(self._monitor_thunderbolt_bridge_status, 10)
-                tg.start_soon(self._monitor_rdma_ctl_status, 10)
-            if not IS_DARWIN:
-                tg.start_soon(self._monitor_memory_usage, 1)
-            tg.start_soon(self._watch_system_info, 10)
-            tg.start_soon(self._monitor_misc, 60)
-            tg.start_soon(self._monitor_static_info, 60)
-            tg.start_soon(self._monitor_disk_usage, 30)
-
-            nc = await NodeConfig.gather()
-            if nc is not None:
-                await self.info_sender.send(nc)
-
-            await self.info_sender.send(await NodeBackends.gather())
-
-    def shutdown(self):
-        self._tg.cancel_tasks()
-
     async def _monitor_static_info(self, static_info_poll_interval: float):
         while True:
             try:
                 with fail_after(30):
-                    await self.info_sender.send(await StaticNodeInformation.gather())
+                    await self.send(await StaticNodeInformation.gather())
             except Exception as e:
                 logger.opt(exception=e).warning("Error gathering static node info")
             await anyio.sleep(static_info_poll_interval)
@@ -477,7 +487,7 @@ class InfoGatherer:
         while True:
             try:
                 with fail_after(10):
-                    await self.info_sender.send(await MiscData.gather())
+                    await self.send(await MiscData.gather())
             except Exception as e:
                 logger.opt(exception=e).warning("Error gathering misc data")
             await anyio.sleep(misc_poll_interval)
@@ -498,12 +508,10 @@ class InfoGatherer:
                     idents = [
                         it for i in data if (it := i.ident(iface_map)) is not None
                     ]
-                    await self.info_sender.send(
-                        MacThunderboltIdentifiers(idents=idents)
-                    )
+                    await self.send(MacThunderboltIdentifiers(idents=idents))
 
                     conns = [it for i in data if (it := i.conn()) is not None]
-                    await self.info_sender.send(MacThunderboltConnections(conns=conns))
+                    await self.send(MacThunderboltConnections(conns=conns))
             except Exception as e:
                 logger.opt(exception=e).warning("Error gathering Thunderbolt data")
             await anyio.sleep(system_profiler_interval)
@@ -520,7 +528,7 @@ class InfoGatherer:
         )
         while True:
             try:
-                await self.info_sender.send(
+                await self.send(
                     MemoryUsage.from_psutil(override_memory=override_memory)
                 )
             except Exception as e:
@@ -532,7 +540,7 @@ class InfoGatherer:
             try:
                 with fail_after(10):
                     nics = await get_network_interfaces()
-                    await self.info_sender.send(NodeNetworkInterfaces(ifaces=nics))
+                    await self.send(NodeNetworkInterfaces(ifaces=nics))
             except Exception as e:
                 logger.opt(exception=e).warning("Error gathering network interfaces")
             await anyio.sleep(interface_watcher_interval)
@@ -545,7 +553,7 @@ class InfoGatherer:
                 with fail_after(30):
                     curr = await ThunderboltBridgeInfo.gather()
                     if curr is not None:
-                        await self.info_sender.send(curr)
+                        await self.send(curr)
             except Exception as e:
                 logger.opt(exception=e).warning(
                     "Error gathering Thunderbolt Bridge status"
@@ -557,7 +565,7 @@ class InfoGatherer:
             try:
                 curr = await RdmaCtlStatus.gather()
                 if curr is not None:
-                    await self.info_sender.send(curr)
+                    await self.send(curr)
             except Exception as e:
                 logger.opt(exception=e).warning("Error gathering RDMA ctl status")
             await anyio.sleep(rdma_ctl_poll_interval)
@@ -566,7 +574,7 @@ class InfoGatherer:
         while True:
             try:
                 with fail_after(5):
-                    await self.info_sender.send(await NodeDiskUsage.gather())
+                    await self.send(await NodeDiskUsage.gather())
             except Exception as e:
                 logger.opt(exception=e).warning("Error gathering disk usage")
             await anyio.sleep(disk_poll_interval)
@@ -611,7 +619,7 @@ class InfoGatherer:
                             )
                             text = data.decode("utf-8", errors="replace").strip()
                             metrics = MacmonMetrics.from_raw_json(text)
-                        await self.info_sender.send(metrics)
+                        await self.send(metrics)
             except TimeoutError:
                 logger.warning(
                     f"MacMon produced no output for {read_timeout}s, restarting"

src/exo/utils/pydantic_ext.py

@@ -19,18 +19,22 @@ class FrozenModel(BaseModel):
 
 
 class TaggedModel(FrozenModel):
+    @classmethod
+    def tag(cls) -> str:
+        return cls.__name__
+
     @model_serializer(mode="wrap")
     def _serialize(self, handler: SerializerFunctionWrapHandler):
         inner = handler(self)  # pyright: ignore[reportAny]
-        return {self.__class__.__name__: inner}
+        return {self.tag(): inner}
 
     @model_validator(mode="wrap")
     @classmethod
     def _validate(cls, v: Any, handler: ValidatorFunctionWrapHandler) -> Self:  # pyright: ignore[reportAny]
-        if isinstance(v, dict) and len(v) == 1 and cls.__name__ in v:  # pyright: ignore[reportUnknownArgumentType]
-            return handler(v[cls.__name__])  # pyright: ignore[reportAny]
+        if isinstance(v, dict) and len(v) == 1 and cls.tag() in v:  # pyright: ignore[reportUnknownArgumentType]
+            return handler(v[cls.tag()])  # pyright: ignore[reportAny]
 
         return handler(v)  # pyright: ignore[reportAny]
 
     def __str__(self) -> str:
-        return f"{self.__class__.__name__}({super().__str__()})"
+        return f"{self.tag()}({super().__str__()})"

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

@@ -30,8 +30,6 @@ from exo.shared.types.worker.runner_response import (
     ModelLoadingResponse,
 )
 from exo.shared.types.worker.shards import (
-    CfgShardMetadata,
-    PipelineShardMetadata,
     ShardMetadata,
 )
 from exo.utils.channels import MpReceiver, MpSender
@@ -49,22 +47,6 @@ from exo.worker.engines.mlx.utils_mlx import (
 )
 
 
-def _is_primary_output_node(shard_metadata: ShardMetadata) -> bool:
-    """Check if this node is the primary output node for image generation.
-
-    For CFG models: the last pipeline stage in CFG group 0 (positive prompt).
-    For non-CFG models: the last pipeline stage.
-    """
-    if isinstance(shard_metadata, CfgShardMetadata):
-        is_pipeline_last = (
-            shard_metadata.pipeline_rank == shard_metadata.pipeline_world_size - 1
-        )
-        return is_pipeline_last and shard_metadata.cfg_rank == 0
-    elif isinstance(shard_metadata, PipelineShardMetadata):
-        return shard_metadata.device_rank == shard_metadata.world_size - 1
-    return False
-
-
 def _send_traces_if_enabled(
     event_sender: MpSender[Event],
     task_id: TaskId,
@@ -171,7 +153,7 @@ class ImageEngine(Engine):
             resp = next(self.current_gen, None)
         return (
             (resp,)
-            if resp is not None and _is_primary_output_node(self.shard_metadata)
+            if resp is not None and self.shard_metadata.is_primary_output()
             else ()
         )
 
@@ -202,10 +184,10 @@ class ImageEngine(Engine):
                 task=task_params,
                 cancel_checker=cancel_checker,
             ):
-                if _is_primary_output_node(self.shard_metadata):
+                if self.shard_metadata.is_primary_output():
                     yield (task_id, response)
         except Exception as e:
-            if _is_primary_output_node(self.shard_metadata):
+            if self.shard_metadata.is_primary_output():
                 yield (
                     task_id,
                     ErrorChunk(

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

@@ -154,7 +154,7 @@ def initialize_mlx(
     # TODO: pass in seed from params
     mx.random.seed(42)
 
-    assert len(bound_instance.instance.shard_assignments.node_to_runner) > 1, (
+    assert len(bound_instance.instance.shard_assignments.shards) > 1, (
         "Tried to initialize mlx for a single node instance"
     )
     return mlx_distributed_init(bound_instance)

src/exo/worker/main.py

@@ -4,7 +4,9 @@ from datetime import datetime, timezone
 
 import anyio
 from anyio import fail_after, to_thread
+from exo_rs import LVAggregator, SessionHandle
 from loguru import logger
+from pydantic import ValidationError
 
 from exo.api.types import ImageEditsTaskParams
 from exo.download.download_utils import is_read_only_model_dir, resolve_existing_model
@@ -14,7 +16,8 @@ from exo.routing.event_router import (
 )
 from exo.shared.apply import apply
 from exo.shared.constants import EXO_MAX_INSTANCE_RETRIES
-from exo.shared.models.model_cards import ModelId, card_cache
+from exo.shared.models import model_cards
+from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.chunks import InputImageChunk
 from exo.shared.types.commands import (
     DeleteInstance,
@@ -53,7 +56,7 @@ from exo.shared.types.topology import Connection, SocketConnection
 from exo.shared.types.worker.downloads import DownloadCompleted
 from exo.shared.types.worker.instances import InstanceId
 from exo.shared.types.worker.runners import RunnerId
-from exo.utils.channels import Receiver, Sender, channel
+from exo.utils.channels import Receiver, Sender
 from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer
 from exo.utils.info_gatherer.net_profile import check_reachable
 from exo.utils.keyed_backoff import KeyedBackoff
@@ -73,6 +76,7 @@ class Worker:
         # but I think it's the correct way to be thinking about commands
         command_sender: Sender[ForwarderCommand],
         download_command_sender: Sender[ForwarderDownloadCommand],
+        session_handle: SessionHandle,
         api_port: int,
     ):
         self.node_id: NodeId = node_id
@@ -98,17 +102,17 @@ class Worker:
             base=0.5, cap=10.0
         )
         self._stopped: anyio.Event = anyio.Event()
+        self._sh: SessionHandle = session_handle
+        self.aggregator: LVAggregator = session_handle.last_value_aggregator("metrics")
 
     async def run(self):
         logger.info("Starting Worker")
 
-        info_send, info_recv = channel[GatheredInfo]()
-        info_gatherer: InfoGatherer = InfoGatherer(info_send)
+        info_gatherer: InfoGatherer = InfoGatherer(self._sh, self.node_id)
 
         try:
             async with self._tg as tg:
                 tg.start_soon(info_gatherer.run)
-                tg.start_soon(self._forward_info, info_recv)
                 tg.start_soon(self.plan_step)
                 tg.start_soon(self._event_applier)
                 tg.start_soon(self._poll_connection_updates)
@@ -178,17 +182,24 @@ class Worker:
                             ] = img
 
     async def _reconcile_custom_cards(self) -> None:
+        storage = self._sh.storage_interface()
         while True:
-            await anyio.sleep(1)
-            target = dict(self.state.custom_model_cards)
-            for model_id, card in target.items():
-                if card_cache.get(model_id) == card:
+            await anyio.sleep(10)
+            target: list[ModelId] = []
+            for _, value in (await storage.dump("custom_model_cards/")).items():
+                try:
+                    card = ModelCard.model_validate_json(value)
+                except ValidationError:
                     continue
-                await card_cache.save(card)
+                target.append(card.model_id)
+                if model_cards.card_cache.get(card.model_id) == card:
+                    continue
+                logger.info(f"Registered new custom model card for {card.model_id}")
+                await model_cards.card_cache.save(card)
 
-            for card in await card_cache.list_all():
+            for card in await model_cards.card_cache.list_all():
                 if card.model_id not in target:
-                    await card_cache.pop(card.model_id)
+                    await model_cards.card_cache.delete(card.model_id)
 
     async def plan_step(self):
         while True:
@@ -372,9 +383,8 @@ class Worker:
 
     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)
+            for rid in instance.runners_for(self.node_id):
+                await self.runners[rid].start_task(task)
 
     async def _create_supervisor(self, task: CreateRunner) -> RunnerSupervisor:
         """Creates and stores a new AssignedRunner with initial downloading status."""
@@ -388,14 +398,13 @@ class Worker:
 
     async def _poll_connection_updates(self):
         while True:
-            edges = set(
-                conn.edge for conn in self.state.topology.out_edges(self.node_id)
-            )
+            state = self.state.with_aggregator(self.aggregator)
+            edges = set(conn.edge for conn in state.topology.out_edges(self.node_id))
             conns: defaultdict[NodeId, set[str]] = defaultdict(set)
             async for ip, nid in check_reachable(
-                self.state.topology,
+                state.topology,
                 self.node_id,
-                self.state.node_network,
+                state.node_network,
                 api_port=self.api_port,
             ):
                 if ip in conns[nid]:
@@ -416,7 +425,7 @@ class Worker:
                         )
                     )
 
-            for conn in self.state.topology.out_edges(self.node_id):
+            for conn in state.topology.out_edges(self.node_id):
                 if not isinstance(conn.edge, SocketConnection):
                     continue
                 # ignore mDNS discovered connections

src/exo/worker/plan.py

@@ -40,6 +40,7 @@ from exo.shared.types.worker.runners import (
     RunnerStatus,
     RunnerWarmingUp,
 )
+from exo.utils import fmap
 from exo.utils.keyed_backoff import KeyedBackoff
 from exo.worker.runner.supervisor import RunnerSupervisor
 
@@ -88,8 +89,10 @@ def _kill_runner(
             )
 
         for (
-            global_runner_id
-        ) in runner.bound_instance.instance.shard_assignments.node_to_runner.values():
+            _,
+            global_runner_id,
+            _,
+        ) in runner.bound_instance.instance.shard_assignments.shards:
             if runner_id == global_runner_id:
                 continue
 
@@ -108,7 +111,13 @@ def _create_runner(
     instance_backoff: KeyedBackoff[InstanceId],
 ) -> CreateRunner | None:
     for instance in instances.values():
-        runner_id = instance.shard_assignments.node_to_runner.get(node_id, None)
+        runner_id = next(
+            fmap(
+                lambda it: it.runner_id if it.node_id == node_id else None,
+                instance.shard_assignments.shards,
+            ),
+            None,
+        )
         if runner_id is None:
             continue
 
@@ -118,7 +127,7 @@ def _create_runner(
         # don't create runners if any other nodes have runners that have failed - wait for them to fix themselves first.
         instance_has_failed_runner = any(
             isinstance(all_runners.get(remote_runner_id), RunnerFailed)
-            for remote_runner_id in instance.shard_assignments.node_to_runner.values()
+            for (_, remote_runner_id, _) in instance.shard_assignments.shards
             if remote_runner_id != runner_id
         )
         we_have_failed_before = isinstance(all_runners.get(runner_id), RunnerFailed)
@@ -175,7 +184,7 @@ def _init_distributed_backend(
         instance = runner.bound_instance.instance
         shard_assignments = instance.shard_assignments
 
-        is_single_node_instance = len(shard_assignments.runner_to_shard) == 1
+        is_single_node_instance = len(shard_assignments.shards) == 1
         if is_single_node_instance:
             continue
 
@@ -185,7 +194,7 @@ def _init_distributed_backend(
                 all_runners.get(global_runner_id),
                 (RunnerConnecting, RunnerIdle),
             )
-            for global_runner_id in shard_assignments.runner_to_shard
+            for (_, global_runner_id, _) in shard_assignments.shards
         )
 
         if not (runner_is_idle and all_runners_connecting):
@@ -205,7 +214,7 @@ def _init_distributed_backend(
         # Rank = n-1
         connecting_rank_ready = device_rank == world_size - 1 and all(
             isinstance(all_runners.get(global_runner_id, None), RunnerConnecting)
-            for global_runner_id in shard_assignments.runner_to_shard
+            for (_, global_runner_id, _) in shard_assignments.shards
             if global_runner_id != runner_id
         )
 
@@ -233,12 +242,12 @@ def _load_model(
                 and dp.shard_metadata.model_card.model_id == shard_assignments.model_id
                 for dp in global_download_status[nid]
             )
-            for nid in shard_assignments.node_to_runner
+            for (nid, _, _) in shard_assignments.shards
         )
         if not all_local_downloads_complete:
             continue
 
-        is_single_node_instance = len(instance.shard_assignments.runner_to_shard) == 1
+        is_single_node_instance = len(instance.shard_assignments.shards) == 1
         if is_single_node_instance and isinstance(runner.status, RunnerIdle):
             return LoadModel(instance_id=instance.instance_id)
 
@@ -249,7 +258,7 @@ def _load_model(
                 all_runners.get(global_runner_id, None),
                 (RunnerConnected, RunnerLoading, RunnerLoaded),
             )
-            for global_runner_id in shard_assignments.runner_to_shard
+            for (_, global_runner_id, _) in shard_assignments.shards
         )
 
         if is_runner_waiting and all_ready_for_model:
@@ -281,13 +290,13 @@ def _ready_to_warmup(
                 all_runners.get(global_runner_id, None),
                 (RunnerLoaded, RunnerWarmingUp),
             )
-            for global_runner_id in shard_assignments.runner_to_shard
+            for (_, global_runner_id, _) in shard_assignments.shards
         )
 
         # Rank = 0
         connecting_rank_ready = device_rank == 0 and all(
             isinstance(all_runners.get(global_runner_id, None), RunnerWarmingUp)
-            for global_runner_id in shard_assignments.runner_to_shard
+            for (_, global_runner_id, _) in shard_assignments.shards
             if global_runner_id != runner_id
         )
 
@@ -338,7 +347,11 @@ def _pending_tasks(
 
             if isinstance(runner.status, (RunnerReady, RunnerRunning)) and all(
                 isinstance(all_runners[global_runner_id], (RunnerReady, RunnerRunning))
-                for global_runner_id in runner.bound_instance.instance.shard_assignments.runner_to_shard
+                for (
+                    _,
+                    global_runner_id,
+                    _,
+                ) in runner.bound_instance.instance.shard_assignments.shards
             ):
                 return task
 

src/exo/worker/runner/runner.py

@@ -102,10 +102,6 @@ class Runner:
         self.device_rank = self.shard_metadata.device_rank
 
         logger.info("hello from the runner")
-        if getattr(self.shard_metadata, "immediate_exception", False):
-            raise Exception("Fake exception - runner failed to spin up.")
-        if timeout := getattr(self.shard_metadata, "should_timeout", 0):
-            time.sleep(timeout)
 
         self.setup_start_time = time.time()
 

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

@@ -11,7 +11,12 @@ from exo.shared.types.worker.instances import (
     InstanceId,
     MlxRingInstance,
 )
-from exo.shared.types.worker.runners import RunnerId, RunnerStatus, ShardAssignments
+from exo.shared.types.worker.runners import (
+    RunnerId,
+    RunnerStatus,
+    ShardAssignments,
+    ShardWithId,
+)
 from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
 
 
@@ -52,16 +57,22 @@ def get_pipeline_shard_metadata(
     )
 
 
+# todo: clean up legacy formatted shards
 def get_shard_assignments(
     model_id: ModelId,
     node_to_runner: dict[NodeId, RunnerId],
     runner_to_shard: dict[RunnerId, ShardMetadata],
 ) -> ShardAssignments:
-    return ShardAssignments(
-        model_id=model_id,
-        node_to_runner=node_to_runner,
-        runner_to_shard=runner_to_shard,
-    )
+    pon = 0
+    shards = [
+        ShardWithId(nid, rid := node_to_runner[nid], runner_to_shard[rid])
+        for nid in node_to_runner
+    ]
+    for i, (_, _, shard) in enumerate(shards):
+        if shard.is_primary_output():
+            pon = i
+
+    return ShardAssignments(model_id=model_id, shards=shards, primary_output_node=pon)
 
 
 def get_mlx_ring_instance(

tmp/run_exo_on.sh

@@ -42,7 +42,7 @@ i=0
 for host; do
   colour=${colours[i++ % 4]}
   ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
-    "EXO_ZENOH_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix run $remote_installable" 2>&1 |
+    "ENABLE_DISAGGREGATION=true EXO_ZENOH_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix run $remote_installable" 2>&1 |
     awk -v p="${colour}[${host}]${reset}" '{ print p $0; fflush() }' &
 done
 

tools/src/exo_tools/harness.py

@@ -52,18 +52,17 @@ def instance_id_from_instance(instance: dict[str, Any]) -> str:
 
 def nodes_used_in_instance(instance: dict[str, Any]) -> int:
     inner = unwrap_instance(instance)
-    return len(inner["shardAssignments"]["nodeToRunner"])
+    return len(inner["shardAssignments"]["shards"])
 
 
 def runner_ids_from_instance(instance: dict[str, Any]) -> list[str]:
     inner = unwrap_instance(instance)
-    runner_to_shard = inner["shardAssignments"]["runnerToShard"]
-    return list(runner_to_shard.keys())
+    return [r for (_, r, _) in inner["shardAssignments"]["shards"]]
 
 
 def node_ids_from_instance(instance: dict[str, Any]) -> list[str]:
     inner = unwrap_instance(instance)
-    return list(inner["shardAssignments"]["nodeToRunner"].keys())
+    return [n for (n, _, _) in inner["shardAssignments"]["shards"]]
 
 
 def runner_ready(runner: dict[str, Any]) -> bool:
@@ -322,8 +321,8 @@ def run_planning_phase(
 
     # Get nodes from preview
     inner = unwrap_instance(preview["instance"])
-    node_ids = list(inner["shardAssignments"]["nodeToRunner"].keys())
-    runner_to_shard = inner["shardAssignments"]["runnerToShard"]
+    node_ids = [n for (n, _, _) in inner["shardAssignments"]["shards"]]
+    shards = inner["shardAssignments"]["shards"]
 
     needs_download = False
 
@@ -391,9 +390,7 @@ def run_planning_phase(
 
     # Start downloads (idempotent)
     download_t0 = time.perf_counter() if needs_download else None
-    for node_id in node_ids:
-        runner_id = inner["shardAssignments"]["nodeToRunner"][node_id]
-        shard = runner_to_shard[runner_id]
+    for node_id, _, shard in shards:
         client.request_json(
             "POST",
             "/download/start",

uv.lock

@@ -450,7 +450,7 @@ build = [
 ]
 mlx = [
     { name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
-    { name = "mlx", version = "0.32.0.dev20260506+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
+    { name = "mlx", version = "0.32.0.dev20260512+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl" }, marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
@@ -461,7 +461,7 @@ mlx = [
 ]
 mlx-cpu = [
     { name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
-    { name = "mlx", version = "0.32.0.dev20260506+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-mlx-cpu') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
+    { name = "mlx", version = "0.32.0.dev20260512+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-mlx-cpu') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cpu') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra != 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra != 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra != 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl" }, marker = "(platform_machine != 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cpu') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra != 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra != 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra != 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "mlx-cpu", marker = "sys_platform == 'linux'" },
@@ -473,7 +473,7 @@ mlx-cpu = [
 ]
 mlx-cuda12 = [
     { name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
-    { name = "mlx", version = "0.32.0.dev20260506+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
+    { name = "mlx", version = "0.32.0.dev20260512+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra != 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl" }, marker = "(platform_machine != 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra != 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "mlx-cuda-12", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx_cuda_12-0.32.0-py3-none-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra != 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
@@ -487,7 +487,7 @@ mlx-cuda12 = [
 ]
 mlx-cuda13 = [
     { name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
-    { name = "mlx", version = "0.32.0.dev20260506+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none')" },
+    { name = "mlx", version = "0.32.0.dev20260512+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none')" },
     { name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none')" },
     { name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl" }, marker = "(platform_machine != 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none')" },
     { name = "mlx-cuda-13", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx_cuda_13-0.32.0-py3-none-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none')" },
@@ -1343,7 +1343,7 @@ dependencies = [
     { name = "hf-transfer", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "matplotlib", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
-    { name = "mlx", version = "0.32.0.dev20260506+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
+    { name = "mlx", version = "0.32.0.dev20260512+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl" }, marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
@@ -1374,7 +1374,7 @@ sdist = { url = "https://files.pythonhosted.org/packages/55/fa/96d4cc7ada2833571
 
 [[package]]
 name = "mlx"
-version = "0.32.0.dev20260506+cc3f3e60"
+version = "0.32.0.dev20260512+cc3f3e60"
 source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }
 resolution-markers = [
     "sys_platform == 'darwin'",
@@ -1549,7 +1549,7 @@ version = "0.31.3"
 source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Fdeepseek-v4#6a3df6cd6b00a347ee40f12d97a182aaf86ea599" }
 dependencies = [
     { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
-    { name = "mlx", version = "0.32.0.dev20260506+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
+    { name = "mlx", version = "0.32.0.dev20260512+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
@@ -1565,7 +1565,7 @@ dependencies = [
     { name = "datasets", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "fastapi", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "miniaudio", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
-    { name = "mlx", version = "0.32.0.dev20260506+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
+    { name = "mlx", version = "0.32.0.dev20260512+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl" }, marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },