rebase fix

## Motivation

EXO should be PID file locked, to prevent duplicate processes from
clobbering the log, right now this isn't the case.

## Changes

I added a wrapper around a Rust PID file lock library, and used it to
implement PID locking for EXO, with the PID file being in exo cache
directory.

## Test Plan

### Manual Testing
Tested on e11, trying to spawn duplicate EXO processes prevented.

.gitignore

@@ -40,3 +40,4 @@ bench/**/*.json
 tmp/models
 /build/exo
 /.claude/skills
+/.claude

Cargo.toml

@@ -1,6 +1,6 @@
 [workspace]
 resolver = "3"
-members = ["rust/networking", "rust/exo_pyo3_bindings", "rust/util"]
+members = ["rust/exo_net", "rust/networking"]
 
 [workspace.package]
 version = "0.0.1"
@@ -20,30 +20,34 @@ opt-level = 3
 [workspace.dependencies]
 ## Crate members as common dependencies
 networking = { path = "rust/networking" }
-util = { path = "rust/util" }
 
-# Macro dependecies
+# pyo3
+pyo3 = "0.27.2"
+pyo3-async-runtimes = "0.27.0"
+pyo3-log = "0.13.2"
+pyo3-stub-gen = "0.22.2"
+
+# util
 extend = "1.2"
-delegate = "0.13"
-
-# Utility dependencies
-keccak-const = "0.2"
-
-# Async dependencies
-async-stream = "0.3"
 tokio = "1.46"
 futures-lite = "2.6.1"
-futures-timer = "3.0"
-
-# Data structures
-either = "1.15"
+async-stream = "0.3.6"
+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"
+env_logger = "0.11.10"
 
 # networking
-libp2p = "0.56"
-libp2p-tcp = "0.44"
+zenoh = "=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"
 
 [workspace.lints.rust]
 static_mut_refs = "warn"      # Or use "warn" instead of deny

bench/eval_tool_calls.py

@@ -15,9 +15,8 @@ from pathlib import Path
 from typing import Any, Literal
 
 import httpx
-from harness import (
-    ExoClient,
-    ExoHttpError,
+from exo_tools.client import ExoClient, ExoHttpError
+from exo_tools.harness import (
     add_common_instance_args,
     capture_cluster_snapshot,
     instance_id_from_instance,

bench/exo_bench.py

@@ -30,9 +30,8 @@ from pathlib import Path
 from statistics import mean
 from typing import Any
 
-from harness import (
-    ExoClient,
-    ExoHttpError,
+from exo_tools.client import ExoClient, ExoHttpError
+from exo_tools.harness import (
     add_common_instance_args,
     capture_cluster_snapshot,
     find_existing_instance,

bench/exo_eval.py

@@ -42,9 +42,8 @@ from pathlib import Path
 from typing import Any
 
 import httpx
-from harness import (
-    ExoClient,
-    ExoHttpError,
+from exo_tools.client import ExoClient, ExoHttpError
+from exo_tools.harness import (
     add_common_instance_args,
     capture_cluster_snapshot,
     find_existing_instance,

bench/prefill_decode_bench.py

@@ -35,9 +35,8 @@ from exo_bench import (
     load_tokenizer_for_bench,
     parse_int_list,
 )
-from harness import (
-    ExoClient,
-    ExoHttpError,
+from exo_tools.client import ExoClient, ExoHttpError
+from exo_tools.harness import (
     add_common_instance_args,
     instance_id_from_instance,
     node_ids_from_instance,

flake.lock

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

flake.nix

@@ -110,7 +110,7 @@
               nixpkgs-fmt.enable = true;
               ruff-format = {
                 enable = true;
-                excludes = [ "rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi" ];
+                excludes = [ "rust/exo_net/exo_net.pyi" ];
               };
               rustfmt = {
                 enable = true;
@@ -146,7 +146,7 @@
                 config.treefmt.build.wrapper
 
                 # PYTHON
-                self'.packages.editableVenv
+                #self'.packages.editableVenv
                 uv
 
                 # RUST

justfile

@@ -23,7 +23,7 @@ sync-clean:
 
 rust-rebuild:
     PYO3_PYTHON="$(uv run python -c 'import sys; print(sys.executable)')" cargo run --bin stub_gen
-    uv sync --reinstall-package exo_pyo3_bindings
+    uv sync --reinstall-package exo_net
 
 build-dashboard:
     #!/usr/bin/env bash

pyproject.toml

@@ -15,7 +15,7 @@ dependencies = [
   "huggingface-hub>=1.8.0",
   "psutil>=7.0.0",
   "loguru>=0.7.3",
-  "exo-pyo3-bindings",                         # rust bindings
+  "exo-net",                                   # rust bindings
   "anyio==4.11.0",
   "mlx==0.31.2; sys_platform == 'darwin'",
   "mlx-lm; sys_platform=='darwin'",
@@ -40,6 +40,7 @@ exo = "exo.main:main"
 dev = [
   "basedpyright>=1.29.0",
   "pyinstaller>=6.17.0",
+  "playwright>=1.52.0",
   "pytest>=8.4.0",
   "pytest-asyncio>=1.0.0",
   "pytest-env",
@@ -75,10 +76,10 @@ cuda13 = [
 ###
 
 [tool.uv.workspace]
-members = ["rust/exo_pyo3_bindings", "bench"]
+members = ["rust/exo_net", "bench", "tools"]
 
 [tool.uv.sources]
-exo-pyo3-bindings = { workspace = true }
+exo-net = { workspace = true }
 mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
 mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/deepseek-v4" }
 torch = [
@@ -112,7 +113,7 @@ build-backend = "uv_build"
 ###
 
 [tool.basedpyright]
-include = ["src", "bench"]
+include = ["src", "bench", "tools"]
 typeCheckingMode = "strict"
 failOnWarnings = true
 
@@ -146,6 +147,13 @@ reportMissingModuleSource = false
 [[tool.basedpyright.executionEnvironments]]
 root = "src"
 
+[[tool.basedpyright.executionEnvironments]]
+root = "bench"
+extraPaths = ["tools/src"]
+
+[[tool.basedpyright.executionEnvironments]]
+root = "tools/src"
+
 
 ###
 # uv configuration
@@ -206,11 +214,7 @@ torchaudio = ["torch"]
 ###
 
 [tool.ruff]
-extend-exclude = [
-  "*mlx_typings/**",
-  "rust/exo_pyo3_bindings/**",
-  "bench/vendor/**",
-]
+extend-exclude = ["*mlx_typings/**", "rust/exo_net/**", "bench/vendor/**"]
 
 [tool.ruff.lint]
 extend-select = ["I", "N", "B", "A", "PIE", "SIM"]
@@ -220,5 +224,5 @@ pythonpath = "."
 asyncio_mode = "auto"
 markers = ["slow: marks tests as slow (deselected by default)"]
 env = ["EXO_TESTS=1"]
-addopts = "-m 'not slow' --ignore=tests/start_distributed_test.py"
+addopts = "-m 'not slow' --ignore=tests"
 filterwarnings = ["ignore:builtin type Swig:DeprecationWarning"]

python/parts.nix

@@ -35,17 +35,18 @@ let
         # Replace workspace exo_pyo3_bindings with Nix-built wheel.
         # Preserve passthru so mkVirtualEnv can resolve dependency groups.
         # Copy .pyi stub + py.typed marker so basedpyright can find the types.
-        exo-pyo3-bindings = pkgs.stdenv.mkDerivation {
-          pname = "exo-pyo3-bindings";
+        exo-net = pkgs.stdenv.mkDerivation {
+          pname = "exo-net";
           version = "0.1.0";
-          src = self'.packages.exo_pyo3_bindings;
+          src = self'.packages.exo-net;
           # Install from pre-built wheel
           nativeBuildInputs = [ final.pyprojectWheelHook ];
           dontStrip = true;
           passthru = prev.exo-pyo3-bindings.passthru or { };
+
           postInstall = ''
-            local siteDir=$out/${final.python.sitePackages}/exo_pyo3_bindings
-            cp ${inputs.self}/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi $siteDir/
+            local siteDir=$out/${final.python.sitePackages}/exo_net
+            cp ${inputs.self}/rust/exo_net/exo_net.pyi $siteDir/
             touch $siteDir/py.typed
           '';
         };

rust/exo_net/Cargo.toml

@@ -0,0 +1,52 @@
+[package]
+name = "exo_net"
+version = { workspace = true }
+edition = { workspace = true }
+publish = false
+
+[lib]
+doctest = false
+path = "src/lib.rs"
+name = "exo_net"
+
+# "cdylib" needed to produce shared library for Python to import
+# "rlib" needed for stub-gen to run
+crate-type = ["cdylib", "rlib"]
+
+[[bin]]
+path = "src/bin/stub_gen.rs"
+name = "stub_gen"
+doc = false
+
+[lints]
+workspace = true
+
+[dependencies]
+networking.workspace = true
+extend.workspace = true
+
+# interop
+pyo3 = { workspace = true, features = ["experimental-async"] }
+pyo3-stub-gen.workspace = true
+pyo3-async-runtimes = { workspace = true, features = [
+  "attributes",
+  "tokio-runtime",
+  "testing",
+] }
+pyo3-log.workspace = true
+
+# async runtime
+tokio = { workspace = true, features = ["full"] }
+futures-lite.workspace = true
+pin-project.workspace = true
+
+# Tracing
+log.workspace = true
+env_logger.workspace = true
+
+# Networking
+zenoh.workspace = true
+rand.workspace = true
+serde_json.workspace = true
+parking_lot.workspace = true
+pidfile-rs.workspace = true

rust/exo_net/exo_net.pyi

@@ -0,0 +1,122 @@
+# This file is automatically generated by pyo3_stub_gen
+# ruff: noqa: E501, F401, F403, F405
+
+import builtins
+import collections.abc
+import os
+import pathlib
+import typing
+__all__ = [
+    "NetReceiver",
+    "NetSender",
+    "NetworkingHandle",
+    "Pidfile",
+    "PidfileError",
+    "PyFromSwarm",
+    "PySession",
+    "StateProxy",
+]
+
+@typing.final
+class NetReceiver:
+    def recv(self) -> collections.abc.Awaitable[bytes | None]: ...
+
+@typing.final
+class NetSender:
+    def send(self, data: bytes) -> collections.abc.Awaitable[bool]: ...
+
+@typing.final
+class NetworkingHandle:
+    @staticmethod
+    def new(identity: bytes, bootstrap_peers: typing.Sequence[builtins.str], listen_port: builtins.int) -> tuple[NetworkingHandle, PySession]: ...
+    async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
+        r"""
+        Subscribe to a `GossipSub` topic.
+        
+        Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
+        """
+    async def gossipsub_unsubscribe(self, topic: builtins.str) -> builtins.bool:
+        r"""
+        Unsubscribes from a `GossipSub` topic.
+        
+        Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
+        """
+    async def gossipsub_publish(self, topic: builtins.str, data: bytes) -> None:
+        r"""
+        Publishes a message with multiple topics to the `GossipSub` network.
+        
+        If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
+        """
+    async def recv(self) -> PyFromSwarm: ...
+
+@typing.final
+class Pidfile:
+    r"""
+    A PID file protected with a lock.
+    
+    An instance of `Pidfile` can be used to manage a PID file: create it,
+    lock it, detect already running daemons. It is backed by [`pidfile`][]
+    functions of `libbsd`/`libutil` which use `flopen` to lock the PID
+    file.
+    
+    When a PID file is created, the process ID of the current process is
+    *not* written there, making it possible to lock the PID file before
+    forking and only write the ID of the forked process when it is ready.
+    
+    The PID file is deleted automatically when the `Pidfile` comes out of
+    the scope. To close the PID file without deleting it, for example, in
+    the parent process of a forked daemon, call `close()`.
+    
+    [`exit`]: https://doc.rust-lang.org/std/process/fn.exit.html
+    [`pidfile`]: https://linux.die.net/man/3/pidfile
+    [`daemon`(3)]: https://linux.die.net/man/3/daemon
+    """
+    def __new__(cls, path: builtins.str | os.PathLike | pathlib.Path, mode: builtins.int) -> Pidfile:
+        r"""
+        Creates a new PID file and locks it.
+        
+        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
+        the PID file yet.
+        """
+    def write(self) -> None:
+        r"""
+        Writes the current process ID to the PID file.
+        
+        The file is truncated before writing.
+        """
+
+@typing.final
+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 PySession:
+    def net_receiver(self, key: builtins.str) -> NetReceiver: ...
+    def net_sender(self, key: builtins.str) -> NetSender: ...
+    def state_proxy(self) -> StateProxy: ...
+
+@typing.final
+class StateProxy:
+    def snapshot(self) -> collections.abc.Awaitable[str]: ...
+

rust/exo_net/pyproject.toml

@@ -3,24 +3,22 @@ requires = ["maturin>=1.0,<2.0"]
 build-backend = "maturin"
 
 [project]
-name = "exo_pyo3_bindings"
-version = "0.2.1"
+name = "exo_net"
+version = "0.3.0"
 description = "Add your description here"
 readme = "README.md"
 authors = [
-  { name = "Andrei Cravtov", email = "the.andrei.cravtov@gmail.com" },
   { name = "Evan Quiney", email = "evanev7@gmail.com" },
+  { name = "Andrei Cravtov", email = "the.andrei.cravtov@gmail.com" },
 ]
 requires-python = ">=3.13"
 dependencies = []
 
 [dependency-groups]
-dev = ["exo_pyo3_bindings", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"]
+dev = ["exo-net", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"]
 
 [tool.maturin]
-#purelib = true
-#python-source = "python"
-module-name = "exo_pyo3_bindings"
+module-name = "exo_net"
 features = ["pyo3/extension-module", "pyo3/experimental-async"]
 
 [tool.pytest.ini_options]

rust/exo_net/src/bin/stub_gen.rs

@@ -2,7 +2,7 @@ use pyo3_stub_gen::Result;
 
 fn main() -> Result<()> {
     env_logger::Builder::from_env(env_logger::Env::default().filter_or("RUST_LOG", "info")).init();
-    let stub = exo_pyo3_bindings::stub_info()?;
+    let stub = exo_net::stub_info()?;
     stub.generate()?;
     Ok(())
 }

rust/exo_net/src/ident.rs

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

rust/exo_net/src/lib.rs

@@ -5,21 +5,23 @@
 //!
 
 mod allow_threading;
-mod ident;
+mod pidfile;
+// mod ident;
 mod networking;
+mod point_to_point;
+mod session;
+mod state;
 
-use crate::ident::PyKeypair;
 use crate::networking::networking_submodule;
+use crate::pidfile::pidfile_submodule;
+use crate::point_to_point::{NetReceiver, NetSender};
+use crate::session::PySession;
+use crate::state::StateProxy;
 use pyo3::prelude::PyModule;
 use pyo3::types::PyModuleMethods;
-use pyo3::{Bound, PyResult, pyclass, pymodule};
+use pyo3::{Bound, PyResult, pymodule};
 use pyo3_stub_gen::define_stub_info_gatherer;
 
-/// Namespace for all the constants used by this crate.
-pub(crate) mod r#const {
-    pub const MPSC_CHANNEL_SIZE: usize = 1024;
-}
-
 /// Namespace for crate-wide extension traits/methods
 pub(crate) mod ext {
     use crate::allow_threading::AllowThreads;
@@ -151,7 +153,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_pyo3_bindings")]
+#[pymodule(name = "exo_net")]
 fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     // install logger
     pyo3_log::init();
@@ -162,7 +164,13 @@ fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     // TODO: for now this is all NOT a submodule, but figure out how to make the submodule system
     //       work with maturin, where the types generate correctly, in the right folder, without
     //       too many importing issues...
-    m.add_class::<PyKeypair>()?;
+    pidfile_submodule(m)?;
+    // m.add_class::<PyKeypair>()?;
+    // networking_submodule(m)?;
+    m.add_class::<StateProxy>()?;
+    m.add_class::<PySession>()?;
+    m.add_class::<NetReceiver>()?;
+    m.add_class::<NetSender>()?;
     networking_submodule(m)?;
 
     // top-level constructs

rust/exo_net/src/networking.rs

@@ -0,0 +1,192 @@
+use std::pin::Pin;
+use std::sync::Arc;
+
+use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
+use crate::ext::{ResultExt as _, TokioMpscSenderExt as _};
+use crate::session::PySession;
+use futures_lite::{Stream, StreamExt as _};
+use networking::swarm::{FromSwarm, ToSwarm, create_swarm};
+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_methods_from_python, gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods,
+};
+use tokio::sync::{Mutex, mpsc, oneshot};
+
+#[gen_stub_pyclass]
+#[pyclass(name = "NetworkingHandle")]
+struct PyNetworkingHandle {
+    // channels
+    pub to_swarm: mpsc::Sender<ToSwarm>,
+    pub swarm: Arc<Mutex<Pin<Box<dyn Stream<Item = FromSwarm> + Send>>>>,
+}
+
+#[gen_stub_pyclass_complex_enum]
+#[pyclass]
+enum PyFromSwarm {
+    Connection { connected: bool },
+    Message { topic: String, data: Py<PyBytes> },
+}
+impl From<FromSwarm> for PyFromSwarm {
+    fn from(value: FromSwarm) -> Self {
+        match value {
+            FromSwarm::Discovered {} => Self::Connection { connected: true },
+            FromSwarm::Expired {} => Self::Connection { connected: false },
+            FromSwarm::Message { topic, data } => Self::Message {
+                topic: topic,
+                data: data.pybytes(),
+            },
+        }
+    }
+}
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl PyNetworkingHandle {
+    // 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]
+    fn new<'py>(
+        identity: Bound<'py, PyBytes>,
+        bootstrap_peers: Vec<String>,
+        listen_port: u16,
+    ) -> PyResult<(PyNetworkingHandle, PySession)> {
+        // create communication channels
+        let (to_swarm, from_client) = mpsc::channel(1024);
+
+        // get identity
+        let identity = u128::from_le_bytes(
+            identity
+                .extract::<'_, Vec<u8>>()?
+                .try_into()
+                .map_err(|_| PyValueError::new_err("invalid identity bytes"))?,
+        );
+
+        // create networking swarm (within tokio context!! or it crashes)
+        let swarm = pyo3_async_runtimes::tokio::get_runtime()
+            .block_on(create_swarm(
+                identity,
+                from_client,
+                bootstrap_peers,
+                listen_port,
+            ))
+            .pyerr()?;
+
+        let session = swarm.session.z.clone();
+
+        Ok((
+            PyNetworkingHandle {
+                swarm: Arc::new(Mutex::new(swarm.into_stream())),
+                to_swarm,
+            },
+            PySession { session },
+        ))
+    }
+
+    #[gen_stub(skip)]
+    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 {
+            swarm
+                .try_lock()
+                .map_err(|_| PyRuntimeError::new_err("called recv twice concurrently"))?
+                .next()
+                .await
+                .ok_or(PyErr::receiver_channel_closed())
+                .map(PyFromSwarm::from)
+        })
+    }
+
+    // ---- Gossipsub management methods ----
+
+    /// Subscribe to a `GossipSub` topic.
+    ///
+    /// Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
+    async fn gossipsub_subscribe(&self, topic: String) -> PyResult<bool> {
+        let (tx, rx) = oneshot::channel();
+
+        // send off request to subscribe
+        self.to_swarm
+            .send_py(ToSwarm::Subscribe {
+                topic,
+                result_sender: tx,
+            })
+            .allow_threads_py() // allow-threads-aware async call
+            .await?;
+
+        // wait for response & return any errors
+        rx.allow_threads_py() // allow-threads-aware async call
+            .await
+            .map_err(|_| PyErr::receiver_channel_closed())?
+            .pyerr()
+    }
+
+    /// Unsubscribes from a `GossipSub` topic.
+    ///
+    /// Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
+    async fn gossipsub_unsubscribe(&self, topic: String) -> PyResult<bool> {
+        let (tx, rx) = oneshot::channel();
+
+        // send off request to unsubscribe
+        self.to_swarm
+            .send_py(ToSwarm::Unsubscribe {
+                topic,
+                result_sender: tx,
+            })
+            .allow_threads_py() // allow-threads-aware async call
+            .await?;
+
+        // wait for response & convert any errors
+        rx.allow_threads_py() // allow-threads-aware async call
+            .await
+            .map_err(|_| PyErr::receiver_channel_closed())
+    }
+
+    /// Publishes a message with multiple topics to the `GossipSub` network.
+    ///
+    /// If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
+    async fn gossipsub_publish(&self, topic: String, data: Py<PyBytes>) -> PyResult<()> {
+        let (tx, rx) = oneshot::channel();
+
+        // send off request to subscribe
+        let data = Python::attach(|py| Vec::from(data.as_bytes(py)));
+        self.to_swarm
+            .send_py(ToSwarm::Publish {
+                topic,
+                data,
+                result_sender: tx,
+            })
+            .allow_threads_py() // allow-threads-aware async call
+            .await?;
+
+        // wait for response & return any errors => ignore messageID for now!!!
+        let _ = rx
+            .allow_threads_py() // allow-threads-aware async call
+            .await
+            .map_err(|_| PyErr::receiver_channel_closed())?
+            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+        Ok(())
+    }
+}
+
+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>()?;
+
+    Ok(())
+}

rust/exo_net/src/pidfile.rs

@@ -0,0 +1,87 @@
+use pidfile_rs::{Pidfile, PidfileError};
+use pyo3::exceptions::PyException;
+use pyo3::prelude::{PyModule, PyModuleMethods};
+use pyo3::{Bound, PyErr, PyResult, Python, pyclass, pymethods};
+use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
+use std::fs::Permissions;
+use std::os::unix::prelude::PermissionsExt;
+use std::path::PathBuf;
+
+#[gen_stub_pyclass]
+#[pyclass(frozen, extends=PyException, name="PidfileError")]
+pub struct PyPidfileError(PidfileError);
+
+impl PyPidfileError {
+    // TODO: I actually like this pattern a LOT more but how to abstract??
+    fn into_pyerr(self, py: Python) -> PyErr {
+        match Bound::new(py, self) {
+            Ok(err) => PyErr::from_value(err.into_any()),
+            Err(err) => err,
+        }
+    }
+}
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl PyPidfileError {
+    fn __repr__(&self) -> String {
+        format!("PidfileError(\"{}\")", self.0)
+    }
+
+    fn __str__(&self) -> String {
+        self.0.to_string()
+    }
+}
+
+/// A PID file protected with a lock.
+///
+/// An instance of `Pidfile` can be used to manage a PID file: create it,
+/// lock it, detect already running daemons. It is backed by [`pidfile`][]
+/// functions of `libbsd`/`libutil` which use `flopen` to lock the PID
+/// file.
+///
+/// When a PID file is created, the process ID of the current process is
+/// *not* written there, making it possible to lock the PID file before
+/// forking and only write the ID of the forked process when it is ready.
+///
+/// The PID file is deleted automatically when the `Pidfile` comes out of
+/// the scope. To close the PID file without deleting it, for example, in
+/// the parent process of a forked daemon, call `close()`.
+///
+/// [`exit`]: https://doc.rust-lang.org/std/process/fn.exit.html
+/// [`pidfile`]: https://linux.die.net/man/3/pidfile
+/// [`daemon`(3)]: https://linux.die.net/man/3/daemon
+#[gen_stub_pyclass]
+#[pyclass(name = "Pidfile")]
+pub struct PyPidfile(Pidfile);
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl PyPidfile {
+    /// Creates a new PID file and locks it.
+    ///
+    /// 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
+    /// the PID file yet.
+    #[new]
+    fn py_new(py: Python, path: PathBuf, mode: u32) -> PyResult<Self> {
+        Ok(Self(
+            Pidfile::new(&path, Permissions::from_mode(mode))
+                .map_err(|e| PyPidfileError(e).into_pyerr(py))?,
+        ))
+    }
+
+    /// Writes the current process ID to the PID file.
+    ///
+    /// The file is truncated before writing.
+    fn write<'py>(&mut self, py: Python<'py>) -> PyResult<()> {
+        self.0.write().map_err(|e| PyPidfileError(e).into_pyerr(py))
+    }
+}
+
+pub fn pidfile_submodule(m: &Bound<PyModule>) -> PyResult<()> {
+    m.add_class::<PyPidfileError>()?;
+    m.add_class::<PyPidfile>()?;
+
+    Ok(())
+}

rust/exo_net/src/point_to_point.rs

@@ -0,0 +1,108 @@
+use std::sync::Arc;
+
+use pyo3::exceptions::PyConnectionError;
+use pyo3::types::PyBytes;
+use pyo3::types::PyNone;
+use pyo3::{BoundObject, prelude::*};
+use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
+use zenoh::Result;
+use zenoh::{
+    handlers::FifoChannelHandler,
+    pubsub::{Publisher, Subscriber},
+    sample::Sample,
+};
+
+use crate::ext::ByteArrayExt;
+
+#[gen_stub_pyclass]
+#[pyclass]
+pub struct NetReceiver {
+    pub subscriber: Subscriber<FifoChannelHandler<Sample>>,
+}
+#[gen_stub_pymethods]
+#[pymethods]
+impl NetReceiver {
+    #[gen_stub(override_return_type(
+        type_repr="collections.abc.Awaitable[bytes | 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, {
+            assert!(
+                self.subscriber.receiver_count() == 1,
+                "tried to receive twice on the same receiver"
+            );
+            let subscriber = self.subscriber.clone();
+            async move {
+                match subscriber.recv_async().await {
+                    Err(_) => {
+                        // stream closed;
+                        Ok(Python::attach(|py| PyNone::get(py).unbind()).into_any())
+                    }
+                    Ok(sample) => Ok(sample.payload().to_bytes().to_vec().pybytes().into_any()),
+                }
+            }
+        })
+    }
+}
+
+#[gen_stub_pyclass]
+#[pyclass]
+pub struct NetSender {
+    pub publisher: Arc<Publisher<'static>>,
+    pub first: bool,
+}
+#[gen_stub_pymethods]
+#[pymethods]
+impl NetSender {
+    #[gen_stub(override_return_type(
+        type_repr="collections.abc.Awaitable[bool]",
+        imports=("collections.abc")
+    ))]
+    pub fn send<'py>(
+        &'py mut self,
+        py: Python<'py>,
+        data: Bound<'py, PyBytes>,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        let is_first = self.first;
+        self.first = false;
+        pyo3_async_runtimes::tokio::future_into_py(py, {
+            let publisher = Arc::clone(&self.publisher);
+            // clone the data so py can have it back
+            let bytes = data.as_bytes().to_vec();
+            async move {
+                if is_first {
+                    wait_for_listener(&*publisher)
+                        .await
+                        .map_err(|e| PyConnectionError::new_err(e.to_string()))?;
+                }
+                if !publisher
+                    .matching_status()
+                    .await
+                    .map_err(|e| PyConnectionError::new_err(e.to_string()))?
+                    .matching()
+                {
+                    return Ok(false);
+                }
+                publisher
+                    .put(&bytes)
+                    .await
+                    .map_err(|e| PyConnectionError::new_err(e.to_string()))?;
+                Ok(true)
+            }
+        })
+    }
+}
+
+async fn wait_for_listener<'a>(publisher: &Publisher<'a>) -> Result<()> {
+    let matcher = publisher.matching_listener().await?;
+    if publisher.matching_status().await?.matching() {
+        return Ok(());
+    }
+    while let Ok(status) = matcher.recv_async().await {
+        if status.matching() {
+            break;
+        }
+    }
+    Ok(())
+}

rust/exo_net/src/session.rs

@@ -0,0 +1,72 @@
+use std::sync::Arc;
+
+use pyo3::{exceptions::PyValueError, prelude::*};
+use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
+
+use zenoh::Session;
+use zenoh::Wait;
+use zenoh::qos::CongestionControl;
+
+use crate::{
+    point_to_point::{NetReceiver, NetSender},
+    state::StateProxy,
+};
+
+#[gen_stub_pyclass]
+#[pyclass]
+pub struct PySession {
+    pub session: Session,
+}
+#[gen_stub_pymethods]
+#[pymethods]
+impl PySession {
+    /* for now construct with NetworkingHandle
+    #[staticmethod]
+    #[gen_stub(override_return_type(
+        type_repr="collections.abc.Awaitable[PySession]",
+        imports=("collections.abc")
+    ))]
+    pub fn init<'py>(py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+            Ok(Self {
+                session: networking::open(
+                    networking::cfg(rand::random(), 0).expect("default cfg is valid"),
+                )
+                .await
+                .map_err(|e| PyRuntimeError::new_err(e.to_string()))?,
+            })
+        })
+    }
+    */
+
+    pub fn net_receiver<'py>(&self, key: String) -> PyResult<NetReceiver> {
+        Ok(NetReceiver {
+            subscriber: self
+                .session
+                .declare_subscriber(key)
+                .wait()
+                // C5: key format error
+                .map_err(|e| PyValueError::new_err(e.to_string()))?,
+        })
+    }
+
+    pub fn net_sender<'py>(&self, key: String) -> PyResult<NetSender> {
+        Ok(NetSender {
+            publisher: Arc::new(
+                self.session
+                    .declare_publisher(key)
+                    .congestion_control(CongestionControl::Block)
+                    .wait()
+                    // C5: key format error, could be declaration error
+                    .map_err(|e| PyValueError::new_err(e.to_string()))?,
+            ),
+            first: true,
+        })
+    }
+
+    pub fn state_proxy(&self) -> StateProxy {
+        StateProxy {
+            session: self.session.clone(),
+        }
+    }
+}

rust/exo_net/src/state.rs

@@ -0,0 +1,78 @@
+use pyo3::{exceptions::PyValueError, prelude::*};
+use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
+use serde_json::{Map, Value};
+use zenoh::{Result, Session, sample::SampleFields};
+
+#[gen_stub_pyclass]
+#[pyclass]
+pub struct StateProxy {
+    pub session: Session,
+}
+#[gen_stub_pymethods]
+#[pymethods]
+impl StateProxy {
+    #[gen_stub(override_return_type(
+        type_repr="collections.abc.Awaitable[str]",
+        imports=("collections.abc")
+    ))]
+    pub fn snapshot<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        pyo3_async_runtimes::tokio::future_into_py(py, {
+            let session = self.session.clone();
+            async move {
+                Self::_snapshot(session)
+                    .await
+                    .map_err(|e| PyValueError::new_err(e.to_string()))
+                    .map(|v| v.to_string())
+            }
+        })
+    }
+}
+
+impl StateProxy {
+    async fn _snapshot(session: Session) -> Result<Value> {
+        let q = session.get("storage/mem1/**").await?;
+
+        let mut v = Value::Object(Map::default());
+
+        while let Ok(sample) = q.recv_async().await {
+            let mut cur_v = &mut v;
+            let Ok(sample) = sample.into_result() else {
+                continue;
+            };
+            // skip storage/mem1
+            let SampleFields {
+                payload, key_expr, ..
+            } = sample.into();
+            let mut iter = key_expr.split('/').skip(2).peekable();
+            loop {
+                let Some(p) = iter.next() else {
+                    break;
+                };
+                if iter.peek().is_none() {
+                    // terminal; write value into json
+                    let existing = cur_v
+                        .as_object_mut()
+                        .expect("path terminated unexpectedly - value stored at some/path and some/path/two")
+                        .insert(p.to_owned(), Value::String(payload.try_to_string()?.to_string()));
+
+                    if let Some(value) = existing {
+                        assert!(value.is_string())
+                        // could log, but string overwrites are fine
+                    }
+                } else {
+                    // non-terminal; ensure key exists in v, then replace cur with that object
+                    cur_v = cur_v
+                        .as_object_mut()
+                        .expect("path terminated unexpectedly - value stored at some/path and some/path/two")
+                        .entry(p)
+                        .or_insert(Value::Object(Map::default()));
+                    assert!(
+                        cur_v.is_object(),
+                        "path terminated unexpectedly - value stored at some/path and some/path/two"
+                    )
+                }
+            }
+        }
+        Ok(v)
+    }
+}

rust/exo_net/tests/test_python.py

@@ -1,10 +1,11 @@
 import asyncio
 
 import pytest
+from _pytest.capture import CaptureFixture
 from exo_pyo3_bindings import (
     Keypair,
     NetworkingHandle,
-    NoPeersSubscribedToTopicError,
+    Pidfile,
     PyFromSwarm,
 )
 
@@ -13,17 +14,22 @@ from exo_pyo3_bindings import (
 async def test_sleep_on_multiple_items() -> None:
     print("PYTHON: starting handle")
     h = NetworkingHandle(Keypair.generate(), [], 0)
+    print("PYTHON: handle started")
 
     rt = asyncio.create_task(_await_recv(h))
 
     # sleep for 4 ticks
-    for i in range(4):
+    for i in range(10):
         await asyncio.sleep(1)
 
-        try:
-            await h.gossipsub_publish("topic", b"somehting or other")
-        except NoPeersSubscribedToTopicError as e:
-            print("caught it", e)
+        await h.gossipsub_publish("topic", b"somehting or other")
+
+
+def test_pidfile(capsys: CaptureFixture[str]):
+    with capsys.disabled():
+        print("\nbefore python")
+        scoped_lock_file()
+        print("after python")
 
 
 async def _await_recv(h: NetworkingHandle):
@@ -34,3 +40,11 @@ async def _await_recv(h: NetworkingHandle):
                 print(f"PYTHON: connection update: {c}")
             case PyFromSwarm.Message() as m:
                 print(f"PYTHON: message: {m}")
+
+
+def scoped_lock_file():
+    a = Pidfile("/tmp/lock.pid", 0o0600)
+
+
+if __name__ == "__main__":
+    asyncio.run(test_sleep_on_multiple_items())

rust/exo_pyo3_bindings/Cargo.toml

@@ -1,66 +0,0 @@
-[package]
-name = "exo_pyo3_bindings"
-version = { workspace = true }
-edition = { workspace = true }
-publish = false
-
-[lib]
-doctest = false
-path = "src/lib.rs"
-name = "exo_pyo3_bindings"
-
-# "cdylib" needed to produce shared library for Python to import
-# "rlib" needed for stub-gen to run
-crate-type = ["cdylib", "rlib"]
-
-[[bin]]
-path = "src/bin/stub_gen.rs"
-name = "stub_gen"
-doc = false
-
-[lints]
-workspace = true
-
-[dependencies]
-networking = { workspace = true }
-
-# interop
-pyo3 = { version = "0.27.2", features = [
-  # "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13
-  # "nightly", # enables better-supported GIL integration
-  "experimental-async", # async support in #[pyfunction] & #[pymethods]
-  #"experimental-inspect", # inspection of generated binary => easier to automate type-hint generation
-  #"py-clone", # adding Clone-ing of `Py<T>` without GIL (may cause panics - remove if panics happen)
-  # "multiple-pymethods", # allows multiple #[pymethods] sections per class
-
-  # integrations with other libraries
-  # "arc_lock", "bigdecimal", "either", "hashbrown", "indexmap", "num-bigint", "num-complex", "num-rational",
-  # "ordered-float", "rust_decimal", "smallvec",
-  # "anyhow", "chrono", "chrono-local", "chrono-tz", "eyre", "jiff-02", "lock_api", "parking-lot", "time",  "serde",
-] }
-pyo3-stub-gen = { version = "0.17.2" }
-pyo3-async-runtimes = { version = "0.27.0", features = [
-  "attributes",
-  "tokio-runtime",
-  "testing",
-] }
-pyo3-log = "0.13.2"
-
-# macro dependencies
-extend = { workspace = true }
-delegate = { workspace = true }
-
-# async runtime
-tokio = { workspace = true, features = ["full", "tracing"] }
-futures-lite = { workspace = true }
-
-# utility dependencies
-util = { workspace = true }
-
-# Tracing
-log = { workspace = true }
-env_logger = "0.11"
-
-# Networking
-libp2p = { workspace = true, features = ["full"] }
-pin-project = "1.1.10"

rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi

@@ -1,94 +0,0 @@
-# This file is automatically generated by pyo3_stub_gen
-# ruff: noqa: E501, F401
-
-import builtins
-import typing
-
-@typing.final
-class AllQueuesFullError(builtins.Exception):
-    def __new__(cls, *args: typing.Any) -> AllQueuesFullError: ...
-    def __repr__(self) -> builtins.str: ...
-    def __str__(self) -> builtins.str: ...
-
-@typing.final
-class Keypair:
-    r"""
-    Identity keypair of a node.
-    """
-    @staticmethod
-    def generate() -> Keypair:
-        r"""
-        Generate a new Ed25519 keypair.
-        """
-    @staticmethod
-    def from_bytes(bytes: bytes) -> Keypair:
-        r"""
-        Construct an Ed25519 keypair from secret key bytes
-        """
-    def to_bytes(self) -> bytes:
-        r"""
-        Get the secret key bytes underlying the keypair
-        """
-    def to_node_id(self) -> builtins.str:
-        r"""
-        Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
-        """
-
-@typing.final
-class MessageTooLargeError(builtins.Exception):
-    def __new__(cls, *args: typing.Any) -> MessageTooLargeError: ...
-    def __repr__(self) -> builtins.str: ...
-    def __str__(self) -> builtins.str: ...
-
-@typing.final
-class NetworkingHandle:
-    def __new__(cls, identity: Keypair, bootstrap_peers: typing.Sequence[builtins.str], listen_port: builtins.int) -> NetworkingHandle: ...
-    async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
-        r"""
-        Subscribe to a `GossipSub` topic.
-        
-        Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
-        """
-    async def gossipsub_unsubscribe(self, topic: builtins.str) -> builtins.bool:
-        r"""
-        Unsubscribes from a `GossipSub` topic.
-        
-        Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
-        """
-    async def gossipsub_publish(self, topic: builtins.str, data: bytes) -> None:
-        r"""
-        Publishes a message with multiple topics to the `GossipSub` network.
-        
-        If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
-        """
-    async def recv(self) -> PyFromSwarm: ...
-
-@typing.final
-class NoPeersSubscribedToTopicError(builtins.Exception):
-    def __new__(cls, *args: typing.Any) -> NoPeersSubscribedToTopicError: ...
-    def __repr__(self) -> builtins.str: ...
-    def __str__(self) -> builtins.str: ...
-
-class PyFromSwarm:
-    @typing.final
-    class Connection(PyFromSwarm):
-        __match_args__ = ("peer_id", "connected",)
-        @property
-        def peer_id(self) -> builtins.str: ...
-        @property
-        def connected(self) -> builtins.bool: ...
-        def __new__(cls, peer_id: builtins.str, connected: builtins.bool) -> PyFromSwarm.Connection: ...
-    
-    @typing.final
-    class Message(PyFromSwarm):
-        __match_args__ = ("origin", "topic", "data",)
-        @property
-        def origin(self) -> builtins.str: ...
-        @property
-        def topic(self) -> builtins.str: ...
-        @property
-        def data(self) -> bytes: ...
-        def __new__(cls, origin: builtins.str, topic: builtins.str, data: bytes) -> PyFromSwarm.Message: ...
-    
-    ...
-

rust/exo_pyo3_bindings/src/networking.rs

@@ -1,318 +0,0 @@
-use std::pin::Pin;
-use std::sync::Arc;
-
-use crate::r#const::MPSC_CHANNEL_SIZE;
-use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
-use crate::ext::{ResultExt as _, TokioMpscSenderExt as _};
-use crate::ident::PyKeypair;
-use crate::networking::exception::{
-    PyAllQueuesFullError, PyMessageTooLargeError, PyNoPeersSubscribedToTopicError,
-};
-use crate::pyclass;
-use futures_lite::{Stream, StreamExt as _};
-use libp2p::gossipsub::PublishError;
-use networking::swarm::{FromSwarm, ToSwarm, create_swarm};
-use pyo3::exceptions::PyRuntimeError;
-use pyo3::prelude::{PyModule, PyModuleMethods as _};
-use pyo3::types::PyBytes;
-use pyo3::{Bound, Py, PyAny, PyErr, PyResult, Python, 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};
-
-mod exception {
-    use pyo3::types::PyTuple;
-    use pyo3::{exceptions::PyException, prelude::*};
-    use pyo3_stub_gen::derive::*;
-
-    #[gen_stub_pyclass]
-    #[pyclass(frozen, extends=PyException, name="NoPeersSubscribedToTopicError")]
-    pub struct PyNoPeersSubscribedToTopicError {}
-
-    impl PyNoPeersSubscribedToTopicError {
-        const MSG: &'static str = "\
-        No peers are currently subscribed to receive messages on this topic. \
-        Wait for peers to subscribe or check your network connectivity.";
-
-        ///   Creates a new  [ `PyErr` ]  of this type.
-        ///
-        ///   [`PyErr`] :  https://docs.rs/pyo3/latest/pyo3/struct.PyErr.html   "PyErr in pyo3"
-        pub(crate) fn new_err() -> PyErr {
-            PyErr::new::<Self, _>(()) // TODO: check if this needs to be replaced???
-        }
-    }
-
-    #[gen_stub_pymethods]
-    #[pymethods]
-    impl PyNoPeersSubscribedToTopicError {
-        #[new]
-        #[pyo3(signature = (*args))]
-        #[allow(unused_variables)]
-        pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self {
-            Self {}
-        }
-
-        fn __repr__(&self) -> String {
-            format!("PeerId(\"{}\")", Self::MSG)
-        }
-
-        fn __str__(&self) -> String {
-            Self::MSG.to_string()
-        }
-    }
-
-    #[gen_stub_pyclass]
-    #[pyclass(frozen, extends=PyException, name="AllQueuesFullError")]
-    pub struct PyAllQueuesFullError {}
-
-    impl PyAllQueuesFullError {
-        const MSG: &'static str =
-            "All libp2p peers are unresponsive, resend the message or reconnect.";
-
-        ///   Creates a new  [ `PyErr` ]  of this type.
-        ///
-        ///   [`PyErr`] :  https://docs.rs/pyo3/latest/pyo3/struct.PyErr.html   "PyErr in pyo3"
-        pub(crate) fn new_err() -> PyErr {
-            PyErr::new::<Self, _>(()) // TODO: check if this needs to be replaced???
-        }
-    }
-
-    #[gen_stub_pymethods]
-    #[pymethods]
-    impl PyAllQueuesFullError {
-        #[new]
-        #[pyo3(signature = (*args))]
-        #[allow(unused_variables)]
-        pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self {
-            Self {}
-        }
-
-        fn __repr__(&self) -> String {
-            format!("PeerId(\"{}\")", Self::MSG)
-        }
-
-        fn __str__(&self) -> String {
-            Self::MSG.to_string()
-        }
-    }
-
-    #[gen_stub_pyclass]
-    #[pyclass(frozen, extends=PyException, name="MessageTooLargeError")]
-    pub struct PyMessageTooLargeError {}
-
-    impl PyMessageTooLargeError {
-        const MSG: &'static str = "Gossipsub message exceeds max_transmit_size. Reduce prompt length or increase the limit.";
-
-        pub(crate) fn new_err() -> PyErr {
-            PyErr::new::<Self, _>(())
-        }
-    }
-
-    #[gen_stub_pymethods]
-    #[pymethods]
-    impl PyMessageTooLargeError {
-        #[new]
-        #[pyo3(signature = (*args))]
-        #[allow(unused_variables)]
-        pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self {
-            Self {}
-        }
-
-        fn __repr__(&self) -> String {
-            format!("MessageTooLargeError(\"{}\")", Self::MSG)
-        }
-
-        fn __str__(&self) -> String {
-            Self::MSG.to_string()
-        }
-    }
-}
-
-#[gen_stub_pyclass]
-#[pyclass(name = "NetworkingHandle")]
-struct PyNetworkingHandle {
-    // channels
-    pub to_swarm: mpsc::Sender<ToSwarm>,
-    pub swarm: Arc<Mutex<Pin<Box<dyn Stream<Item = FromSwarm> + Send>>>>,
-}
-
-#[gen_stub_pyclass_complex_enum]
-#[pyclass]
-enum PyFromSwarm {
-    Connection {
-        peer_id: String,
-        connected: bool,
-    },
-    Message {
-        origin: String,
-        topic: String,
-        data: Py<PyBytes>,
-    },
-}
-impl From<FromSwarm> for PyFromSwarm {
-    fn from(value: FromSwarm) -> Self {
-        match value {
-            FromSwarm::Discovered { peer_id } => Self::Connection {
-                peer_id: peer_id.to_base58(),
-                connected: true,
-            },
-            FromSwarm::Expired { peer_id } => Self::Connection {
-                peer_id: peer_id.to_base58(),
-                connected: false,
-            },
-            FromSwarm::Message { from, topic, data } => Self::Message {
-                origin: from.to_base58(),
-                topic: topic,
-                data: data.pybytes(),
-            },
-        }
-    }
-}
-
-#[gen_stub_pymethods]
-#[pymethods]
-impl PyNetworkingHandle {
-    // 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 ----
-
-    #[new]
-    #[pyo3(signature = (identity, bootstrap_peers, listen_port))]
-    fn py_new(
-        identity: Bound<'_, PyKeypair>,
-        bootstrap_peers: Vec<String>,
-        listen_port: u16,
-    ) -> PyResult<Self> {
-        // create communication channels
-        let (to_swarm, from_client) = mpsc::channel(MPSC_CHANNEL_SIZE);
-
-        // get identity
-        let identity = identity.borrow().0.clone();
-
-        // create networking swarm (within tokio context!! or it crashes)
-        let _guard = pyo3_async_runtimes::tokio::get_runtime().enter();
-        let swarm = create_swarm(identity, from_client, bootstrap_peers, listen_port)
-            .pyerr()?
-            .into_stream();
-
-        Ok(Self {
-            swarm: Arc::new(Mutex::new(swarm)),
-            to_swarm,
-        })
-    }
-
-    #[gen_stub(skip)]
-    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 {
-            swarm
-                .try_lock()
-                .map_err(|_| PyRuntimeError::new_err("called recv twice concurrently"))?
-                .next()
-                .await
-                .ok_or(PyErr::receiver_channel_closed())
-                .map(PyFromSwarm::from)
-        })
-    }
-
-    // ---- Gossipsub management methods ----
-
-    /// Subscribe to a `GossipSub` topic.
-    ///
-    /// Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
-    async fn gossipsub_subscribe(&self, topic: String) -> PyResult<bool> {
-        let (tx, rx) = oneshot::channel();
-
-        // send off request to subscribe
-        self.to_swarm
-            .send_py(ToSwarm::Subscribe {
-                topic,
-                result_sender: tx,
-            })
-            .allow_threads_py() // allow-threads-aware async call
-            .await?;
-
-        // wait for response & return any errors
-        rx.allow_threads_py() // allow-threads-aware async call
-            .await
-            .map_err(|_| PyErr::receiver_channel_closed())?
-            .pyerr()
-    }
-
-    /// Unsubscribes from a `GossipSub` topic.
-    ///
-    /// Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
-    async fn gossipsub_unsubscribe(&self, topic: String) -> PyResult<bool> {
-        let (tx, rx) = oneshot::channel();
-
-        // send off request to unsubscribe
-        self.to_swarm
-            .send_py(ToSwarm::Unsubscribe {
-                topic,
-                result_sender: tx,
-            })
-            .allow_threads_py() // allow-threads-aware async call
-            .await?;
-
-        // wait for response & convert any errors
-        rx.allow_threads_py() // allow-threads-aware async call
-            .await
-            .map_err(|_| PyErr::receiver_channel_closed())
-    }
-
-    /// Publishes a message with multiple topics to the `GossipSub` network.
-    ///
-    /// If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
-    async fn gossipsub_publish(&self, topic: String, data: Py<PyBytes>) -> PyResult<()> {
-        let (tx, rx) = oneshot::channel();
-
-        // send off request to subscribe
-        let data = Python::attach(|py| Vec::from(data.as_bytes(py)));
-        self.to_swarm
-            .send_py(ToSwarm::Publish {
-                topic,
-                data,
-                result_sender: tx,
-            })
-            .allow_threads_py() // allow-threads-aware async call
-            .await?;
-
-        // wait for response & return any errors => ignore messageID for now!!!
-        let _ = rx
-            .allow_threads_py() // allow-threads-aware async call
-            .await
-            .map_err(|_| PyErr::receiver_channel_closed())?
-            .map_err(|e| match e {
-                PublishError::AllQueuesFull(_) => PyAllQueuesFullError::new_err(),
-                PublishError::MessageTooLarge => PyMessageTooLargeError::new_err(),
-                PublishError::NoPeersSubscribedToTopic => {
-                    PyNoPeersSubscribedToTopicError::new_err()
-                }
-                e => PyRuntimeError::new_err(e.to_string()),
-            })?;
-        Ok(())
-    }
-}
-
-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::<exception::PyNoPeersSubscribedToTopicError>()?;
-    m.add_class::<exception::PyAllQueuesFullError>()?;
-    m.add_class::<exception::PyMessageTooLargeError>()?;
-
-    m.add_class::<PyNetworkingHandle>()?;
-    m.add_class::<PyFromSwarm>()?;
-
-    Ok(())
-}

rust/networking/Cargo.toml

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

rust/networking/examples/chatroom.rs

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

rust/networking/examples/put_string.rs

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

rust/networking/examples/serve_storage.rs

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

rust/networking/src/RESEARCH_NOTES.txt

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

rust/networking/src/discovery.rs

@@ -1,390 +1,311 @@
-use crate::ext::MultiaddrExt;
-use delegate::delegate;
-use either::Either;
-use futures_lite::FutureExt;
-use futures_timer::Delay;
-use libp2p::core::transport::PortUse;
-use libp2p::core::{ConnectedPoint, Endpoint};
-use libp2p::swarm::behaviour::ConnectionEstablished;
-use libp2p::swarm::dial_opts::DialOpts;
-use libp2p::swarm::{
-    CloseConnection, ConnectionClosed, ConnectionDenied, ConnectionHandler,
-    ConnectionHandlerSelect, ConnectionId, FromSwarm, NetworkBehaviour, THandler, THandlerInEvent,
-    THandlerOutEvent, ToSwarm, dummy,
+use std::{
+    io::{self, ErrorKind},
+    net::{Ipv6Addr, SocketAddr, SocketAddrV6},
+    sync::Arc,
+    time::Duration,
 };
-use libp2p::{Multiaddr, PeerId, identity, mdns};
-use std::collections::{BTreeSet, HashMap};
-use std::convert::Infallible;
-use std::io;
-use std::net::IpAddr;
-use std::task::{Context, Poll};
-use std::time::Duration;
-use util::wakerdeque::WakerDeque;
 
-const RETRY_CONNECT_INTERVAL: Duration = Duration::from_secs(5);
+use bytemuck::{Pod, Zeroable};
+use log::{debug, trace, warn};
+use netwatcher::WatchHandle;
+use parking_lot::Mutex;
+use tokio::{
+    net::UdpSocket,
+    time::{Interval, interval},
+};
+use zenoh::config::ZenohId;
 
-mod managed {
-    use libp2p::swarm::NetworkBehaviour;
-    use libp2p::{identity, mdns, ping};
-    use std::io;
-    use std::time::Duration;
+const GROUP: Ipv6Addr = Ipv6Addr::new(0xff12, 0, 0, 0, 0, 0, 0xe0a1, 0xde89);
 
-    const MDNS_RECORD_TTL: Duration = Duration::from_secs(2_500);
-    const MDNS_QUERY_INTERVAL: Duration = Duration::from_secs(1_500);
-    const PING_TIMEOUT: Duration = Duration::from_millis(2_500);
-    const PING_INTERVAL: Duration = Duration::from_millis(2_500);
-
-    #[derive(NetworkBehaviour)]
-    pub struct Behaviour {
-        mdns: mdns::tokio::Behaviour,
-        ping: ping::Behaviour,
-    }
-
-    impl Behaviour {
-        pub fn new(keypair: &identity::Keypair) -> io::Result<Self> {
-            Ok(Self {
-                mdns: mdns_behaviour(keypair)?,
-                ping: ping_behaviour(),
-            })
-        }
-    }
-
-    fn mdns_behaviour(keypair: &identity::Keypair) -> io::Result<mdns::tokio::Behaviour> {
-        use mdns::{Config, tokio};
-
-        // mDNS config => enable IPv6
-        let mdns_config = Config {
-            ttl: MDNS_RECORD_TTL,
-            query_interval: MDNS_QUERY_INTERVAL,
-
-            // enable_ipv6: true, // TODO: for some reason, TCP+mDNS don't work well with ipv6?? figure out how to make work
-            ..Default::default()
-        };
-
-        let mdns_behaviour = tokio::Behaviour::new(mdns_config, keypair.public().to_peer_id());
-        Ok(mdns_behaviour?)
-    }
-
-    fn ping_behaviour() -> ping::Behaviour {
-        ping::Behaviour::new(
-            ping::Config::new()
-                .with_timeout(PING_TIMEOUT)
-                .with_interval(PING_INTERVAL),
-        )
-    }
+pub struct Discovery {
+    sock: Arc<UdpSocket>,
+    ifaces: Arc<Mutex<Vec<SocketAddr>>>,
+    last_nonce: Mutex<[u8; 8]>,
+    /// the port of the service we are doing discovery for - transmitted to peers
+    listen_port: u16,
+    zid: ZenohId,
+    tick: Interval,
+    _sync: Mutex<WatchHandle>,
 }
 
-/// Events for when a listening connection is truly established and truly closed.
-#[derive(Debug, Clone)]
-pub enum Event {
-    ConnectionEstablished {
-        peer_id: PeerId,
-        connection_id: ConnectionId,
-        remote_ip: IpAddr,
-        remote_tcp_port: u16,
-    },
-    ConnectionClosed {
-        peer_id: PeerId,
-        connection_id: ConnectionId,
-        remote_ip: IpAddr,
-        remote_tcp_port: u16,
-    },
-}
-
-/// Discovery behavior that wraps mDNS to produce truly discovered durable peer-connections.
-///
-/// The behaviour operates as such:
-///  1) All true (listening) connections/disconnections are tracked, emitting corresponding events
-///     to the swarm.
-///  1) mDNS discovered/expired peers are tracked; discovered but not connected peers are dialed
-///     immediately, and expired but connected peers are disconnected from immediately.
-///  2) Every fixed interval: discovered but not connected peers are dialed, and expired but
-///     connected peers are disconnected from.
-pub struct Behaviour {
-    // state-tracking for managed behaviors & mDNS-discovered peers
-    managed: managed::Behaviour,
-    mdns_discovered: HashMap<PeerId, BTreeSet<Multiaddr>>,
-    bootstrap_peers: Vec<Multiaddr>,
-
-    retry_delay: Delay, // retry interval
-
-    // pending events to emmit => waker-backed Deque to control polling
-    pending_events: WakerDeque<ToSwarm<Event, Infallible>>,
-}
-
-impl Behaviour {
-    pub fn new(keypair: &identity::Keypair, bootstrap_peers: Vec<Multiaddr>) -> io::Result<Self> {
-        Ok(Self {
-            managed: managed::Behaviour::new(keypair)?,
-            mdns_discovered: HashMap::new(),
-            bootstrap_peers,
-            retry_delay: Delay::new(RETRY_CONNECT_INTERVAL),
-            pending_events: WakerDeque::new(),
-        })
-    }
-
-    fn dial(&mut self, peer_id: PeerId, addr: Multiaddr) {
-        self.pending_events.push_back(ToSwarm::Dial {
-            opts: DialOpts::peer_id(peer_id).addresses(vec![addr]).build(),
-        })
-    }
-
-    fn close_connection(&mut self, peer_id: PeerId, connection: ConnectionId) {
-        // push front to make this IMMEDIATE
-        self.pending_events.push_front(ToSwarm::CloseConnection {
-            peer_id,
-            connection: CloseConnection::One(connection),
-        })
-    }
-
-    fn handle_mdns_discovered(&mut self, peers: Vec<(PeerId, Multiaddr)>) {
-        for (p, ma) in peers {
-            self.dial(p, ma.clone()); // always connect
-
-            // get peer's multi-addresses or insert if missing
-            let Some(mas) = self.mdns_discovered.get_mut(&p) else {
-                self.mdns_discovered.insert(p, BTreeSet::from([ma]));
-                continue;
-            };
-
-            // multiaddress should never already be present - else something has gone wrong
-            let is_new_addr = mas.insert(ma);
-            assert!(is_new_addr, "cannot discover a discovered peer");
-        }
-    }
-
-    fn handle_mdns_expired(&mut self, peers: Vec<(PeerId, Multiaddr)>) {
-        for (p, ma) in peers {
-            // at this point, we *must* have the peer
-            let mas = self
-                .mdns_discovered
-                .get_mut(&p)
-                .expect("nonexistent peer cannot expire");
-
-            // at this point, we *must* have the multiaddress
-            let was_present = mas.remove(&ma);
-            assert!(was_present, "nonexistent multiaddress cannot expire");
-
-            // if empty, remove the peer-id entirely
-            if mas.is_empty() {
-                self.mdns_discovered.remove(&p);
-            }
-        }
-    }
-
-    fn on_connection_established(
-        &mut self,
-        peer_id: PeerId,
-        connection_id: ConnectionId,
-        remote_ip: IpAddr,
-        remote_tcp_port: u16,
-    ) {
-        // send out connected event
-        self.pending_events
-            .push_back(ToSwarm::GenerateEvent(Event::ConnectionEstablished {
-                peer_id,
-                connection_id,
-                remote_ip,
-                remote_tcp_port,
-            }));
-    }
-
-    fn on_connection_closed(
-        &mut self,
-        peer_id: PeerId,
-        connection_id: ConnectionId,
-        remote_ip: IpAddr,
-        remote_tcp_port: u16,
-    ) {
-        // send out disconnected event
-        self.pending_events
-            .push_back(ToSwarm::GenerateEvent(Event::ConnectionClosed {
-                peer_id,
-                connection_id,
-                remote_ip,
-                remote_tcp_port,
-            }));
-    }
-}
-
-impl NetworkBehaviour for Behaviour {
-    type ConnectionHandler =
-        ConnectionHandlerSelect<dummy::ConnectionHandler, THandler<managed::Behaviour>>;
-    type ToSwarm = Event;
-
-    // simply delegate to underlying mDNS behaviour
-
-    delegate! {
-        to self.managed {
-            fn handle_pending_inbound_connection(&mut self, connection_id: ConnectionId, local_addr: &Multiaddr, remote_addr: &Multiaddr) -> Result<(), ConnectionDenied>;
-            fn handle_pending_outbound_connection(&mut self, connection_id: ConnectionId, maybe_peer: Option<PeerId>, addresses: &[Multiaddr], effective_role: Endpoint) -> Result<Vec<Multiaddr>, ConnectionDenied>;
-        }
-    }
-
-    fn handle_established_inbound_connection(
-        &mut self,
-        connection_id: ConnectionId,
-        peer: PeerId,
-        local_addr: &Multiaddr,
-        remote_addr: &Multiaddr,
-    ) -> Result<THandler<Self>, ConnectionDenied> {
-        Ok(ConnectionHandler::select(
-            dummy::ConnectionHandler,
-            self.managed.handle_established_inbound_connection(
-                connection_id,
-                peer,
-                local_addr,
-                remote_addr,
-            )?,
-        ))
-    }
-
-    #[allow(clippy::needless_question_mark)]
-    fn handle_established_outbound_connection(
-        &mut self,
-        connection_id: ConnectionId,
-        peer: PeerId,
-        addr: &Multiaddr,
-        role_override: Endpoint,
-        port_use: PortUse,
-    ) -> Result<THandler<Self>, ConnectionDenied> {
-        Ok(ConnectionHandler::select(
-            dummy::ConnectionHandler,
-            self.managed.handle_established_outbound_connection(
-                connection_id,
-                peer,
-                addr,
-                role_override,
-                port_use,
-            )?,
-        ))
-    }
-
-    fn on_connection_handler_event(
-        &mut self,
-        peer_id: PeerId,
-        connection_id: ConnectionId,
-        event: THandlerOutEvent<Self>,
-    ) {
-        match event {
-            Either::Left(ev) => libp2p::core::util::unreachable(ev),
-            Either::Right(ev) => {
-                self.managed
-                    .on_connection_handler_event(peer_id, connection_id, ev)
-            }
-        }
-    }
-
-    // hook into these methods to drive behavior
-
-    fn on_swarm_event(&mut self, event: FromSwarm) {
-        self.managed.on_swarm_event(event); // let mDNS handle swarm events
-
-        // handle swarm events to update internal state:
-        match event {
-            FromSwarm::ConnectionEstablished(ConnectionEstablished {
-                peer_id,
-                connection_id,
-                endpoint,
-                ..
-            }) => {
-                let remote_address = match endpoint {
-                    ConnectedPoint::Dialer { address, .. } => address,
-                    ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr,
-                };
-
-                if let Some((ip, port)) = remote_address.try_to_tcp_addr() {
-                    // handle connection established event which is filtered correctly
-                    self.on_connection_established(peer_id, connection_id, ip, port)
-                }
-            }
-            FromSwarm::ConnectionClosed(ConnectionClosed {
-                peer_id,
-                connection_id,
-                endpoint,
-                ..
-            }) => {
-                let remote_address = match endpoint {
-                    ConnectedPoint::Dialer { address, .. } => address,
-                    ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr,
-                };
-
-                if let Some((ip, port)) = remote_address.try_to_tcp_addr() {
-                    // handle connection closed event which is filtered correctly
-                    self.on_connection_closed(peer_id, connection_id, ip, port)
-                }
-            }
-
-            // since we are running TCP/IP transport layer, we are assuming that
-            // no address changes can occur, hence encountering one is a fatal error
-            FromSwarm::AddressChange(a) => {
-                unreachable!("unhandlable: address change encountered: {:?}", a)
-            }
-            _ => {}
-        }
-    }
-
-    fn poll(&mut self, cx: &mut Context) -> Poll<ToSwarm<Self::ToSwarm, THandlerInEvent<Self>>> {
-        // delegate to managed behaviors for any behaviors they need to perform
-        match self.managed.poll(cx) {
-            Poll::Ready(ToSwarm::GenerateEvent(e)) => {
-                match e {
-                    // handle discovered and expired events from mDNS
-                    managed::BehaviourEvent::Mdns(e) => match e.clone() {
-                        mdns::Event::Discovered(peers) => {
-                            self.handle_mdns_discovered(peers);
+impl Discovery {
+    pub async fn new(zid: ZenohId, listen_port: u16) -> io::Result<Self> {
+        let discovery_port = 52413;
+        let sock = Arc::new(UdpSocket::bind(format!("[::]:{discovery_port}")).await?);
+        //sock.set_multicast_loop_v6(false)?;
+        let ifaces: Arc<Mutex<Vec<SocketAddr>>> = Default::default();
+        let _sync = Mutex::new(
+            netwatcher::watch_interfaces_with_callback({
+                let sock = sock.clone();
+                let ifaces = ifaces.clone();
+                move |update| {
+                    for (iface_idx, iface) in update.interfaces.iter() {
+                        if iface
+                            .ipv6_ips()
+                            .all(|addr| addr.is_loopback() || addr.is_unspecified())
+                        {
+                            continue;
                         }
-                        mdns::Event::Expired(peers) => {
-                            self.handle_mdns_expired(peers);
+                        if let Err(e) = sock.join_multicast_v6(&GROUP, *iface_idx).inspect(|_| {
+                            ifaces.lock().push(SocketAddr::V6(SocketAddrV6::new(
+                                GROUP, 52413, 0, *iface_idx,
+                            )))
+                        }) {
+                            if let Some(iface) = update.interfaces.get(&iface_idx) {
+                                warn!(
+                                    "failed to join multicast v6 for interface {}: {e}",
+                                    iface.name
+                                )
+                            }
                         }
-                    },
-
-                    // handle ping events => if error then disconnect
-                    managed::BehaviourEvent::Ping(e) => {
-                        if let Err(_) = e.result {
-                            self.close_connection(e.peer, e.connection.clone())
+                    }
+                    for iface_idx in update.diff.removed {
+                        ifaces.lock().retain(|addr| {
+                            if let SocketAddr::V6(v6) = addr {
+                                v6.scope_id() != iface_idx
+                            } else {
+                                true
+                            }
+                        });
+                        if let Err(e) = sock.leave_multicast_v6(&GROUP, iface_idx) {
+                            if let Some(iface) = update.interfaces.get(&iface_idx) {
+                                warn!(
+                                    "failed to leave multicast v6 for interface {}: {e}",
+                                    iface.name
+                                )
+                            }
                         }
                     }
                 }
+            })
+            // todo: better error handling here
+            .expect("failed to bind discovery watcher"),
+        );
+        Ok(Self {
+            sock,
+            ifaces,
+            last_nonce: Default::default(),
+            listen_port,
+            zid,
+            tick: interval(Duration::from_secs(1)),
+            _sync,
+        })
+    }
 
-                // since we just consumed an event, we should immediately wake just in case
-                // there are more events to come where that came from
-                cx.waker().wake_by_ref();
-            }
-
-            // forward any other mDNS event to the swarm or its connection handler(s)
-            Poll::Ready(e) => {
-                return Poll::Ready(
-                    e.map_out(|_| unreachable!("events returning to swarm already handled"))
-                        .map_in(Either::Right),
-                );
-            }
-
-            Poll::Pending => {}
-        }
-
-        // retry connecting to all mDNS peers periodically (fails safely if already connected)
-        if self.retry_delay.poll(cx).is_ready() {
-            for (p, mas) in self.mdns_discovered.clone() {
-                for ma in mas {
-                    self.dial(p, ma)
+    pub async fn next(&mut self) -> io::Result<Discovered> {
+        let mut buf = [0u8; Hello::buf_size() + WhatsUp::buf_size() + 1];
+        loop {
+            tokio::select! {
+                _ = self.tick.tick() => {
+                    self.announce().await?;
+                }
+                res = self.sock.recv_from(&mut buf) => {
+                    let Ok((bytes_read, addr)) = res else { continue; };
+                    if let Some(discovered) = self.respond(bytes_read, addr, &mut buf).await? {
+                        return Ok(discovered)
+                    }
                 }
             }
-            // dial bootstrap peers (for environments where mDNS is unavailable)
-            for addr in &self.bootstrap_peers {
-                self.pending_events.push_back(ToSwarm::Dial {
-                    opts: DialOpts::unknown_peer_id().address(addr.clone()).build(),
-                })
+        }
+    }
+
+    async fn respond(
+        &self,
+        bytes_read: usize,
+        addr: SocketAddr,
+        buf: &mut [u8],
+    ) -> io::Result<Option<Discovered>> {
+        trace!(
+            "raw recv: {bytes_read} bytes from {addr}: {:02x?}",
+            &buf[..bytes_read]
+        );
+        if bytes_read < size_of::<Header>() {
+            trace!("dropped: early EOF");
+            return Ok(None);
+        };
+        let header: &Header = bytemuck::from_bytes(&buf[0..size_of::<Header>()]);
+        if header.magic != *b"EXO" {
+            trace!("dropped: wrong magic");
+            return Ok(None);
+        };
+        let Ok(kind) = header.kind.try_into() else {
+            trace!("dropped: unknown message kind {}", header.kind);
+            return Ok(None);
+        };
+        match kind {
+            Kind::Hello => {
+                let total = Hello::buf_size();
+                if bytes_read != total {
+                    trace!("dropped: hello wrong size");
+                    return Ok(None);
+                }
+                let hello: &Hello = bytemuck::from_bytes(&buf[size_of::<Header>()..total]);
+                if hello.nonce == *self.last_nonce.lock() {
+                    trace!("dropped: local hello nonce");
+                    return Ok(None);
+                }
+
+                // reply
+                let mut reply_buf = [0u8; WhatsUp::buf_size()];
+                let reply = WhatsUp {
+                    nonce: hello.nonce,
+                    zid: self.zid.to_le_bytes(),
+                    port_le: self.listen_port.to_le_bytes(),
+                };
+                reply.write_into(&mut reply_buf);
+
+                for i in 0..4 {
+                    if self
+                        .sock
+                        .send_to(&reply_buf, addr)
+                        .await
+                        .inspect_err(|e| debug!("send to {addr} failed: {e}"))
+                        .is_ok_and(|sent| sent == WhatsUp::buf_size())
+                    {
+                        trace!(
+                            "sent {} bytes to {addr} after {} attempt(s)",
+                            WhatsUp::buf_size(),
+                            i + 1
+                        );
+                        break;
+                    };
+                    tokio::time::sleep(Duration::from_millis(300)).await;
+                }
+                Ok(None)
+            }
+            Kind::WhatsUp => {
+                let total = WhatsUp::buf_size();
+                if bytes_read != total {
+                    trace!("dropped: whatsup wrong size");
+                    return Ok(None);
+                }
+                let whats_up: &WhatsUp = bytemuck::from_bytes(&buf[size_of::<Header>()..total]);
+                if whats_up.nonce == [0u8; 8] || whats_up.nonce != *self.last_nonce.lock() {
+                    trace!("dropped: stale nonce");
+                    return Ok(None);
+                }
+                let SocketAddr::V6(v6) = addr else {
+                    trace!("dropped: v4 addr used");
+                    return Ok(None);
+                };
+                let Ok(zid) = ZenohId::try_from(&whats_up.zid[..]) else {
+                    trace!("dropped: zenoh conversion failed");
+                    return Ok(None);
+                };
+                if zid == self.zid {
+                    trace!("dropped: self zenoh id");
+                    return Ok(None);
+                }
+                // discovered
+                let addr = {
+                    let mut x = v6.clone();
+                    x.set_port(u16::from_le_bytes(whats_up.port_le));
+                    x
+                };
+                Ok(Some(Discovered { addr, zid }))
             }
-            self.retry_delay.reset(RETRY_CONNECT_INTERVAL) // reset timeout
         }
+    }
 
-        // send out any pending events from our own service
-        if let Some(e) = self.pending_events.pop_front(cx) {
-            return Poll::Ready(e.map_in(Either::Left));
+    async fn announce(&self) -> io::Result<()> {
+        let nonce = rand::random();
+        *self.last_nonce.lock() = nonce;
+        let hello = Hello { nonce };
+
+        let mut buf = [0u8; Hello::buf_size()];
+        hello.write_into(&mut buf);
+
+        let addrs = self.ifaces.lock().clone();
+        debug!("announcing {hello:?} to {addrs:?}");
+        // rev so .remove() doesn't break things
+        for (i, addr) in addrs.into_iter().enumerate().rev() {
+            match self.sock.send_to(&buf, addr).await {
+                Ok(bytes) => trace!("sent {bytes} to {addr}"),
+                Err(e) if e.kind() == ErrorKind::HostUnreachable => {
+                    debug!("disabling discovery address {addr}: {e}");
+                    _ = self.ifaces.lock().swap_remove(i);
+                }
+                Err(e) => debug!("failed to reach {addr}: {e}"),
+            }
         }
-
-        // wait for pending events
-        Poll::Pending
+        Ok(())
     }
 }
+
+pub trait Message: Pod {
+    const KIND: Kind;
+    fn header() -> Header {
+        Header {
+            magic: *b"EXO",
+            kind: Self::KIND as u8,
+        }
+    }
+    fn write_into(&self, buf: &mut [u8]) {
+        let total = size_of::<Header>() + size_of::<Self>();
+        assert!(total <= buf.len());
+        buf[0..size_of::<Header>()].copy_from_slice(bytemuck::bytes_of(&Self::header()));
+        buf[size_of::<Header>()..total].copy_from_slice(bytemuck::bytes_of(self));
+    }
+}
+
+#[repr(u8)]
+#[derive(Debug, Clone, Copy)]
+// packet & version
+pub enum Kind {
+    Hello = 0,
+    WhatsUp = 1,
+}
+
+#[derive(Debug, Clone, Copy)]
+pub struct Discovered {
+    pub zid: ZenohId,
+    pub addr: SocketAddrV6,
+}
+
+pub struct UnknownKind;
+impl TryFrom<u8> for Kind {
+    type Error = UnknownKind;
+    fn try_from(value: u8) -> Result<Self, Self::Error> {
+        match value {
+            0 => Ok(Kind::Hello),
+            1 => Ok(Kind::WhatsUp),
+            _ => Err(UnknownKind),
+        }
+    }
+}
+#[repr(C)]
+#[derive(Debug, Clone, Copy, Pod, Zeroable)]
+pub struct Header {
+    magic: [u8; 3],
+    kind: u8,
+}
+
+#[repr(C)]
+#[derive(Debug, Clone, Copy, Pod, Zeroable)]
+pub struct Hello {
+    pub nonce: [u8; 8],
+}
+impl Hello {
+    const fn buf_size() -> usize {
+        size_of::<Header>() + size_of::<Self>()
+    }
+}
+impl Message for Hello {
+    const KIND: Kind = Kind::Hello;
+}
+
+#[repr(C)]
+#[derive(Debug, Clone, Copy, Pod, Zeroable)]
+pub struct WhatsUp {
+    pub nonce: [u8; 8],
+    pub zid: [u8; 16],
+    pub port_le: [u8; 2],
+}
+impl WhatsUp {
+    const fn buf_size() -> usize {
+        size_of::<Header>() + size_of::<Self>()
+    }
+}
+impl Message for WhatsUp {
+    const KIND: Kind = Kind::WhatsUp;
+}

rust/networking/src/lib.rs

@@ -1,44 +1,92 @@
-//! TODO: crate documentation
-//!
-//! this is here as a placeholder documentation
-//!
-//!
+use std::env;
+
+use tokio::task::JoinHandle;
+use zenoh::{Result, Session as ZSession, config::Locator};
+use zenoh_plugin_storage_manager::StoragesPlugin;
+use zenoh_plugin_trait::PluginsManager;
+
+pub use zenoh::{Config, config::ZenohId};
+
+use crate::discovery::Discovery;
+
 pub mod discovery;
 pub mod swarm;
 
-/// Namespace for all the type/trait aliases used by this crate.
-pub(crate) mod alias {
-    use std::error::Error;
-
-    pub type AnyError = Box<dyn Error + Send + Sync + 'static>;
-    pub type AnyResult<T> = Result<T, AnyError>;
+pub fn cfg(identity: u128, listen_port: u16) -> Result<zenoh::Config> {
+    assert!(listen_port != 0, "must used defined listen port port");
+    let namespace = env::var("EXO_ZENOH_NAMESPACE").unwrap_or_else(|_| "exo".to_string());
+    let mut cfg = zenoh::Config::default();
+    // todo: cleanup
+    cfg.insert_json5("id", &format!("\"{identity:x}\""))?;
+    cfg.insert_json5("mode", "\"router\"")?;
+    cfg.insert_json5("listen/endpoints", &format!("[\"tcp/[::]:{listen_port}\"]"))?;
+    cfg.insert_json5("scouting/multicast/enabled", "false")?;
+    cfg.insert_json5("scouting/multicast/autoconnect", "[]")?;
+    cfg.insert_json5("scouting/gossip/multihop", "true")?;
+    cfg.insert_json5("namespace", &format!("{namespace:?}"))?;
+    cfg.insert_json5("transport/link/tx/batch_size", "9216")?;
+    cfg.insert_json5("timestamping/enabled", "true")?;
+    cfg.insert_json5("plugins/storage_manager/__required__", "true")?;
+    cfg.insert_json5(
+        "plugins/storage_manager/storages/mem1",
+        r#"{
+            key_expr: "storage/mem1/**",
+            strip_prefix: "storage/mem1",
+            volume: "memory",
+            replication: {
+                interval: 2,
+            }
+        }"#,
+    )?;
+    Ok(cfg)
 }
 
-/// Namespace for crate-wide extension traits/methods
-pub(crate) mod ext {
-    use extend::ext;
-    use libp2p::Multiaddr;
-    use libp2p::multiaddr::Protocol;
-    use std::net::IpAddr;
+pub async fn open(cfg: zenoh::Config, listen_port: u16) -> Result<Session> {
+    assert!(listen_port != 0, "must used defined listen port");
+    let mut plugins = PluginsManager::static_plugins_only();
+    plugins.declare_static_plugin::<StoragesPlugin, _>("storage_manager", true);
+    let mut runtime = zenoh::internal::runtime::RuntimeBuilder::new(cfg)
+        .plugins_manager(plugins)
+        .build()
+        .await?;
+    let z = zenoh::session::init(runtime.clone().into()).await?;
+    runtime.start().await?;
+    let mut discovery = Discovery::new(z.zid(), listen_port).await?;
+    let _jh = tokio::task::spawn(async move {
+        loop {
+            let Ok(discovered) = discovery.next().await.inspect_err(|e| {
+                log::warn!("discovery error {e}");
+            }) else {
+                continue;
+            };
 
-    #[ext(pub, name = MultiaddrExt)]
-    impl Multiaddr {
-        /// If the multiaddress corresponds to a TCP address, extracts it
-        fn try_to_tcp_addr(&self) -> Option<(IpAddr, u16)> {
-            let mut ps = self.into_iter();
-            let ip = if let Some(p) = ps.next() {
-                match p {
-                    Protocol::Ip4(ip) => IpAddr::V4(ip),
-                    Protocol::Ip6(ip) => IpAddr::V6(ip),
-                    _ => return None,
-                }
-            } else {
-                return None;
+            if discovered.zid > runtime.zid() {
+                log::debug!("not connecting to peer with greater zid");
+                continue;
+            }
+
+            let Ok(locator) =
+                Locator::new("tcp", discovered.addr.to_string(), "").inspect_err(|e| {
+                    log::warn!("failed to pass locator from addr: {e}");
+                })
+            else {
+                continue;
             };
-            let Some(Protocol::Tcp(port)) = ps.next() else {
-                return None;
-            };
-            Some((ip, port))
+
+            runtime
+                .connect_peer(&discovered.zid.into(), &[locator])
+                .await;
         }
+    });
+    Ok(Session { z, _jh })
+}
+
+pub struct Session {
+    pub z: ZSession,
+    _jh: JoinHandle<()>,
+}
+impl Drop for Session {
+    fn drop(&mut self) {
+        self._jh.abort();
     }
 }

rust/networking/src/swarm.rs

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

rust/networking/tests/bootstrap_peers.rs

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

rust/networking/tests/dummy.rs

@@ -1,7 +0,0 @@
-// maybe this will hold test in the future...??
-
-#[cfg(test)]
-mod tests {
-    #[test]
-    fn does_nothing() {}
-}

rust/parts.nix

@@ -55,6 +55,7 @@
         ];
 
         OPENSSL_NO_VENDOR = "1";
+        MATURIN_NO_INSTALL_RUST = "1";
 
         # Required for pyo3 tests to find libpython
         LD_LIBRARY_PATH = lib.makeLibraryPath [ pkgs.python313 ];
@@ -81,11 +82,11 @@
       config = {
         packages = {
           # Python bindings wheel via maturin
-          exo_pyo3_bindings = craneLib.buildPackage (
+          exo-net = craneLib.buildPackage (
             commonArgs
             // {
               inherit cargoArtifacts;
-              pname = "exo_pyo3_bindings";
+              pname = "exo-net";
 
               nativeBuildInputs = commonArgs.nativeBuildInputs ++ [
                 pkgs.maturin
@@ -95,7 +96,7 @@
                 maturin build \
                   --release \
                   --manylinux off \
-                  --manifest-path rust/exo_pyo3_bindings/Cargo.toml \
+                  --manifest-path rust/exo_net/Cargo.toml \
                   --features "pyo3/extension-module,pyo3/experimental-async" \
                   --interpreter ${pkgs.python313}/bin/python \
                   --out dist

rust/util/Cargo.toml

@@ -1,15 +0,0 @@
-[package]
-name = "util"
-version = { workspace = true }
-edition = { workspace = true }
-publish = false
-
-[lib]
-doctest = false
-name = "util"
-path = "src/lib.rs"
-
-[lints]
-workspace = true
-
-[dependencies]

rust/util/src/lib.rs

@@ -1 +0,0 @@
-pub mod wakerdeque;

rust/util/src/wakerdeque.rs

@@ -1,55 +0,0 @@
-use std::collections::VecDeque;
-use std::fmt::{Debug, Formatter};
-use std::task::{Context, Waker};
-
-/// A wrapper around [`VecDeque`] which wakes (if it can) on any `push_*` methods,
-/// and updates the internally stored waker by consuming [`Context`] on any `pop_*` methods.
-pub struct WakerDeque<T> {
-    waker: Option<Waker>,
-    deque: VecDeque<T>,
-}
-
-impl<T: Debug> Debug for WakerDeque<T> {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        self.deque.fmt(f)
-    }
-}
-
-impl<T> WakerDeque<T> {
-    pub fn new() -> Self {
-        Self {
-            waker: None,
-            deque: VecDeque::new(),
-        }
-    }
-
-    fn update(&mut self, cx: &mut Context<'_>) {
-        self.waker = Some(cx.waker().clone());
-    }
-
-    fn wake(&mut self) {
-        let Some(ref mut w) = self.waker else { return };
-        w.wake_by_ref();
-        self.waker = None;
-    }
-
-    pub fn pop_front(&mut self, cx: &mut Context<'_>) -> Option<T> {
-        self.update(cx);
-        self.deque.pop_front()
-    }
-
-    pub fn pop_back(&mut self, cx: &mut Context<'_>) -> Option<T> {
-        self.update(cx);
-        self.deque.pop_back()
-    }
-
-    pub fn push_front(&mut self, value: T) {
-        self.wake();
-        self.deque.push_front(value);
-    }
-
-    pub fn push_back(&mut self, value: T) {
-        self.wake();
-        self.deque.push_back(value);
-    }
-}

src/exo/api/adapters/ollama.py

@@ -20,7 +20,6 @@ from exo.shared.types.chunks import (
     TokenChunk,
     ToolCallChunk,
 )
-from exo.shared.types.common import CommandId
 from exo.shared.types.text_generation import (
     Base64Image,
     InputMessage,
@@ -181,7 +180,6 @@ def ollama_request_to_text_generation(
 
 
 async def generate_ollama_chat_stream(
-    _command_id: CommandId,
     chunk_stream: AsyncGenerator[
         ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
     ],
@@ -264,7 +262,6 @@ async def generate_ollama_chat_stream(
 
 
 async def collect_ollama_chat_response(
-    _command_id: CommandId,
     chunk_stream: AsyncGenerator[
         ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
     ],
@@ -369,7 +366,6 @@ def ollama_generate_request_to_text_generation(
 
 
 async def generate_ollama_generate_stream(
-    _command_id: CommandId,
     chunk_stream: AsyncGenerator[
         ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
     ],
@@ -442,7 +438,6 @@ async def generate_ollama_generate_stream(
 
 
 async def collect_ollama_generate_response(
-    _command_id: CommandId,
     chunk_stream: AsyncGenerator[
         ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
     ],

src/exo/api/tests/test_api_snapshot_bootstrap.py

@@ -1,136 +0,0 @@
-# pyright: reportPrivateUsage=false
-
-import hashlib
-
-import anyio
-import pytest
-import zstandard
-
-from exo.api.main import API
-from exo.routing.event_router import EventRouter
-from exo.shared.types.commands import ForwarderCommand, RequestSnapshot
-from exo.shared.types.common import NodeId, SessionId, SystemId
-from exo.shared.types.events import (
-    Event,
-    GlobalForwarderEvent,
-    IndexedEvent,
-    LocalForwarderEvent,
-    TestEvent,
-)
-from exo.shared.types.snapshots import SnapshotChunk, SnapshotTransferId
-from exo.shared.types.state import State
-from exo.utils.channels import Receiver, Sender, channel
-
-
-class _FakeEventLog:
-    def __init__(self) -> None:
-        self.appended: list[Event] = []
-
-    def append(self, event: Event) -> None:
-        self.appended.append(event)
-
-
-def _snapshot_chunk(
-    state: State, *, requester_node_id: NodeId, session_id: SessionId
-) -> SnapshotChunk:
-    body = zstandard.ZstdCompressor().compress(state.model_dump_json().encode("utf-8"))
-    return SnapshotChunk.from_data(
-        data=body,
-        transfer_id=SnapshotTransferId("transfer-1"),
-        requester_node_id=requester_node_id,
-        session_id=session_id,
-        schema_version=state.schema_version,
-        last_event_applied_idx=state.last_event_applied_idx,
-        chunk_index=0,
-        total_chunks=1,
-        sha256_hex=hashlib.sha256(body).hexdigest(),
-    )
-
-
-def _api(
-    node_id: NodeId, session_id: SessionId
-) -> tuple[
-    API,
-    EventRouter,
-    Receiver[ForwarderCommand],
-    Sender[SnapshotChunk],
-    Sender[IndexedEvent],
-    _FakeEventLog,
-]:
-    router_command_sender, _router_command_receiver = channel[ForwarderCommand]()
-    _global_event_sender, global_event_receiver = channel[GlobalForwarderEvent]()
-    local_event_sender, _local_event_receiver = channel[LocalForwarderEvent]()
-    event_router = EventRouter(
-        session_id=session_id,
-        command_sender=router_command_sender,
-        external_inbound=global_event_receiver,
-        external_outbound=local_event_sender,
-    )
-
-    event_sender, event_receiver = channel[IndexedEvent]()
-    command_sender, command_receiver = channel[ForwarderCommand]()
-    snapshot_sender, snapshot_receiver = channel[SnapshotChunk]()
-
-    api = object.__new__(API)
-    api.node_id = node_id
-    api.session_id = session_id
-    api.event_router = event_router
-    api.event_receiver = event_receiver
-    api.snapshot_chunk_receiver = snapshot_receiver
-    api.command_sender = command_sender
-    api._system_id = SystemId("api-system")
-    api.state = State()
-    event_log = _FakeEventLog()
-    api._event_log = event_log  # pyright: ignore[reportAttributeAccessIssue]
-    api._image_generation_queues = {}
-    api._text_generation_queues = {}
-    return api, event_router, command_receiver, snapshot_sender, event_sender, event_log
-
-
-@pytest.mark.asyncio
-async def test_api_fetch_snapshot_applies_state_and_fast_forwards_router() -> None:
-    node_id = NodeId("api")
-    session_id = SessionId(master_node_id=NodeId("master"), election_clock=1)
-    api, event_router, command_receiver, snapshot_sender, _event_sender, _event_log = (
-        _api(node_id, session_id)
-    )
-    state = State(last_event_applied_idx=7)
-
-    async with anyio.create_task_group() as tg:
-        tg.start_soon(api._fetch_snapshot)
-        command = await command_receiver.receive()
-        assert isinstance(command.command, RequestSnapshot)
-        assert command.command.requester_node_id == node_id
-
-        await snapshot_sender.send(
-            _snapshot_chunk(state, requester_node_id=node_id, session_id=session_id)
-        )
-
-    assert api.state.last_event_applied_idx == 7
-    assert event_router.event_buffer.next_idx_to_release == 8
-
-
-@pytest.mark.asyncio
-async def test_api_apply_state_ignores_events_covered_by_snapshot() -> None:
-    node_id = NodeId("api")
-    session_id = SessionId(master_node_id=NodeId("master"), election_clock=1)
-    (
-        api,
-        _event_router,
-        _command_receiver,
-        _snapshot_sender,
-        event_sender,
-        event_log,
-    ) = _api(node_id, session_id)
-    api.state = State(last_event_applied_idx=7)
-
-    async with anyio.create_task_group() as tg:
-        tg.start_soon(api._apply_state)
-        await event_sender.send(IndexedEvent(idx=7, event=TestEvent()))
-        await event_sender.send(IndexedEvent(idx=8, event=TestEvent()))
-
-        while api.state.last_event_applied_idx != 8:
-            await anyio.sleep(0.001)
-        tg.cancel_scope.cancel()
-
-    assert len(event_log.appended) == 1

src/exo/api/tests/test_cancel_command.py

@@ -1,11 +1,11 @@
 # pyright: reportUnusedFunction=false, reportAny=false
 from typing import Any
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import AsyncMock
 
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
 
-from exo.api.main import API
+from exo.api.main import API, Transport
 from exo.shared.types.common import CommandId
 
 
@@ -15,9 +15,9 @@ def _make_api() -> Any:
     app = FastAPI()
     api = object.__new__(API)
     api.app = app
-    api._text_generation_queues = {}  # pyright: ignore[reportPrivateUsage]
-    api._image_generation_queues = {}  # pyright: ignore[reportPrivateUsage]
-    api._send = AsyncMock()  # pyright: ignore[reportPrivateUsage]
+    api.transport = object.__new__(Transport)
+    api.transport.cancel = AsyncMock()
+    api.transport.send_command = AsyncMock()
     api._setup_exception_handlers()  # pyright: ignore[reportPrivateUsage]
     app.post("/v1/cancel/{command_id}")(api.cancel_command)
     return api
@@ -43,16 +43,14 @@ def test_cancel_active_text_generation() -> None:
     client = TestClient(api.app)
 
     cid = CommandId("text-cmd-123")
-    sender = MagicMock()
-    api._text_generation_queues[cid] = sender
 
     response = client.post(f"/v1/cancel/{cid}")
     assert response.status_code == 200
     data: dict[str, Any] = response.json()
     assert data["message"] == "Command cancelled."
     assert data["command_id"] == str(cid)
-    sender.close.assert_called_once()
-    api._send.assert_called_once()
+    api.transport.cancel.assert_called_once()
+    api.transport.send_command.assert_called_once()
     task_cancelled = api._send.call_args[0][0]
     assert task_cancelled.cancelled_command_id == cid
 
@@ -63,15 +61,13 @@ def test_cancel_active_image_generation() -> None:
     client = TestClient(api.app)
 
     cid = CommandId("img-cmd-456")
-    sender = MagicMock()
-    api._image_generation_queues[cid] = sender
 
     response = client.post(f"/v1/cancel/{cid}")
     assert response.status_code == 200
     data: dict[str, Any] = response.json()
     assert data["message"] == "Command cancelled."
     assert data["command_id"] == str(cid)
-    sender.close.assert_called_once()
-    api._send.assert_called_once()
-    task_cancelled = api._send.call_args[0][0]
+    api.transport.cancel.assert_called_once()
+    api.transport.send_command.assert_called_once()
+    task_cancelled = api.transport.send_command.call_args[0][0]
     assert task_cancelled.cancelled_command_id == cid

src/exo/api/tests/test_instance_deleted_stream_cleanup.py

@@ -1,9 +1,10 @@
 # pyright: reportUnusedFunction=false, reportAny=false
 """Tests that InstanceDeleted events close active generation streams."""
 
+from typing import Any
 from unittest.mock import MagicMock
 
-from exo.api.main import API
+from exo.api.main import API, Transport
 from exo.api.types import ImageGenerationTaskParams
 from exo.shared.types.common import CommandId, ModelId
 from exo.shared.types.state import State
@@ -16,12 +17,11 @@ from exo.shared.types.text_generation import (
 from exo.shared.types.worker.instances import InstanceId
 
 
-def _make_api_with_state(state: State) -> API:
+def _make_api_with_state(state: State) -> Any:
     """Create a minimal API instance with pre-set state."""
     api = object.__new__(API)
     api.state = state
-    api._text_generation_queues = {}  # pyright: ignore[reportPrivateUsage]
-    api._image_generation_queues = {}  # pyright: ignore[reportPrivateUsage]
+    api.transport = object.__new__(Transport)
     return api
 
 
@@ -47,13 +47,10 @@ def test_close_streams_for_deleted_instance() -> None:
     state = State(tasks={task.task_id: task})
     api = _make_api_with_state(state)
 
-    sender = MagicMock()
-    api._text_generation_queues[command_id] = sender  # pyright: ignore[reportPrivateUsage]
+    api._close_streams_for_instance(instance_id)
 
-    api._close_streams_for_instance(instance_id)  # pyright: ignore[reportPrivateUsage]
-
-    sender.close.assert_called_once()
-    assert command_id not in api._text_generation_queues  # pyright: ignore[reportPrivateUsage]
+    api.transport.cancel.assert_called_once()
+    assert api.transport.cancel.call_args[0][0] == command_id
 
 
 def test_close_streams_ignores_unrelated_instances() -> None:
@@ -72,7 +69,6 @@ def test_close_streams_ignores_unrelated_instances() -> None:
     api._close_streams_for_instance(target_id)  # pyright: ignore[reportPrivateUsage]
 
     sender.close.assert_not_called()
-    assert other_cmd in api._text_generation_queues  # pyright: ignore[reportPrivateUsage]
 
 
 def test_close_streams_for_deleted_instance_image_generation() -> None:

src/exo/download/coordinator.py

@@ -16,7 +16,8 @@ from exo.download.download_utils import (
 )
 from exo.download.shard_downloader import ShardDownloader
 from exo.shared.constants import EXO_DEFAULT_MODELS_DIR, EXO_MODELS_READ_ONLY_DIRS
-from exo.shared.models.model_cards import ModelId, get_model_cards
+from exo.shared.models import model_cards
+from exo.shared.models.model_cards import ModelId
 from exo.shared.types.commands import (
     CancelDownload,
     DeleteDownload,
@@ -422,7 +423,7 @@ class DownloadCoordinator:
                     )
                 # Scan read-only directories for pre-downloaded models
                 if EXO_MODELS_READ_ONLY_DIRS:
-                    for card in await get_model_cards():
+                    for card in await model_cards.card_cache.list_all():
                         mid = card.model_id
                         if mid in self.active_downloads:
                             continue

src/exo/download/impl_shard_downloader.py

@@ -11,11 +11,11 @@ from exo.download.download_utils import (
     download_shard,
 )
 from exo.download.shard_downloader import ShardDownloader
+from exo.shared.models import model_cards
 from exo.shared.models.model_cards import (
     ModelCard,
     ModelId,
     ModelTask,
-    get_model_cards,
 )
 from exo.shared.types.memory import Memory
 from exo.shared.types.worker.shards import (
@@ -258,7 +258,7 @@ class ResumableShardDownloader(ShardDownloader):
 
         tasks = [
             create_task(download_with_semaphore(model_card))
-            for model_card in await get_model_cards()
+            for model_card in await model_cards.card_cache.list_all()
         ]
 
         for task in asyncio.as_completed(tasks):

src/exo/main.py

@@ -3,10 +3,13 @@ import multiprocessing as mp
 import os
 import resource
 import signal
+import sys
 from dataclasses import dataclass, field
 from typing import Self
+from uuid import uuid4
 
 import anyio
+from exo_net import Pidfile, PidfileError, PySession
 from loguru import logger
 from pydantic import PositiveInt
 
@@ -16,8 +19,8 @@ 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_id_keypair
-from exo.shared.constants import EXO_LOG
+from exo.routing.router import Router
+from exo.shared.constants import EXO_LOG, EXO_PID_FILE
 from exo.shared.election import Election, ElectionResult
 from exo.shared.logging import logger_cleanup, logger_setup
 from exo.shared.types.common import NodeId, SessionId
@@ -39,32 +42,31 @@ class Node:
     api: API | None
 
     node_id: NodeId
+    session: PySession
     offline: bool
     _api_port: int
     _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
 
     @classmethod
     async def create(cls, args: "Args") -> Self:
-        keypair = get_node_id_keypair()
-        node_id = NodeId(keypair.to_node_id())
+        node_id_bytes = uuid4()
+        node_id = NodeId(str(node_id_bytes))
         session_id = SessionId(master_node_id=node_id, election_clock=0)
-        router = Router.create(
-            keypair,
+        router, session = Router.create(
+            node_id_bytes.bytes,
             bootstrap_peers=args.bootstrap_peers,
             listen_port=args.libp2p_port,
         )
         await router.register_topic(topics.GLOBAL_EVENTS)
         await router.register_topic(topics.LOCAL_EVENTS)
-        await router.register_topic(topics.COMMANDS)
         await router.register_topic(topics.ELECTION_MESSAGES)
         await router.register_topic(topics.CONNECTION_MESSAGES)
         await router.register_topic(topics.DOWNLOAD_COMMANDS)
-        await router.register_topic(topics.SNAPSHOT_RESPONSES)
         event_router = EventRouter(
             session_id,
-            command_sender=router.sender(topics.COMMANDS),
             external_outbound=router.sender(topics.LOCAL_EVENTS),
             external_inbound=router.receiver(topics.GLOBAL_EVENTS),
+            command_sender=session.net_sender("orchestrator"),
         )
 
         logger.info(f"Starting node {node_id}")
@@ -84,14 +86,11 @@ class Node:
         if args.spawn_api:
             api = API(
                 node_id,
-                session_id,
                 port=args.api_port,
-                event_router=event_router,
                 event_receiver=event_router.receiver(),
-                snapshot_chunk_receiver=router.receiver(topics.SNAPSHOT_RESPONSES),
-                command_sender=router.sender(topics.COMMANDS),
                 download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
                 election_receiver=router.receiver(topics.ELECTION_MESSAGES),
+                session=session,
             )
         else:
             api = None
@@ -99,14 +98,11 @@ class Node:
         if not args.no_worker:
             worker = Worker(
                 node_id,
-                session_id,
-                event_router=event_router,
                 event_receiver=event_router.receiver(),
                 event_sender=event_router.sender(),
-                snapshot_chunk_receiver=router.receiver(topics.SNAPSHOT_RESPONSES),
-                command_sender=router.sender(topics.COMMANDS),
                 download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
                 api_port=args.api_port,
+                session=session,
             )
         else:
             worker = None
@@ -118,9 +114,8 @@ class Node:
             event_sender=event_router.sender(),
             global_event_sender=router.sender(topics.GLOBAL_EVENTS),
             local_event_receiver=router.receiver(topics.LOCAL_EVENTS),
-            command_receiver=router.receiver(topics.COMMANDS),
-            snapshot_chunk_sender=router.sender(topics.SNAPSHOT_RESPONSES),
             download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
+            command_receiver=session.net_receiver("orchestrator"),
         )
 
         er_send, er_recv = channel[ElectionResult]()
@@ -133,7 +128,6 @@ class Node:
             election_message_sender=router.sender(topics.ELECTION_MESSAGES),
             election_message_receiver=router.receiver(topics.ELECTION_MESSAGES),
             connection_message_receiver=router.receiver(topics.CONNECTION_MESSAGES),
-            command_receiver=router.receiver(topics.COMMANDS),
             election_result_sender=er_send,
         )
 
@@ -147,6 +141,7 @@ class Node:
             master,
             api,
             node_id,
+            session,
             args.offline,
             args.api_port,
         )
@@ -196,7 +191,7 @@ class Node:
                     self.event_router.shutdown()
                     self.event_router = EventRouter(
                         result.session_id,
-                        self.router.sender(topics.COMMANDS),
+                        self.session.net_sender("orchestrator"),
                         self.router.receiver(topics.GLOBAL_EVENTS),
                         self.router.sender(topics.LOCAL_EVENTS),
                     )
@@ -217,13 +212,10 @@ class Node:
                         event_sender=self.event_router.sender(),
                         global_event_sender=self.router.sender(topics.GLOBAL_EVENTS),
                         local_event_receiver=self.router.receiver(topics.LOCAL_EVENTS),
-                        command_receiver=self.router.receiver(topics.COMMANDS),
-                        snapshot_chunk_sender=self.router.sender(
-                            topics.SNAPSHOT_RESPONSES
-                        ),
                         download_command_sender=self.router.sender(
                             topics.DOWNLOAD_COMMANDS
                         ),
+                        command_receiver=self.session.net_receiver("orchestrator"),
                     )
                     self._tg.start_soon(self.master.run)
                 elif (
@@ -257,28 +249,17 @@ class Node:
                         # TODO: add profiling etc to resource monitor
                         self.worker = Worker(
                             self.node_id,
-                            result.session_id,
-                            event_router=self.event_router,
                             event_receiver=self.event_router.receiver(),
                             event_sender=self.event_router.sender(),
-                            snapshot_chunk_receiver=self.router.receiver(
-                                topics.SNAPSHOT_RESPONSES
-                            ),
-                            command_sender=self.router.sender(topics.COMMANDS),
                             download_command_sender=self.router.sender(
                                 topics.DOWNLOAD_COMMANDS
                             ),
                             api_port=self._api_port,
+                            session=self.session,
                         )
                         self._tg.start_soon(self.worker.run)
                     if self.api:
-                        self.api.reset(
-                            result.won_clock,
-                            result.session_id,
-                            self.event_router,
-                            self.event_router.receiver(),
-                            self.router.receiver(topics.SNAPSHOT_RESPONSES),
-                        )
+                        self.api.reset(result.won_clock, self.event_router.receiver())
                     self._tg.start_soon(self.event_router.run)
                 else:
                     if self.api:
@@ -286,12 +267,21 @@ class Node:
 
 
 def main():
+    # Exit early if no PID file (not compatible with double-for daemonization yet)
+    try:
+        pidfile = Pidfile(EXO_PID_FILE, 0o0600)
+        pidfile.write()
+    except (PidfileError, OSError) as exception:
+        print(exception, file=sys.stderr)
+        raise SystemExit(1) from exception
+
     args = Args.parse()
     soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
     target = min(max(soft, 65535), hard)
     resource.setrlimit(resource.RLIMIT_NOFILE, (target, hard))
 
     mp.set_start_method("spawn", force=True)
+
     # TODO: Refactor the current verbosity system
     logger_setup(EXO_LOG, args.verbosity)
     logger.info(f"{'=' * 40}")
@@ -328,6 +318,7 @@ def main():
     finally:
         logger.info("EXO Shutdown complete")
         logger_cleanup()
+        del pidfile
 
 
 class Args(FrozenModel):

src/exo/master/main.py

@@ -1,9 +1,9 @@
-import hashlib
 from datetime import datetime, timedelta, timezone
 
 import anyio
-from anyio import to_thread
+from exo_net import NetReceiver
 from loguru import logger
+from pydantic import TypeAdapter
 
 from exo.master.placement import (
     add_instance_to_placements,
@@ -17,17 +17,16 @@ 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,
+    Command,
     CreateInstance,
     DeleteCustomModelCard,
     DeleteInstance,
     DeleteInstanceLink,
-    ForwarderCommand,
     ForwarderDownloadCommand,
     ImageEdits,
     ImageGeneration,
     PlaceInstance,
     RequestEventLog,
-    RequestSnapshot,
     SendInputChunk,
     SetInstanceLink,
     TaskCancelled,
@@ -57,7 +56,6 @@ from exo.shared.types.events import (
     TracesMerged,
 )
 from exo.shared.types.instance_link import InstanceLink
-from exo.shared.types.snapshots import SnapshotChunk, SnapshotTransferId
 from exo.shared.types.state import State
 from exo.shared.types.tasks import (
     ImageEdits as ImageEditsTask,
@@ -78,15 +76,6 @@ from exo.utils.disk_event_log import DiskEventLog
 from exo.utils.event_buffer import MultiSourceBuffer
 from exo.utils.task_group import TaskGroup
 
-_SNAPSHOT_CHUNK_BYTES = 512 * 1024
-_MAX_EVENT_LOG_REPLAY_BATCH = 1000
-
-
-def _encode_state_for_transfer(state: State) -> bytes:
-    import zstandard
-
-    return zstandard.ZstdCompressor().compress(state.model_dump_json().encode("utf-8"))
-
 
 def _prefill_endpoint_for(state: State, decode_instance_id: InstanceId) -> str | None:
     decode = state.instances.get(decode_instance_id)
@@ -134,11 +123,10 @@ class Master:
         node_id: NodeId,
         session_id: SessionId,
         *,
-        command_receiver: Receiver[ForwarderCommand],
+        command_receiver: NetReceiver,  # todo: not this type
         event_sender: Sender[Event],
         local_event_receiver: Receiver[LocalForwarderEvent],
         global_event_sender: Sender[GlobalForwarderEvent],
-        snapshot_chunk_sender: Sender[SnapshotChunk],
         download_command_sender: Sender[ForwarderDownloadCommand],
     ):
         self.node_id = node_id
@@ -149,7 +137,6 @@ class Master:
         self.command_receiver = command_receiver
         self.local_event_receiver = local_event_receiver
         self.global_event_sender = global_event_sender
-        self.snapshot_chunk_sender = snapshot_chunk_sender
         self.download_command_sender = download_command_sender
         self.event_sender = event_sender
         self._system_id = SystemId()
@@ -170,310 +157,290 @@ class Master:
             self._event_log.close()
             self.global_event_sender.close()
             self.local_event_receiver.close()
-            self.snapshot_chunk_sender.close()
-            self.command_receiver.close()
 
     async def shutdown(self):
         logger.info("Stopping Master")
         self._tg.cancel_tasks()
 
     async def _command_processor(self) -> None:
-        with self.command_receiver as commands:
-            async for forwarder_command in commands:
-                try:
-                    logger.info(f"Executing command: {forwarder_command.command}")
+        while True:
+            data = await self.command_receiver.recv()
+            if not data:
+                break
+            try:
+                command = TypeAdapter[Command](Command).validate_json(data)
+                logger.info(f"Executing command: {command}")
 
-                    generated_events: list[Event] = []
-                    command = forwarder_command.command
-                    instance_task_counts: dict[InstanceId, int] = {}
-                    match command:
-                        case TestCommand():
-                            pass
-                        case TextGeneration():
-                            prefill_only: set[InstanceId] = set()
-                            for link in self.state.instance_links.values():
-                                prefill_only.update(link.prefill_instances)
-                            for link in self.state.instance_links.values():
-                                prefill_only.difference_update(link.decode_instances)
+                generated_events: list[Event] = []
+                instance_task_counts: dict[InstanceId, int] = {}
+                match command:
+                    case TestCommand():
+                        pass
+                    case TextGeneration():
+                        prefill_only: set[InstanceId] = set()
+                        for link in self.state.instance_links.values():
+                            prefill_only.update(link.prefill_instances)
+                        for link in self.state.instance_links.values():
+                            prefill_only.difference_update(link.decode_instances)
 
-                            for instance in self.state.instances.values():
-                                if (
-                                    instance.shard_assignments.model_id
-                                    == command.task_params.model
-                                    and instance.instance_id not in prefill_only
-                                ):
-                                    in_flight = {TaskStatus.Pending, TaskStatus.Running}
-                                    task_count = sum(
-                                        1
-                                        for task in self.state.tasks.values()
-                                        if task.instance_id == instance.instance_id
-                                        and task.task_status in in_flight
-                                    )
-                                    instance_task_counts[instance.instance_id] = (
-                                        task_count
-                                    )
-
-                            if not instance_task_counts:
-                                raise ValueError(
-                                    f"No instance found for model {command.task_params.model}"
-                                )
-
-                            available_instance_ids = sorted(
-                                instance_task_counts.keys(),
-                                key=lambda instance_id: instance_task_counts[
-                                    instance_id
-                                ],
-                            )
-
-                            decode_instance_id = available_instance_ids[0]
-                            task_id = TaskId()
-                            params = command.task_params.model_copy(
-                                update={
-                                    "prefill_endpoint": _prefill_endpoint_for(
-                                        self.state, decode_instance_id
-                                    ),
-                                }
-                            )
-                            generated_events.append(
-                                TaskCreated(
-                                    task_id=task_id,
-                                    task=TextGenerationTask(
-                                        task_id=task_id,
-                                        command_id=command.command_id,
-                                        instance_id=decode_instance_id,
-                                        task_status=TaskStatus.Pending,
-                                        task_params=params,
-                                    ),
-                                )
-                            )
-                            self.command_task_mapping[command.command_id] = task_id
-                        case ImageGeneration():
-                            for instance in self.state.instances.values():
-                                if (
-                                    instance.shard_assignments.model_id
-                                    == command.task_params.model
-                                ):
-                                    in_flight = {TaskStatus.Pending, TaskStatus.Running}
-                                    task_count = sum(
-                                        1
-                                        for task in self.state.tasks.values()
-                                        if task.instance_id == instance.instance_id
-                                        and task.task_status in in_flight
-                                    )
-                                    instance_task_counts[instance.instance_id] = (
-                                        task_count
-                                    )
-
-                            if not instance_task_counts:
-                                raise ValueError(
-                                    f"No instance found for model {command.task_params.model}"
-                                )
-
-                            available_instance_ids = sorted(
-                                instance_task_counts.keys(),
-                                key=lambda instance_id: instance_task_counts[
-                                    instance_id
-                                ],
-                            )
-
-                            task_id = TaskId()
-                            selected_instance_id = available_instance_ids[0]
-                            generated_events.append(
-                                TaskCreated(
-                                    task_id=task_id,
-                                    task=ImageGenerationTask(
-                                        task_id=task_id,
-                                        command_id=command.command_id,
-                                        instance_id=selected_instance_id,
-                                        task_status=TaskStatus.Pending,
-                                        task_params=command.task_params,
-                                    ),
-                                )
-                            )
-
-                            self.command_task_mapping[command.command_id] = task_id
-
-                            if EXO_TRACING_ENABLED:
-                                selected_instance = self.state.instances.get(
-                                    selected_instance_id
-                                )
-                                if selected_instance:
-                                    ranks = set(
-                                        shard.device_rank
-                                        for shard in selected_instance.shard_assignments.runner_to_shard.values()
-                                    )
-                                    self._expected_ranks[task_id] = ranks
-                        case ImageEdits():
-                            for instance in self.state.instances.values():
-                                if (
-                                    instance.shard_assignments.model_id
-                                    == command.task_params.model
-                                ):
-                                    in_flight = {TaskStatus.Pending, TaskStatus.Running}
-                                    task_count = sum(
-                                        1
-                                        for task in self.state.tasks.values()
-                                        if task.instance_id == instance.instance_id
-                                        and task.task_status in in_flight
-                                    )
-                                    instance_task_counts[instance.instance_id] = (
-                                        task_count
-                                    )
-
-                            if not instance_task_counts:
-                                raise ValueError(
-                                    f"No instance found for model {command.task_params.model}"
-                                )
-
-                            available_instance_ids = sorted(
-                                instance_task_counts.keys(),
-                                key=lambda instance_id: instance_task_counts[
-                                    instance_id
-                                ],
-                            )
-
-                            task_id = TaskId()
-                            selected_instance_id = available_instance_ids[0]
-                            generated_events.append(
-                                TaskCreated(
-                                    task_id=task_id,
-                                    task=ImageEditsTask(
-                                        task_id=task_id,
-                                        command_id=command.command_id,
-                                        instance_id=selected_instance_id,
-                                        task_status=TaskStatus.Pending,
-                                        task_params=command.task_params,
-                                    ),
-                                )
-                            )
-
-                            self.command_task_mapping[command.command_id] = task_id
-
-                            if EXO_TRACING_ENABLED:
-                                selected_instance = self.state.instances.get(
-                                    selected_instance_id
-                                )
-                                if selected_instance:
-                                    ranks = set(
-                                        shard.device_rank
-                                        for shard in selected_instance.shard_assignments.runner_to_shard.values()
-                                    )
-                                    self._expected_ranks[task_id] = ranks
-                        case DeleteInstance():
-                            placement = delete_instance(command, self.state.instances)
-                            transition_events = get_transition_events(
-                                self.state.instances, placement, self.state.tasks
-                            )
-                            for cmd in cancel_unnecessary_downloads(
-                                placement, self.state.downloads
+                        for instance in self.state.instances.values():
+                            if (
+                                instance.shard_assignments.model_id
+                                == command.task_params.model
+                                and instance.instance_id not in prefill_only
                             ):
-                                await self.download_command_sender.send(
-                                    ForwarderDownloadCommand(
-                                        origin=self._system_id, command=cmd
-                                    )
-                                )
-                            generated_events.extend(transition_events)
-                        case PlaceInstance():
-                            placement = place_instance(
-                                command,
-                                self.state.topology,
-                                self.state.instances,
-                                self.state.node_memory,
-                                self.state.node_network,
-                                download_status=self.state.downloads,
-                            )
-                            transition_events = get_transition_events(
-                                self.state.instances, placement, self.state.tasks
-                            )
-                            generated_events.extend(transition_events)
-                        case CreateInstance():
-                            placement = add_instance_to_placements(
-                                command,
-                                self.state.topology,
-                                self.state.instances,
-                            )
-                            transition_events = get_transition_events(
-                                self.state.instances, placement, self.state.tasks
-                            )
-                            generated_events.extend(transition_events)
-                        case SendInputChunk(chunk=chunk):
-                            generated_events.append(
-                                InputChunkReceived(
-                                    command_id=chunk.command_id,
-                                    chunk=chunk,
-                                )
-                            )
-                        case TaskCancelled():
-                            if (
-                                task_id := self.command_task_mapping.get(
-                                    command.cancelled_command_id
-                                )
-                            ) is not None:
-                                generated_events.append(
-                                    TaskStatusUpdated(
-                                        task_status=TaskStatus.Cancelled,
-                                        task_id=task_id,
-                                    )
-                                )
-                            else:
-                                logger.warning(
-                                    f"Nonexistent command {command.cancelled_command_id} cancelled"
-                                )
-                        case TaskFinished():
-                            if (
-                                task_id := self.command_task_mapping.pop(
-                                    command.finished_command_id, None
-                                )
-                            ) is not None:
-                                generated_events.append(TaskDeleted(task_id=task_id))
-                            else:
-                                logger.warning(
-                                    f"Finished command {command.finished_command_id} finished"
+                                in_flight = {TaskStatus.Pending, TaskStatus.Running}
+                                task_count = sum(
+                                    1
+                                    for task in self.state.tasks.values()
+                                    if task.instance_id == instance.instance_id
+                                    and task.task_status in in_flight
                                 )
+                                instance_task_counts[instance.instance_id] = task_count
 
-                        case AddCustomModelCard():
-                            generated_events.append(
-                                CustomModelCardAdded(model_card=command.model_card)
+                        if not instance_task_counts:
+                            raise ValueError(
+                                f"No instance found for model {command.task_params.model}"
                             )
-                        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)
+
+                        available_instance_ids = sorted(
+                            instance_task_counts.keys(),
+                            key=lambda instance_id: instance_task_counts[instance_id],
+                        )
+
+                        decode_instance_id = available_instance_ids[0]
+                        task_id = TaskId()
+                        params = command.task_params.model_copy(
+                            update={
+                                "prefill_endpoint": _prefill_endpoint_for(
+                                    self.state, decode_instance_id
                                 ),
-                                decode_instances=list(
-                                    dict.fromkeys(command.decode_instances)
+                            }
+                        )
+                        generated_events.append(
+                            TaskCreated(
+                                task_id=task_id,
+                                task=TextGenerationTask(
+                                    task_id=task_id,
+                                    command_id=command.command_id,
+                                    instance_id=decode_instance_id,
+                                    task_status=TaskStatus.Pending,
+                                    task_params=params,
                                 ),
                             )
-                            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
-                            end = min(
-                                command.since_idx + _MAX_EVENT_LOG_REPLAY_BATCH,
-                                len(self._event_log),
-                            )
-                            for i, event in enumerate(
-                                self._event_log.read_range(command.since_idx, end),
-                                start=command.since_idx,
+                        )
+                        self.command_task_mapping[command.command_id] = task_id
+                    case ImageGeneration():
+                        for instance in self.state.instances.values():
+                            if (
+                                instance.shard_assignments.model_id
+                                == command.task_params.model
                             ):
-                                await self._send_event(IndexedEvent(idx=i, event=event))
-                        case RequestSnapshot():
-                            self._tg.start_soon(
-                                self._serve_snapshot, command.requester_node_id
+                                in_flight = {TaskStatus.Pending, TaskStatus.Running}
+                                task_count = sum(
+                                    1
+                                    for task in self.state.tasks.values()
+                                    if task.instance_id == instance.instance_id
+                                    and task.task_status in in_flight
+                                )
+                                instance_task_counts[instance.instance_id] = task_count
+
+                        if not instance_task_counts:
+                            raise ValueError(
+                                f"No instance found for model {command.task_params.model}"
                             )
-                    for event in generated_events:
-                        await self.event_sender.send(event)
-                except ValueError as e:
-                    logger.opt(exception=e).warning("Error in command processor")
+
+                        available_instance_ids = sorted(
+                            instance_task_counts.keys(),
+                            key=lambda instance_id: instance_task_counts[instance_id],
+                        )
+
+                        task_id = TaskId()
+                        selected_instance_id = available_instance_ids[0]
+                        generated_events.append(
+                            TaskCreated(
+                                task_id=task_id,
+                                task=ImageGenerationTask(
+                                    task_id=task_id,
+                                    command_id=command.command_id,
+                                    instance_id=selected_instance_id,
+                                    task_status=TaskStatus.Pending,
+                                    task_params=command.task_params,
+                                ),
+                            )
+                        )
+
+                        self.command_task_mapping[command.command_id] = task_id
+
+                        if EXO_TRACING_ENABLED:
+                            selected_instance = self.state.instances.get(
+                                selected_instance_id
+                            )
+                            if selected_instance:
+                                ranks = set(
+                                    shard.device_rank
+                                    for shard in selected_instance.shard_assignments.runner_to_shard.values()
+                                )
+                                self._expected_ranks[task_id] = ranks
+                    case ImageEdits():
+                        for instance in self.state.instances.values():
+                            if (
+                                instance.shard_assignments.model_id
+                                == command.task_params.model
+                            ):
+                                in_flight = {TaskStatus.Pending, TaskStatus.Running}
+                                task_count = sum(
+                                    1
+                                    for task in self.state.tasks.values()
+                                    if task.instance_id == instance.instance_id
+                                    and task.task_status in in_flight
+                                )
+                                instance_task_counts[instance.instance_id] = task_count
+
+                        if not instance_task_counts:
+                            raise ValueError(
+                                f"No instance found for model {command.task_params.model}"
+                            )
+
+                        available_instance_ids = sorted(
+                            instance_task_counts.keys(),
+                            key=lambda instance_id: instance_task_counts[instance_id],
+                        )
+
+                        task_id = TaskId()
+                        selected_instance_id = available_instance_ids[0]
+                        generated_events.append(
+                            TaskCreated(
+                                task_id=task_id,
+                                task=ImageEditsTask(
+                                    task_id=task_id,
+                                    command_id=command.command_id,
+                                    instance_id=selected_instance_id,
+                                    task_status=TaskStatus.Pending,
+                                    task_params=command.task_params,
+                                ),
+                            )
+                        )
+
+                        self.command_task_mapping[command.command_id] = task_id
+
+                        if EXO_TRACING_ENABLED:
+                            selected_instance = self.state.instances.get(
+                                selected_instance_id
+                            )
+                            if selected_instance:
+                                ranks = set(
+                                    shard.device_rank
+                                    for shard in selected_instance.shard_assignments.runner_to_shard.values()
+                                )
+                                self._expected_ranks[task_id] = ranks
+                    case DeleteInstance():
+                        placement = delete_instance(command, self.state.instances)
+                        transition_events = get_transition_events(
+                            self.state.instances, placement, self.state.tasks
+                        )
+                        for cmd in cancel_unnecessary_downloads(
+                            placement, self.state.downloads
+                        ):
+                            await self.download_command_sender.send(
+                                ForwarderDownloadCommand(
+                                    origin=self._system_id, command=cmd
+                                )
+                            )
+                        generated_events.extend(transition_events)
+                    case PlaceInstance():
+                        placement = place_instance(
+                            command,
+                            self.state.topology,
+                            self.state.instances,
+                            self.state.node_memory,
+                            self.state.node_network,
+                            download_status=self.state.downloads,
+                            node_rdma_ctl=self.state.node_rdma_ctl,
+                        )
+                        transition_events = get_transition_events(
+                            self.state.instances, placement, self.state.tasks
+                        )
+                        generated_events.extend(transition_events)
+                    case CreateInstance():
+                        placement = add_instance_to_placements(
+                            command,
+                            self.state.topology,
+                            self.state.instances,
+                        )
+                        transition_events = get_transition_events(
+                            self.state.instances, placement, self.state.tasks
+                        )
+                        generated_events.extend(transition_events)
+                    case SendInputChunk(chunk=chunk):
+                        generated_events.append(
+                            InputChunkReceived(
+                                command_id=chunk.command_id,
+                                chunk=chunk,
+                            )
+                        )
+                    case TaskCancelled():
+                        if (
+                            task_id := self.command_task_mapping.get(
+                                command.cancelled_command_id
+                            )
+                        ) is not None:
+                            generated_events.append(
+                                TaskStatusUpdated(
+                                    task_status=TaskStatus.Cancelled,
+                                    task_id=task_id,
+                                )
+                            )
+                        else:
+                            logger.warning(
+                                f"Nonexistent command {command.cancelled_command_id} cancelled"
+                            )
+                    case TaskFinished():
+                        if (
+                            task_id := self.command_task_mapping.pop(
+                                command.finished_command_id, None
+                            )
+                        ) is not None:
+                            generated_events.append(TaskDeleted(task_id=task_id))
+                        else:
+                            logger.warning(
+                                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():
+                        end = len(self._event_log)
+                        for i, event in enumerate(
+                            self._event_log.read_range(command.since_idx, end),
+                            start=command.since_idx,
+                        ):
+                            await self._send_event(IndexedEvent(idx=i, event=event))
+                for event in generated_events:
+                    await self.event_sender.send(event)
+            except ValueError as e:
+                logger.opt(exception=e).warning("Error in command processor")
 
     # These plan loops are the cracks showing in our event sourcing architecture - more things could be commands
     async def _plan(self) -> None:
@@ -529,42 +496,6 @@ class Master:
                     self._event_log.append(event)
                     await self._send_event(indexed)
 
-    async def _serve_snapshot(self, requester_node_id: NodeId) -> None:
-        state = self.state
-        if state.last_event_applied_idx < 0:
-            logger.info(
-                f"RequestSnapshot from {requester_node_id} but master has no events yet"
-            )
-            return
-
-        body = await to_thread.run_sync(_encode_state_for_transfer, state)
-        sha256 = hashlib.sha256(body).hexdigest()
-        chunks = [
-            body[i : i + _SNAPSHOT_CHUNK_BYTES]
-            for i in range(0, len(body), _SNAPSHOT_CHUNK_BYTES)
-        ] or [b""]
-        transfer_id = SnapshotTransferId()
-
-        logger.info(
-            f"Serving snapshot to {requester_node_id}: "
-            f"idx={state.last_event_applied_idx}, "
-            f"{len(chunks)} chunk(s), {len(body)} bytes total"
-        )
-        for index, chunk in enumerate(chunks):
-            await self.snapshot_chunk_sender.send(
-                SnapshotChunk.from_data(
-                    data=chunk,
-                    transfer_id=transfer_id,
-                    requester_node_id=requester_node_id,
-                    session_id=self.session_id,
-                    schema_version=state.schema_version,
-                    last_event_applied_idx=state.last_event_applied_idx,
-                    chunk_index=index,
-                    total_chunks=len(chunks),
-                    sha256_hex=sha256,
-                )
-            )
-
     # This function is re-entrant, take care!
     async def _send_event(self, event: IndexedEvent):
         # Convenience method since this line is ugly

src/exo/master/placement.py

@@ -28,7 +28,7 @@ from exo.shared.types.events import (
     TaskStatusUpdated,
 )
 from exo.shared.types.memory import Memory
-from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
+from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo, NodeRdmaCtlStatus
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import (
     DownloadCompleted,
@@ -105,6 +105,7 @@ def place_instance(
     node_network: Mapping[NodeId, NodeNetworkInfo],
     required_nodes: set[NodeId] | None = None,
     download_status: Mapping[NodeId, Sequence[DownloadProgress]] | None = None,
+    node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus] | None = None,
 ) -> dict[InstanceId, Instance]:
     cycles = topology.get_cycles()
     candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles))
@@ -166,8 +167,18 @@ def place_instance(
 
     smallest_cycles = get_smallest_cycles(cycles_with_sufficient_memory)
 
+    rdma_ctl_status = node_rdma_ctl or {}
+
+    def _all_rdma_ctl_enabled(cycle: Cycle) -> bool:
+        return all(
+            ((status := rdma_ctl_status.get(node_id)) is not None and status.enabled)
+            for node_id in cycle
+        )
+
     smallest_rdma_cycles = [
-        cycle for cycle in smallest_cycles if topology.is_rdma_cycle(cycle)
+        cycle
+        for cycle in smallest_cycles
+        if topology.is_rdma_cycle(cycle) and _all_rdma_ctl_enabled(cycle)
     ]
 
     if command.instance_meta == InstanceMeta.MlxJaccl:

src/exo/master/tests/test_master.py

@@ -6,15 +6,12 @@ import pytest
 from loguru import logger
 
 from exo.master.main import Master
-from exo.routing.router import get_node_id_keypair
-from exo.routing.snapshot_receiver import SnapshotReceiver
 from exo.shared.models.model_cards import ModelCard, ModelTask
 from exo.shared.types.commands import (
     CommandId,
     ForwarderCommand,
     ForwarderDownloadCommand,
     PlaceInstance,
-    RequestSnapshot,
     TextGeneration,
 )
 from exo.shared.types.common import ModelId, NodeId, SessionId, SystemId
@@ -31,7 +28,6 @@ from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import (
     MemoryUsage,
 )
-from exo.shared.types.snapshots import SnapshotChunk
 from exo.shared.types.tasks import TaskStatus
 from exo.shared.types.tasks import TextGeneration as TextGenerationTask
 from exo.shared.types.text_generation import (
@@ -50,8 +46,7 @@ from exo.utils.channels import channel
 
 @pytest.mark.asyncio
 async def test_master():
-    keypair = get_node_id_keypair()
-    node_id = NodeId(keypair.to_node_id())
+    node_id = NodeId("master test")
     session_id = SessionId(master_node_id=node_id, election_clock=0)
 
     ge_sender, global_event_receiver = channel[GlobalForwarderEvent]()
@@ -59,7 +54,6 @@ async def test_master():
     local_event_sender, le_receiver = channel[LocalForwarderEvent]()
     fcds, _fcdr = channel[ForwarderDownloadCommand]()
     ev_send, ev_recv = channel[Event]()
-    snapshot_chunk_send, _snapshot_chunk_recv = channel[SnapshotChunk]()
 
     async def mock_event_router():
         idx = 0
@@ -96,7 +90,6 @@ async def test_master():
         global_event_sender=ge_sender,
         local_event_receiver=le_receiver,
         command_receiver=co_receiver,
-        snapshot_chunk_sender=snapshot_chunk_send,
         download_command_sender=fcds,
     )
     logger.info("run the master")
@@ -234,52 +227,3 @@ async def test_master():
 
         ev_send.close()
         await master.shutdown()
-
-
-@pytest.mark.asyncio
-async def test_master_serves_snapshot_for_current_state():
-    node_id = NodeId("master")
-    requester_node_id = NodeId("worker")
-    session_id = SessionId(master_node_id=node_id, election_clock=0)
-
-    ge_sender, _global_event_receiver = channel[GlobalForwarderEvent]()
-    command_sender, command_receiver = channel[ForwarderCommand]()
-    _local_event_sender, local_event_receiver = channel[LocalForwarderEvent]()
-    download_command_sender, _download_command_receiver = channel[
-        ForwarderDownloadCommand
-    ]()
-    event_sender, _event_receiver = channel[Event]()
-    snapshot_chunk_sender, snapshot_chunk_receiver = channel[SnapshotChunk]()
-
-    master = Master(
-        node_id,
-        session_id,
-        event_sender=event_sender,
-        global_event_sender=ge_sender,
-        local_event_receiver=local_event_receiver,
-        command_receiver=command_receiver,
-        snapshot_chunk_sender=snapshot_chunk_sender,
-        download_command_sender=download_command_sender,
-    )
-    master.state = master.state.model_copy(update={"last_event_applied_idx": 12})
-
-    receiver = SnapshotReceiver(requester_node_id, session_id)
-    async with anyio.create_task_group() as tg:
-        tg.start_soon(master.run)
-        await command_sender.send(
-            ForwarderCommand(
-                origin=SystemId("api"),
-                command=RequestSnapshot(requester_node_id=requester_node_id),
-            )
-        )
-
-        received = None
-        while received is None:
-            chunk = await snapshot_chunk_receiver.receive()
-            received = receiver.ingest(chunk)
-
-        assert received.last_event_applied_idx == 12
-        assert received.state.last_event_applied_idx == 12
-
-        await master.shutdown()
-        tg.cancel_scope.cancel()

src/exo/master/tests/test_placement.py

@@ -21,7 +21,11 @@ from exo.shared.types.events import (
 )
 from exo.shared.types.memory import Memory
 from exo.shared.types.multiaddr import Multiaddr
-from exo.shared.types.profiling import NetworkInterfaceInfo, NodeNetworkInfo
+from exo.shared.types.profiling import (
+    NetworkInterfaceInfo,
+    NodeNetworkInfo,
+    NodeRdmaCtlStatus,
+)
 from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
 from exo.shared.types.text_generation import (
     InputMessage,
@@ -439,8 +443,21 @@ def test_tensor_rdma_backend_connectivity_matrix(
         min_nodes=1,
     )
 
+    node_rdma_ctl = {
+        node_a: NodeRdmaCtlStatus(enabled=True),
+        node_b: NodeRdmaCtlStatus(enabled=True),
+        node_c: NodeRdmaCtlStatus(enabled=True),
+    }
+
     # act
-    placements = place_instance(cic, topology, {}, node_memory, node_network)
+    placements = place_instance(
+        cic,
+        topology,
+        {},
+        node_memory,
+        node_network,
+        node_rdma_ctl=node_rdma_ctl,
+    )
 
     # assert
     assert len(placements) == 1
@@ -482,6 +499,131 @@ def test_tensor_rdma_backend_connectivity_matrix(
             assert len(ip_part.split(".")) == 4
 
 
+def _build_three_node_rdma_topology() -> tuple[
+    Topology, NodeId, NodeId, NodeId, dict[NodeId, NodeNetworkInfo]
+]:
+    topology = Topology()
+    node_a = NodeId()
+    node_b = NodeId()
+    node_c = NodeId()
+
+    ethernet_interface = NetworkInterfaceInfo(name="en0", ip_address="10.0.0.1")
+    ethernet_conn = SocketConnection(
+        sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000")
+    )
+    node_network = {
+        node_a: NodeNetworkInfo(interfaces=[ethernet_interface]),
+        node_b: NodeNetworkInfo(interfaces=[ethernet_interface]),
+        node_c: NodeNetworkInfo(interfaces=[ethernet_interface]),
+    }
+
+    for n in (node_a, node_b, node_c):
+        topology.add_node(n)
+
+    rdma_pairs = [
+        (node_a, node_b, 3),
+        (node_b, node_a, 3),
+        (node_b, node_c, 4),
+        (node_c, node_b, 4),
+        (node_a, node_c, 5),
+        (node_c, node_a, 5),
+    ]
+    for src, sink, iface in rdma_pairs:
+        topology.add_connection(
+            Connection(source=src, sink=sink, edge=create_rdma_connection(iface))
+        )
+
+    socket_pairs = [
+        (node_a, node_b),
+        (node_b, node_c),
+        (node_c, node_a),
+        (node_a, node_c),
+        (node_b, node_a),
+        (node_c, node_b),
+    ]
+    for src, sink in socket_pairs:
+        topology.add_connection(Connection(source=src, sink=sink, edge=ethernet_conn))
+
+    return topology, node_a, node_b, node_c, node_network
+
+
+def test_place_mlx_jaccl_rejects_when_a_node_has_rdma_ctl_disabled(
+    model_card: ModelCard,
+):
+    # arrange
+    model_card = model_card.model_copy(
+        update={"n_layers": 12, "storage_size": Memory.from_bytes(1500)}
+    )
+    topology, node_a, node_b, node_c, node_network = _build_three_node_rdma_topology()
+    node_memory = {
+        node_a: create_node_memory(500),
+        node_b: create_node_memory(500),
+        node_c: create_node_memory(500),
+    }
+    node_rdma_ctl = {
+        node_a: NodeRdmaCtlStatus(enabled=True),
+        node_b: NodeRdmaCtlStatus(enabled=True),
+        node_c: NodeRdmaCtlStatus(enabled=False),
+    }
+    cic = PlaceInstance(
+        sharding=Sharding.Tensor,
+        instance_meta=InstanceMeta.MlxJaccl,
+        command_id=CommandId(),
+        model_card=model_card,
+        min_nodes=3,
+    )
+
+    # act / assert
+    with pytest.raises(
+        ValueError, match="Requested RDMA \\(MlxJaccl\\) but no RDMA-connected cycles"
+    ):
+        place_instance(
+            cic,
+            topology,
+            {},
+            node_memory,
+            node_network,
+            node_rdma_ctl=node_rdma_ctl,
+        )
+
+
+def test_place_mlx_jaccl_rejects_when_node_rdma_ctl_missing(model_card: ModelCard):
+    """A node with no observed rdma_ctl status must not participate in RDMA placement."""
+    # arrange
+    model_card = model_card.model_copy(
+        update={"n_layers": 12, "storage_size": Memory.from_bytes(1500)}
+    )
+    topology, node_a, node_b, node_c, node_network = _build_three_node_rdma_topology()
+    node_memory = {
+        node_a: create_node_memory(500),
+        node_b: create_node_memory(500),
+        node_c: create_node_memory(500),
+    }
+    # node_c has no rdma_ctl entry at all
+    node_rdma_ctl = {
+        node_a: NodeRdmaCtlStatus(enabled=True),
+        node_b: NodeRdmaCtlStatus(enabled=True),
+    }
+    cic = PlaceInstance(
+        sharding=Sharding.Tensor,
+        instance_meta=InstanceMeta.MlxJaccl,
+        command_id=CommandId(),
+        model_card=model_card,
+        min_nodes=3,
+    )
+
+    # act / assert
+    with pytest.raises(ValueError):
+        place_instance(
+            cic,
+            topology,
+            {},
+            node_memory,
+            node_network,
+            node_rdma_ctl=node_rdma_ctl,
+        )
+
+
 def _make_task(
     instance_id: InstanceId,
     status: TaskStatus = TaskStatus.Running,

src/exo/routing/connection_message.py

@@ -1,15 +1,13 @@
-from exo_pyo3_bindings import PyFromSwarm
+from exo_net import PyFromSwarm
 
-from exo.shared.types.common import NodeId
 from exo.utils.pydantic_ext import FrozenModel
 
 """Serialisable types for Connection Updates/Messages"""
 
 
 class ConnectionMessage(FrozenModel):
-    node_id: NodeId
     connected: bool
 
     @classmethod
     def from_update(cls, update: PyFromSwarm.Connection) -> "ConnectionMessage":
-        return cls(node_id=NodeId(update.peer_id), connected=update.connected)
+        return cls(connected=update.connected)

src/exo/routing/event_router.py

@@ -4,9 +4,10 @@ from random import random
 import anyio
 from anyio import BrokenResourceError, ClosedResourceError
 from anyio.abc import CancelScope
+from exo_net import NetSender
 from loguru import logger
 
-from exo.shared.types.commands import ForwarderCommand, RequestEventLog
+from exo.shared.types.commands import RequestEventLog
 from exo.shared.types.common import SessionId, SystemId
 from exo.shared.types.events import (
     Event,
@@ -23,7 +24,7 @@ from exo.utils.task_group import TaskGroup
 @dataclass
 class EventRouter:
     session_id: SessionId
-    command_sender: Sender[ForwarderCommand]
+    command_sender: NetSender
     external_inbound: Receiver[GlobalForwarderEvent]
     external_outbound: Sender[LocalForwarderEvent]
     _system_id: SystemId = field(init=False, default_factory=SystemId)
@@ -80,9 +81,6 @@ class EventRouter:
     def shutdown(self) -> None:
         self._tg.cancel_tasks()
 
-    def set_buffer_start(self, idx: int) -> None:
-        self.event_buffer.fast_forward_to(idx)
-
     async def _ingest(self, system_id: SystemId, recv: Receiver[Event]):
         idx = 0
         with recv as events:
@@ -98,6 +96,7 @@ class EventRouter:
                 self.out_for_delivery[event.event_id] = (anyio.current_time(), f_ev)
 
     async def _run_ext_in(self):
+        buf = OrderedBuffer[Event]()
         with self.external_inbound as events:
             async for event in events:
                 if event.session != self.session_id:
@@ -105,12 +104,12 @@ class EventRouter:
                 if event.origin != self.session_id.master_node_id:
                     continue
 
-                self.event_buffer.ingest(event.origin_idx, event.event)
+                buf.ingest(event.origin_idx, event.event)
                 event_id = event.event.event_id
                 if event_id in self.out_for_delivery:
                     self.out_for_delivery.pop(event_id)
 
-                drained = self.event_buffer.drain_indexed()
+                drained = buf.drain_indexed()
                 if drained:
                     self._nack_attempts = 0
                     if self._nack_cancel_scope:
@@ -121,9 +120,7 @@ class EventRouter:
                     or self._nack_cancel_scope.cancel_called
                 ):
                     # Request the next index.
-                    self._tg.start_soon(
-                        self._nack_request, self.event_buffer.next_idx_to_release
-                    )
+                    self._tg.start_soon(self._nack_request, buf.next_idx_to_release)
                     continue
 
                 for idx, event in drained:
@@ -156,10 +153,9 @@ class EventRouter:
                     f"Nack attempt {self._nack_attempts}: Requesting Event Log from {since_idx}"
                 )
                 await self.command_sender.send(
-                    ForwarderCommand(
-                        origin=self._system_id,
-                        command=RequestEventLog(since_idx=since_idx),
-                    )
+                    RequestEventLog(since_idx=since_idx)
+                    .model_dump_json()
+                    .encode("utf-8")
                 )
             finally:
                 if self._nack_cancel_scope is scope:

src/exo/routing/router.py

@@ -2,8 +2,6 @@ from collections.abc import Sequence
 from copy import copy
 from itertools import count
 from math import inf
-from os import PathLike
-from pathlib import Path
 from typing import cast
 
 from anyio import (
@@ -12,18 +10,9 @@ from anyio import (
     move_on_after,
     sleep_forever,
 )
-from exo_pyo3_bindings import (
-    AllQueuesFullError,
-    Keypair,
-    MessageTooLargeError,
-    NetworkingHandle,
-    NoPeersSubscribedToTopicError,
-    PyFromSwarm,
-)
-from filelock import FileLock
+from exo_net import NetworkingHandle, PyFromSwarm, PySession
 from loguru import logger
 
-from exo.shared.constants import EXO_NODE_ID_KEYPAIR
 from exo.utils.channels import Receiver, Sender, channel
 from exo.utils.pydantic_ext import FrozenModel
 from exo.utils.task_group import TaskGroup
@@ -105,13 +94,14 @@ class Router:
     @classmethod
     def create(
         cls,
-        identity: Keypair,
+        identity: bytes,
         bootstrap_peers: Sequence[str] = (),
         listen_port: int = 0,
-    ) -> "Router":
-        return cls(
-            handle=NetworkingHandle(identity, list(bootstrap_peers), listen_port)
+    ) -> "tuple[Router, PySession]":
+        handle, session = NetworkingHandle.new(
+            identity, list(bootstrap_peers), listen_port
         )
+        return cls(handle=handle), session
 
     def __init__(self, handle: NetworkingHandle):
         self.topic_routers: dict[str, TopicRouter[FrozenModel]] = {}
@@ -191,10 +181,8 @@ class Router:
                 from_swarm = await self._net.recv()
                 logger.debug(from_swarm)
                 match from_swarm:
-                    case PyFromSwarm.Message(origin, topic, data):
-                        logger.trace(
-                            f"Received message on {topic} from {origin} with payload {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(
                                 f"Received message on unknown or inactive topic {topic}"
@@ -225,51 +213,9 @@ class Router:
     async def _networking_publish(self):
         with self.networking_receiver as networked_items:
             async for topic, data in networked_items:
-                try:
-                    logger.trace(f"Sending message on {topic} with payload {data}")
-                    if len(data) > 1024 * 1024:
-                        logger.warning(
-                            "Sending overlarge payload, network performance may be temporarily degraded"
-                        )
-                    await self._net.gossipsub_publish(topic, data)
-                except NoPeersSubscribedToTopicError:
-                    pass
-                except AllQueuesFullError:
-                    logger.warning(f"All peer queues full, dropping message on {topic}")
-                except MessageTooLargeError:
+                logger.trace(f"Sending message on {topic} with payload {data}")
+                if len(data) > 1024 * 1024:
                     logger.warning(
-                        f"Message too large for gossipsub on {topic} ({len(data)} bytes), dropping"
+                        "Sending overlarge payload, network performance may be temporarily degraded"
                     )
-
-
-def get_node_id_keypair(
-    path: str | bytes | PathLike[str] | PathLike[bytes] = EXO_NODE_ID_KEYPAIR,
-) -> Keypair:
-    """
-    Obtains the :class:`Keypair` associated with this node-ID.
-    Obtain the :class:`PeerId` by from it.
-    """
-    # TODO(evan): bring back node id persistence once we figure out how to deal with duplicates
-    return Keypair.generate()
-
-    def lock_path(path: str | bytes | PathLike[str] | PathLike[bytes]) -> Path:
-        return Path(str(path) + ".lock")
-
-    # operate with cross-process lock to avoid race conditions
-    with FileLock(lock_path(path)):
-        with open(path, "a+b") as f:  # opens in append-mode => starts at EOF
-            # if non-zero EOF, then file exists => use to get node-ID
-            if f.tell() != 0:
-                f.seek(0)  # go to start & read protobuf-encoded bytes
-                protobuf_encoded = f.read()
-
-                try:  # if decoded successfully, save & return
-                    return Keypair.from_bytes(protobuf_encoded)
-                except ValueError as e:  # on runtime error, assume corrupt file
-                    logger.warning(f"Encountered error when trying to get keypair: {e}")
-
-        # if no valid credentials, create new ones and persist
-        with open(path, "w+b") as f:
-            keypair = Keypair.generate()
-            f.write(keypair.to_bytes())
-            return keypair
+                await self._net.gossipsub_publish(topic, data)

src/exo/routing/snapshot_receiver.py

@@ -1,109 +0,0 @@
-"""Reassembles a snapshot from a stream of `SnapshotChunk`s.
-
-A receiver belongs to one node; it ignores chunks addressed to other
-requesters and chunks from prior sessions. Once a transfer's chunks have
-all been collected and the SHA-256 checks out, the snapshot is decoded into
-a `State`. Concurrent transfers (for the same requester) are tolerated:
-each is keyed by `transfer_id`.
-"""
-
-from __future__ import annotations
-
-import hashlib
-from dataclasses import dataclass, field
-from typing import final
-
-import zstandard
-from loguru import logger
-
-from exo.shared.types.common import NodeId, SessionId
-from exo.shared.types.snapshots import SnapshotChunk, SnapshotTransferId
-from exo.shared.types.state import State
-
-
-@final
-@dataclass
-class _Assembly:
-    """Partial state for one in-flight snapshot transfer."""
-
-    total_chunks: int
-    sha256_hex: str
-    schema_version: int
-    last_event_applied_idx: int
-    chunks: dict[int, bytes] = field(default_factory=dict)
-
-    def is_complete(self) -> bool:
-        return len(self.chunks) == self.total_chunks
-
-    def assemble(self) -> bytes:
-        return b"".join(self.chunks[i] for i in range(self.total_chunks))
-
-
-@dataclass
-class ReceivedSnapshot:
-    last_event_applied_idx: int
-    state: State
-
-
-class SnapshotReceiver:
-    """Filters and reassembles inbound chunks into a `ReceivedSnapshot`.
-
-    Stateless w.r.t. delivery: callers feed `SnapshotChunk`s in via `ingest`
-    and check the return value for completion.
-    """
-
-    def __init__(self, my_node_id: NodeId, session_id: SessionId) -> None:
-        self._my_node_id = my_node_id
-        self._session_id = session_id
-        self._assemblies: dict[SnapshotTransferId, _Assembly] = {}
-
-    def ingest(self, chunk: SnapshotChunk) -> ReceivedSnapshot | None:
-        """Absorb a chunk; return the snapshot once a transfer completes.
-
-        Returns None for partial transfers, mismatched recipients, stale
-        sessions, version mismatches, or corrupt payloads.
-        """
-        if chunk.requester_node_id != self._my_node_id:
-            return None
-        if chunk.session_id != self._session_id:
-            return None
-
-        existing = self._assemblies.get(chunk.transfer_id)
-        if existing is None:
-            existing = _Assembly(
-                total_chunks=chunk.total_chunks,
-                sha256_hex=chunk.sha256_hex,
-                schema_version=chunk.schema_version,
-                last_event_applied_idx=chunk.last_event_applied_idx,
-            )
-            self._assemblies[chunk.transfer_id] = existing
-        existing.chunks[chunk.chunk_index] = chunk.data
-
-        if not existing.is_complete():
-            return None
-
-        # Transfer complete — finalise and remove from the in-flight map.
-        del self._assemblies[chunk.transfer_id]
-        body = existing.assemble()
-        if hashlib.sha256(body).hexdigest() != existing.sha256_hex:
-            logger.warning(f"Snapshot {chunk.transfer_id} failed checksum; discarding")
-            return None
-        try:
-            decompressed = zstandard.ZstdDecompressor().decompress(body)
-            state = State.model_validate_json(decompressed.decode("utf-8"))
-        except (zstandard.ZstdError, ValueError) as e:
-            logger.opt(exception=e).warning(
-                f"Snapshot {chunk.transfer_id} could not be decoded; discarding"
-            )
-            return None
-        if state.schema_version != existing.schema_version:
-            # Should not happen — the master writes schema_version into both
-            # the chunk meta and the State payload — but treat it as corrupt.
-            logger.warning(
-                f"Snapshot {chunk.transfer_id} schema version mismatch "
-                f"(chunk={existing.schema_version}, state={state.schema_version})"
-            )
-            return None
-        return ReceivedSnapshot(
-            last_event_applied_idx=existing.last_event_applied_idx, state=state
-        )

src/exo/routing/tests/test_event_buffer.py

@@ -141,28 +141,3 @@ async def test_drain_and_ingest_with_new_sequence(buffer: OrderedBuffer[Event]):
     assert [e[0] for e in drained] == [2]
     assert buffer.next_idx_to_release == 3
     assert 4 in buffer.store
-
-
-@pytest.mark.asyncio
-async def test_fast_forward_discards_buffered_stale_events(
-    buffer: OrderedBuffer[Event],
-):
-    buffer.ingest(*make_indexed_event(0))
-    buffer.ingest(*make_indexed_event(2))
-    buffer.ingest(*make_indexed_event(4))
-
-    buffer.fast_forward_to(3)
-
-    assert buffer.next_idx_to_release == 3
-    assert set(buffer.store) == {4}
-
-
-@pytest.mark.asyncio
-async def test_fast_forward_only_moves_forward(buffer: OrderedBuffer[Event]):
-    buffer.ingest(*make_indexed_event(0))
-    buffer.ingest(*make_indexed_event(1))
-    buffer.drain()
-
-    buffer.fast_forward_to(1)
-
-    assert buffer.next_idx_to_release == 2

src/exo/routing/tests/test_snapshot_receiver.py

@@ -1,151 +0,0 @@
-import hashlib
-
-import pytest
-import zstandard
-
-from exo.routing.snapshot_receiver import SnapshotReceiver
-from exo.shared.types.common import NodeId, SessionId
-from exo.shared.types.snapshots import SnapshotChunk, SnapshotTransferId
-from exo.shared.types.state import State
-
-
-@pytest.fixture
-def session_id() -> SessionId:
-    return SessionId(master_node_id=NodeId("master"), election_clock=0)
-
-
-@pytest.fixture
-def my_node() -> NodeId:
-    return NodeId("worker-1")
-
-
-def _encode(state: State) -> bytes:
-    return zstandard.ZstdCompressor().compress(state.model_dump_json().encode("utf-8"))
-
-
-def _make_chunks(
-    body: bytes,
-    *,
-    chunk_size: int,
-    requester_node_id: NodeId,
-    session_id: SessionId,
-    state: State,
-    transfer_id: SnapshotTransferId | None = None,
-) -> list[SnapshotChunk]:
-    sha256 = hashlib.sha256(body).hexdigest()
-    transfer_id = transfer_id or SnapshotTransferId()
-    pieces = [body[i : i + chunk_size] for i in range(0, len(body), chunk_size)] or [
-        b""
-    ]
-    return [
-        SnapshotChunk.from_data(
-            data=piece,
-            transfer_id=transfer_id,
-            requester_node_id=requester_node_id,
-            session_id=session_id,
-            schema_version=state.schema_version,
-            last_event_applied_idx=state.last_event_applied_idx,
-            chunk_index=i,
-            total_chunks=len(pieces),
-            sha256_hex=sha256,
-        )
-        for i, piece in enumerate(pieces)
-    ]
-
-
-def test_completes_on_full_transfer(my_node: NodeId, session_id: SessionId):
-    state = State(last_event_applied_idx=42)
-    chunks = _make_chunks(
-        _encode(state),
-        chunk_size=64,
-        requester_node_id=my_node,
-        session_id=session_id,
-        state=state,
-    )
-    receiver = SnapshotReceiver(my_node, session_id)
-
-    received = None
-    for chunk in chunks:
-        received = receiver.ingest(chunk)
-    assert received is not None
-    assert received.last_event_applied_idx == 42
-    assert received.state.last_event_applied_idx == 42
-
-
-def test_handles_out_of_order_chunks(my_node: NodeId, session_id: SessionId):
-    state = State(last_event_applied_idx=99)
-    chunks = _make_chunks(
-        _encode(state),
-        chunk_size=32,
-        requester_node_id=my_node,
-        session_id=session_id,
-        state=state,
-    )
-    receiver = SnapshotReceiver(my_node, session_id)
-
-    # Reverse them.
-    received = None
-    for chunk in reversed(chunks):
-        received = receiver.ingest(chunk)
-    assert received is not None
-    assert received.last_event_applied_idx == 99
-
-
-def test_ignores_chunks_for_other_recipients(my_node: NodeId, session_id: SessionId):
-    state = State(last_event_applied_idx=1)
-    other = NodeId("worker-2")
-    chunks = _make_chunks(
-        _encode(state),
-        chunk_size=64,
-        requester_node_id=other,
-        session_id=session_id,
-        state=state,
-    )
-    receiver = SnapshotReceiver(my_node, session_id)
-    for chunk in chunks:
-        assert receiver.ingest(chunk) is None
-
-
-def test_ignores_chunks_from_stale_session(my_node: NodeId, session_id: SessionId):
-    state = State(last_event_applied_idx=1)
-    other_session = SessionId(master_node_id=NodeId("other-master"), election_clock=99)
-    chunks = _make_chunks(
-        _encode(state),
-        chunk_size=64,
-        requester_node_id=my_node,
-        session_id=other_session,
-        state=state,
-    )
-    receiver = SnapshotReceiver(my_node, session_id)
-    for chunk in chunks:
-        assert receiver.ingest(chunk) is None
-
-
-def test_discards_on_checksum_mismatch(my_node: NodeId, session_id: SessionId):
-    state = State(last_event_applied_idx=1)
-    chunks = _make_chunks(
-        _encode(state),
-        chunk_size=64,
-        requester_node_id=my_node,
-        session_id=session_id,
-        state=state,
-    )
-    # Corrupt the last byte of the last chunk.
-    original = chunks[-1]
-    chunks[-1] = SnapshotChunk.from_data(
-        data=original.data + b"\x00garbage",
-        transfer_id=original.transfer_id,
-        requester_node_id=original.requester_node_id,
-        session_id=original.session_id,
-        schema_version=original.schema_version,
-        last_event_applied_idx=original.last_event_applied_idx,
-        chunk_index=original.chunk_index,
-        total_chunks=original.total_chunks,
-        sha256_hex=original.sha256_hex,
-    )
-
-    receiver = SnapshotReceiver(my_node, session_id)
-    received = None
-    for chunk in chunks:
-        received = receiver.ingest(chunk)
-    assert received is None

src/exo/routing/tests/test_snapshot_routing_types.py

@@ -1,37 +0,0 @@
-from exo.routing import topics
-from exo.shared.types.commands import ForwarderCommand, RequestSnapshot
-from exo.shared.types.common import NodeId, SessionId, SystemId
-from exo.shared.types.snapshots import SnapshotChunk, SnapshotTransferId
-
-
-def test_request_snapshot_round_trips_through_forwarder_command() -> None:
-    command = ForwarderCommand(
-        origin=SystemId("system-1"),
-        command=RequestSnapshot(requester_node_id=NodeId("worker-1")),
-    )
-
-    restored = ForwarderCommand.model_validate_json(command.model_dump_json())
-
-    assert isinstance(restored.command, RequestSnapshot)
-    assert restored.command.requester_node_id == NodeId("worker-1")
-
-
-def test_snapshot_response_topic_round_trips_chunk() -> None:
-    chunk = SnapshotChunk.from_data(
-        data=b"snapshot-bytes",
-        transfer_id=SnapshotTransferId("transfer-1"),
-        requester_node_id=NodeId("worker-1"),
-        session_id=SessionId(master_node_id=NodeId("master"), election_clock=1),
-        schema_version=1,
-        last_event_applied_idx=42,
-        chunk_index=0,
-        total_chunks=1,
-        sha256_hex="unused",
-    )
-
-    restored = topics.SNAPSHOT_RESPONSES.deserialize(
-        topics.SNAPSHOT_RESPONSES.serialize(chunk)
-    )
-
-    assert restored == chunk
-    assert restored.data == b"snapshot-bytes"

src/exo/routing/topics.py

@@ -8,7 +8,6 @@ from exo.shared.types.events import (
     GlobalForwarderEvent,
     LocalForwarderEvent,
 )
-from exo.shared.types.snapshots import SnapshotChunk
 from exo.utils.pydantic_ext import FrozenModel
 
 
@@ -50,6 +49,3 @@ CONNECTION_MESSAGES = TypedTopic(
 DOWNLOAD_COMMANDS = TypedTopic(
     "download_commands", PublishPolicy.Always, ForwarderDownloadCommand
 )
-SNAPSHOT_RESPONSES = TypedTopic(
-    "snapshot_responses", PublishPolicy.Always, SnapshotChunk
-)

src/exo/shared/apply.py

@@ -4,7 +4,8 @@ from datetime import datetime
 
 from loguru import logger
 
-from exo.shared.types.common import NodeId
+from exo.shared.models.model_cards import ModelCard
+from exo.shared.types.common import ModelId, NodeId
 from exo.shared.types.events import (
     ChunkGenerated,
     CustomModelCardAdded,
@@ -40,14 +41,7 @@ from exo.shared.types.profiling import (
     ThunderboltBridgeStatus,
 )
 from exo.shared.types.state import State
-from exo.shared.types.tasks import (
-    ImageEdits,
-    ImageGeneration,
-    Task,
-    TaskId,
-    TaskStatus,
-    TextGeneration,
-)
+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
@@ -72,6 +66,18 @@ 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:
@@ -79,12 +85,15 @@ def event_apply(event: Event, state: State) -> State:
             TestEvent()
             | ChunkGenerated()
             | TaskAcknowledged()
+            | InputChunkReceived()
             | TracesCollected()
             | TracesMerged()
-            | CustomModelCardAdded()
-            | CustomModelCardDeleted()
         ):  # 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():
@@ -99,8 +108,6 @@ def event_apply(event: Event, state: State) -> State:
             return apply_runner_status_updated(event, state)
         case TaskCreated():
             return apply_task_created(event, state)
-        case InputChunkReceived():
-            return apply_input_chunk_received(event, state)
         case TaskDeleted():
             return apply_task_deleted(event, state)
         case TaskFailed():
@@ -165,32 +172,10 @@ def apply_task_created(event: TaskCreated, state: State) -> State:
     return state.model_copy(update={"tasks": new_tasks})
 
 
-def apply_input_chunk_received(event: InputChunkReceived, state: State) -> State:
-    command_chunks = {
-        **state.input_chunks.get(event.command_id, {}),
-        event.chunk.chunk_index: event.chunk,
-    }
-    return state.model_copy(
-        update={
-            "input_chunks": {**state.input_chunks, event.command_id: command_chunks}
-        }
-    )
-
-
 def apply_task_deleted(event: TaskDeleted, state: State) -> State:
-    task = state.tasks.get(event.task_id)
     new_tasks: Mapping[TaskId, Task] = {
         tid: task for tid, task in state.tasks.items() if tid != event.task_id
     }
-    if isinstance(task, (TextGeneration, ImageGeneration, ImageEdits)):
-        new_input_chunks = {
-            command_id: chunks
-            for command_id, chunks in state.input_chunks.items()
-            if command_id != task.command_id
-        }
-        return state.model_copy(
-            update={"tasks": new_tasks, "input_chunks": new_input_chunks}
-        )
     return state.model_copy(update={"tasks": new_tasks})
 
 
@@ -427,6 +412,9 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
                 for nid in state.node_thunderbolt
                 for tb_ident in state.node_thunderbolt[nid].interfaces
             }
+            source_is_rdma_enabled = _is_rdma_ctl_enabled(
+                event.node_id, state.node_rdma_ctl
+            )
             as_rdma_conns = [
                 Connection(
                     source=event.node_id,
@@ -439,6 +427,10 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
                 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():
@@ -462,6 +454,12 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
                 **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)
 
     return state.model_copy(update=update)
 
@@ -477,3 +475,22 @@ def apply_topology_edge_deleted(event: TopologyEdgeDeleted, state: State) -> Sta
     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,
+    }
+    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
+    }
+    return state.model_copy(update={"custom_model_cards": new_cards})

src/exo/shared/constants.py

@@ -69,6 +69,7 @@ DASHBOARD_DIR = (
 EXO_LOG_DIR = EXO_CACHE_HOME / "exo_log"
 EXO_LOG = EXO_LOG_DIR / "exo.log"
 EXO_TEST_LOG = EXO_CACHE_HOME / "exo_test.log"
+EXO_PID_FILE = EXO_CACHE_HOME / "exo.pid"
 
 # Identity (config)
 EXO_NODE_ID_KEYPAIR = EXO_CONFIG_HOME / "node_id.keypair"

src/exo/shared/election.py

@@ -9,7 +9,6 @@ from anyio import (
 from loguru import logger
 
 from exo.routing.connection_message import ConnectionMessage
-from exo.shared.types.commands import ForwarderCommand
 from exo.shared.types.common import NodeId, SessionId
 from exo.utils.channels import Receiver, Sender
 from exo.utils.pydantic_ext import FrozenModel
@@ -22,7 +21,6 @@ class ElectionMessage(FrozenModel):
     clock: int
     seniority: int
     proposed_session: SessionId
-    commands_seen: int
 
     # Could eventually include a list of neighbour nodes for centrality
     def __lt__(self, other: Self) -> bool:
@@ -30,8 +28,6 @@ class ElectionMessage(FrozenModel):
             return self.clock < other.clock
         if self.seniority != other.seniority:
             return self.seniority < other.seniority
-        elif self.commands_seen != other.commands_seen:
-            return self.commands_seen < other.commands_seen
         else:
             return (
                 self.proposed_session.master_node_id
@@ -54,7 +50,6 @@ class Election:
         election_message_sender: Sender[ElectionMessage],
         election_result_sender: Sender[ElectionResult],
         connection_message_receiver: Receiver[ConnectionMessage],
-        command_receiver: Receiver[ForwarderCommand],
         is_candidate: bool = True,
         seniority: int = 0,
     ):
@@ -64,7 +59,6 @@ class Election:
         self.seniority = seniority if is_candidate else -1
         self.clock = 0
         self.node_id = node_id
-        self.commands_seen = 0
         # Every node spawns as master
         self.current_session: SessionId = SessionId(
             master_node_id=node_id, election_clock=0
@@ -75,7 +69,6 @@ class Election:
         self._em_receiver = election_message_receiver
         self._er_sender = election_result_sender
         self._cm_receiver = connection_message_receiver
-        self._co_receiver = command_receiver
 
         # Campaign state
         self._candidates: list[ElectionMessage] = []
@@ -89,7 +82,6 @@ class Election:
             async with self._tg as tg:
                 tg.start_soon(self._election_receiver)
                 tg.start_soon(self._connection_receiver)
-                tg.start_soon(self._command_counter)
 
                 # And start an election immediately, that instantly resolves
                 candidates: list[ElectionMessage] = []
@@ -179,11 +171,6 @@ class Election:
                 logger.debug("Campaign started")
                 logger.debug("Connection message added")
 
-    async def _command_counter(self) -> None:
-        with self._co_receiver as commands:
-            async for _command in commands:
-                self.commands_seen += 1
-
     async def _campaign(
         self, candidates: list[ElectionMessage], campaign_timeout: float
     ) -> None:
@@ -261,5 +248,4 @@ class Election:
             ),
             clock=c,
             seniority=self.seniority,
-            commands_seen=self.commands_seen,
         )

src/exo/shared/logging.py

@@ -46,7 +46,8 @@ class _InterceptHandler(logging.Handler):
 def logger_setup(log_file: Path | None, verbosity: int = 0):
     """Set up logging for this process - formatting, file handles, verbosity and output"""
 
-    logging.getLogger("exo_pyo3_bindings").setLevel(logging.WARNING)
+    logging.getLogger("exo_net").setLevel(logging.INFO)
+    logging.getLogger("networking").setLevel(logging.INFO)
     logging.getLogger("httpx").setLevel(logging.WARNING)
     logging.getLogger("httpcore").setLevel(logging.WARNING)
 

src/exo/shared/models/model_cards.py

@@ -39,7 +39,57 @@ _BUILTIN_CARD_DIRS = [
     Path(RESOURCES_DIR) / "image_model_cards",
 ]
 
-_card_cache: dict[ModelId, "ModelCard"] = {}
+
+class _CardCache:
+    def __init__(self):
+        self.cc: dict[ModelId, "ModelCard"] = {}
+
+    def get(self, model_id: ModelId) -> "ModelCard | None":
+        return self.cc.get(model_id)
+
+    async def save(self, card: "ModelCard"):
+        self.cc[card.model_id] = card
+        try:
+            await card.save_to_custom_dir()
+        except OSError as e:
+            logger.warning(f"failed to save custom model card ({e.strerror})")
+
+    async def pop(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:
+            if await card_path.exists():
+                await card_path.unlink()
+                return self.cc.pop(model_id, None)
+        except OSError as e:
+            logger.warning(f"failed to delete custom model card ({e.strerror})")
+
+    async def list_all(self) -> list["ModelCard"]:
+        if len(self.cc) == 0:
+            await self.refresh()
+        if EXO_ENABLE_IMAGE_MODELS:
+            return list(self.cc.values())
+        return [c for c in self.cc.values() if not _is_image_card(c)]
+
+    async def _load_cards_from_dir(self, directory: Path, *, is_custom: bool) -> None:
+        """Load all TOML model cards from a directory into the cache."""
+        async for toml_file in directory.rglob("*.toml"):
+            try:
+                card = await ModelCard.load_from_path(toml_file)
+                if is_custom:
+                    card = card.model_copy(update={"is_custom": True})
+                if self.get(card.model_id) is None:
+                    self.cc[card.model_id] = card
+            except (ValidationError, TOMLKitError):
+                pass
+
+    async def refresh(self) -> None:
+        for path in _BUILTIN_CARD_DIRS:
+            await self._load_cards_from_dir(path, is_custom=False)
+        await self._load_cards_from_dir(_custom_cards_dir, is_custom=True)
+
+
+card_cache = _CardCache()
 
 
 def detect_vision_from_config(model_id: ModelId) -> "VisionCardConfig | None":
@@ -59,42 +109,10 @@ def detect_vision_from_config(model_id: ModelId) -> "VisionCardConfig | None":
     return None
 
 
-async def _load_cards_from_dir(directory: Path, *, is_custom: bool) -> None:
-    """Load all TOML model cards from a directory into the cache."""
-    async for toml_file in directory.rglob("*.toml"):
-        try:
-            card = await ModelCard.load_from_path(toml_file)
-            if is_custom:
-                card = card.model_copy(update={"is_custom": True})
-            if card.model_id not in _card_cache:
-                _card_cache[card.model_id] = card
-        except (ValidationError, TOMLKitError):
-            pass
-
-
-async def _refresh_card_cache() -> None:
-    for path in _BUILTIN_CARD_DIRS:
-        await _load_cards_from_dir(path, is_custom=False)
-    await _load_cards_from_dir(_custom_cards_dir, is_custom=True)
-
-
 def _is_image_card(card: "ModelCard") -> bool:
     return any(t in (ModelTask.TextToImage, ModelTask.ImageToImage) for t in card.tasks)
 
 
-def get_card(model_id: ModelId) -> "ModelCard | None":
-    """Look up a single model card from the cache by ID."""
-    return _card_cache.get(model_id)
-
-
-async def get_model_cards() -> list["ModelCard"]:
-    if len(_card_cache) == 0:
-        await _refresh_card_cache()
-    if EXO_ENABLE_IMAGE_MODELS:
-        return list(_card_cache.values())
-    return [c for c in _card_cache.values() if not _is_image_card(c)]
-
-
 class ModelTask(str, Enum):
     TextGeneration = "TextGeneration"
     TextToImage = "TextToImage"
@@ -196,14 +214,13 @@ class ModelCard(FrozenModel):
     # Is it okay that model card.load defaults to network access if the card doesn't exist? do we want to be more explicit here?
     @staticmethod
     async def load(model_id: ModelId) -> "ModelCard":
-        if model_id not in _card_cache:
-            await _refresh_card_cache()
-        if (mc := _card_cache.get(model_id)) is not None:
+        if card_cache.get(model_id) is None:
+            await card_cache.refresh()
+        if (mc := card_cache.get(model_id)) is not None:
             return mc
 
         mc = await ModelCard.fetch_from_hf(model_id)
         await mc.save_to_custom_dir()
-        _card_cache[model_id] = mc
         return mc
 
     @staticmethod
@@ -233,21 +250,6 @@ class ModelCard(FrozenModel):
         )
 
 
-def add_to_card_cache(card: "ModelCard") -> None:
-    """Add or update a model card in the in-memory cache."""
-    _card_cache[card.model_id] = card
-
-
-async def delete_custom_card(model_id: ModelId) -> bool:
-    """Delete a user-added custom model card. Returns True if deleted."""
-    card_path = _custom_cards_dir / (ModelId(model_id).normalize() + ".toml")
-    if await card_path.exists():
-        await card_path.unlink()
-        _card_cache.pop(model_id, None)
-        return True
-    return False
-
-
 class ConfigData(BaseModel):
     model_config = {"extra": "ignore"}  # Allow unknown fields
 

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

@@ -0,0 +1,44 @@
+from exo.shared.apply import apply
+from exo.shared.models.model_cards import ModelCard, ModelTask
+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],
+    )
+
+
+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_input_chunks.py

@@ -1,85 +0,0 @@
-from exo.shared.apply import apply
-from exo.shared.models.model_cards import ModelId
-from exo.shared.types.chunks import InputImageChunk
-from exo.shared.types.common import CommandId
-from exo.shared.types.events import (
-    IndexedEvent,
-    InputChunkReceived,
-    TaskCreated,
-    TaskDeleted,
-)
-from exo.shared.types.state import State
-from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
-from exo.shared.types.text_generation import (
-    InputMessage,
-    InputMessageContent,
-    TextGenerationTaskParams,
-)
-from exo.shared.types.worker.instances import InstanceId
-
-
-def test_apply_input_chunk_received_stores_chunk_in_state() -> None:
-    command_id = CommandId("command-1")
-    chunk = InputImageChunk(
-        model=ModelId("mlx-community/test-model"),
-        command_id=command_id,
-        data="abc",
-        chunk_index=0,
-        total_chunks=1,
-        image_index=0,
-    )
-
-    state = apply(
-        State(),
-        IndexedEvent(
-            idx=0,
-            event=InputChunkReceived(command_id=command_id, chunk=chunk),
-        ),
-    )
-
-    assert state.input_chunks == {command_id: {0: chunk}}
-
-
-def test_apply_task_deleted_removes_chunks_for_generation_command() -> None:
-    command_id = CommandId("command-1")
-    task_id = TaskId("task-1")
-    chunk = InputImageChunk(
-        model=ModelId("mlx-community/test-model"),
-        command_id=command_id,
-        data="abc",
-        chunk_index=0,
-        total_chunks=1,
-        image_index=0,
-    )
-    task = TextGeneration(
-        task_id=task_id,
-        instance_id=InstanceId("instance-1"),
-        task_status=TaskStatus.Pending,
-        command_id=command_id,
-        task_params=TextGenerationTaskParams(
-            model=ModelId("mlx-community/test-model"),
-            input=[
-                InputMessage(role="user", content=InputMessageContent("hello")),
-            ],
-        ),
-    )
-
-    state = State()
-    state = apply(
-        state,
-        IndexedEvent(
-            idx=0,
-            event=InputChunkReceived(command_id=command_id, chunk=chunk),
-        ),
-    )
-    state = apply(
-        state,
-        IndexedEvent(idx=1, event=TaskCreated(task_id=task_id, task=task)),
-    )
-    state = apply(
-        state,
-        IndexedEvent(idx=2, event=TaskDeleted(task_id=task_id)),
-    )
-
-    assert state.tasks == {}
-    assert state.input_chunks == {}

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

@@ -0,0 +1,231 @@
+from datetime import datetime, timezone
+
+from exo.shared.apply import apply_node_gathered_info
+from exo.shared.topology import Topology
+from exo.shared.types.common import NodeId
+from exo.shared.types.events import NodeGatheredInfo
+from exo.shared.types.profiling import (
+    NodeRdmaCtlStatus,
+    NodeThunderboltInfo,
+)
+from exo.shared.types.state import State
+from exo.shared.types.thunderbolt import ThunderboltConnection, ThunderboltIdentifier
+from exo.shared.types.topology import RDMAConnection
+from exo.utils.info_gatherer.info_gatherer import (
+    MacThunderboltConnections,
+    RdmaCtlStatus,
+)
+
+
+def _now() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+def _make_state_with_thunderbolt_idents(
+    *node_ids_and_uuids: tuple[NodeId, str, str],
+    rdma_ctl: dict[NodeId, NodeRdmaCtlStatus] | None = None,
+) -> State:
+    """Build a State with Thunderbolt identifiers per node so the apply MacThunderboltConnections
+    case can resolve uuid -> (node, iface)."""
+    node_thunderbolt = {
+        nid: NodeThunderboltInfo(
+            interfaces=[ThunderboltIdentifier(rdma_interface=iface, domain_uuid=uuid)]
+        )
+        for nid, uuid, iface in node_ids_and_uuids
+    }
+    return State(
+        node_thunderbolt=node_thunderbolt,
+        node_rdma_ctl=rdma_ctl or {},
+    )
+
+
+def _has_rdma_edge(topology: Topology, source: NodeId, sink: NodeId) -> bool:
+    return any(
+        isinstance(edge, RDMAConnection)
+        for edge in topology.get_all_connections_between(source, sink)
+    )
+
+
+def test_mac_thunderbolt_connections_emits_rdma_when_both_endpoints_enabled():
+    node_a = NodeId()
+    node_b = NodeId()
+    state = _make_state_with_thunderbolt_idents(
+        (node_a, "uuid-a", "rdma_en1"),
+        (node_b, "uuid-b", "rdma_en1"),
+        rdma_ctl={
+            node_a: NodeRdmaCtlStatus(enabled=True),
+            node_b: NodeRdmaCtlStatus(enabled=True),
+        },
+    )
+
+    event = NodeGatheredInfo(
+        node_id=node_a,
+        when=_now(),
+        info=MacThunderboltConnections(
+            conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
+        ),
+    )
+
+    new_state = apply_node_gathered_info(event, state)
+
+    assert _has_rdma_edge(new_state.topology, node_a, node_b)
+
+
+def test_mac_thunderbolt_connections_skips_rdma_when_source_rdma_ctl_disabled():
+    node_a = NodeId()
+    node_b = NodeId()
+    state = _make_state_with_thunderbolt_idents(
+        (node_a, "uuid-a", "rdma_en1"),
+        (node_b, "uuid-b", "rdma_en1"),
+        rdma_ctl={
+            node_a: NodeRdmaCtlStatus(enabled=False),
+            node_b: NodeRdmaCtlStatus(enabled=True),
+        },
+    )
+
+    event = NodeGatheredInfo(
+        node_id=node_a,
+        when=_now(),
+        info=MacThunderboltConnections(
+            conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
+        ),
+    )
+
+    new_state = apply_node_gathered_info(event, state)
+
+    assert not _has_rdma_edge(new_state.topology, node_a, node_b)
+
+
+def test_mac_thunderbolt_connections_skips_rdma_when_sink_rdma_ctl_disabled():
+    node_a = NodeId()
+    node_b = NodeId()
+    state = _make_state_with_thunderbolt_idents(
+        (node_a, "uuid-a", "rdma_en1"),
+        (node_b, "uuid-b", "rdma_en1"),
+        rdma_ctl={
+            node_a: NodeRdmaCtlStatus(enabled=True),
+            node_b: NodeRdmaCtlStatus(enabled=False),
+        },
+    )
+
+    event = NodeGatheredInfo(
+        node_id=node_a,
+        when=_now(),
+        info=MacThunderboltConnections(
+            conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
+        ),
+    )
+
+    new_state = apply_node_gathered_info(event, state)
+
+    assert not _has_rdma_edge(new_state.topology, node_a, node_b)
+
+
+def test_mac_thunderbolt_connections_skips_rdma_when_rdma_ctl_status_missing():
+    """Missing rdma_ctl status defaults to not-enabled — node is RDMA-incapable."""
+    node_a = NodeId()
+    node_b = NodeId()
+    state = _make_state_with_thunderbolt_idents(
+        (node_a, "uuid-a", "rdma_en1"),
+        (node_b, "uuid-b", "rdma_en1"),
+        rdma_ctl={
+            node_a: NodeRdmaCtlStatus(enabled=True),
+            # node_b intentionally absent
+        },
+    )
+
+    event = NodeGatheredInfo(
+        node_id=node_a,
+        when=_now(),
+        info=MacThunderboltConnections(
+            conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
+        ),
+    )
+
+    new_state = apply_node_gathered_info(event, state)
+
+    assert not _has_rdma_edge(new_state.topology, node_a, node_b)
+
+
+def test_rdma_ctl_status_disabled_purges_existing_rdma_edges():
+    """When a node reports rdma_ctl disabled, all RDMA edges touching it must be removed."""
+    node_a = NodeId()
+    node_b = NodeId()
+
+    # Start with both nodes RDMA-enabled and existing RDMA edges in the topology.
+    state = _make_state_with_thunderbolt_idents(
+        (node_a, "uuid-a", "rdma_en1"),
+        (node_b, "uuid-b", "rdma_en1"),
+        rdma_ctl={
+            node_a: NodeRdmaCtlStatus(enabled=True),
+            node_b: NodeRdmaCtlStatus(enabled=True),
+        },
+    )
+    state = apply_node_gathered_info(
+        NodeGatheredInfo(
+            node_id=node_a,
+            when=_now(),
+            info=MacThunderboltConnections(
+                conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
+            ),
+        ),
+        state,
+    )
+    state = apply_node_gathered_info(
+        NodeGatheredInfo(
+            node_id=node_b,
+            when=_now(),
+            info=MacThunderboltConnections(
+                conns=[ThunderboltConnection(source_uuid="uuid-b", sink_uuid="uuid-a")]
+            ),
+        ),
+        state,
+    )
+    assert _has_rdma_edge(state.topology, node_a, node_b)
+    assert _has_rdma_edge(state.topology, node_b, node_a)
+
+    # Now node_a flips to rdma_ctl disabled — both directions of RDMA edge must drop.
+    state = apply_node_gathered_info(
+        NodeGatheredInfo(
+            node_id=node_a, when=_now(), info=RdmaCtlStatus(enabled=False)
+        ),
+        state,
+    )
+
+    assert not _has_rdma_edge(state.topology, node_a, node_b)
+    assert not _has_rdma_edge(state.topology, node_b, node_a)
+    assert state.node_rdma_ctl[node_a].enabled is False
+
+
+def test_topology_remove_all_rdma_connections_touching_keeps_socket_edges():
+    """Purging RDMA edges for a disabled node must not affect non-RDMA edges."""
+    from exo.shared.types.multiaddr import Multiaddr
+    from exo.shared.types.topology import Connection, SocketConnection
+
+    topology = Topology()
+    node_a = NodeId()
+    node_b = NodeId()
+    topology.add_node(node_a)
+    topology.add_node(node_b)
+    topology.add_connection(
+        Connection(
+            source=node_a,
+            sink=node_b,
+            edge=RDMAConnection(
+                source_rdma_iface="rdma_en1", sink_rdma_iface="rdma_en1"
+            ),
+        )
+    )
+    socket_edge = SocketConnection(
+        sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000")
+    )
+    topology.add_connection(Connection(source=node_a, sink=node_b, edge=socket_edge))
+
+    topology.remove_all_rdma_connections_touching(node_a)
+
+    assert not _has_rdma_edge(topology, node_a, node_b)
+    # Socket edge survives.
+    assert any(
+        isinstance(edge, SocketConnection)
+        for edge in topology.get_all_connections_between(node_a, node_b)
+    )

src/exo/shared/tests/test_election.py

@@ -327,7 +327,7 @@ async def test_connection_message_triggers_new_round_broadcast() -> None:
             tg.start_soon(election.run)
 
             # Send any connection message object; we close quickly to cancel before result creation
-            await cm_tx.send(ConnectionMessage(node_id=NodeId(), connected=True))
+            await cm_tx.send(ConnectionMessage(connected=True))
 
             # Expect a broadcast for the new round at clock=1
             while True:

src/exo/shared/tests/test_state_serialization.py

@@ -25,7 +25,6 @@ def test_state_serialization_roundtrip() -> None:
     json_repr = state.model_dump_json()
     restored_state = State.model_validate_json(json_repr)
 
-    assert restored_state.schema_version == state.schema_version
     assert (
         state.topology.to_snapshot().nodes
         == restored_state.topology.to_snapshot().nodes

src/exo/shared/topology.py

@@ -169,6 +169,22 @@ class Topology:
         for conn in new_connections:
             self.add_connection(conn)
 
+    def remove_all_rdma_connections_touching(self, node_id: NodeId) -> None:
+        """Remove every RDMA edge incident to ``node_id`` (incoming or outgoing)."""
+        if node_id not in self._vertex_indices:
+            return
+        rx_idx = self._vertex_indices[node_id]
+        rdma_edge_idxs = [
+            edge_idx
+            for edge_idx in (
+                *self._graph.out_edge_indices(rx_idx),
+                *self._graph.in_edge_indices(rx_idx),
+            )
+            if isinstance(self._graph.get_edge_data_by_index(edge_idx), RDMAConnection)
+        ]
+        for edge_idx in rdma_edge_idxs:
+            self._graph.remove_edge_from_index(edge_idx)
+
     def remove_connection(self, conn: Connection) -> None:
         if (
             conn.source not in self._vertex_indices

src/exo/shared/types/chunks.py

@@ -32,15 +32,21 @@ class TokenChunk(BaseChunk):
 
 class ErrorChunk(BaseChunk):
     error_message: str
-    finish_reason: Literal["error"] = "error"
+
+    @property
+    def finish_reason(self) -> Literal["error"]:
+        return "error"
 
 
 class ToolCallChunk(BaseChunk):
     tool_calls: list[ToolCallItem]
     usage: Usage | None
-    finish_reason: Literal["tool_calls"] = "tool_calls"
     stats: GenerationStats | None = None
 
+    @property
+    def finish_reason(self) -> Literal["tool_calls"]:
+        return "tool_calls"
+
 
 class ImageChunk(BaseChunk):
     data: str
@@ -84,7 +90,13 @@ class PrefillProgressChunk(BaseChunk):
     processed_tokens: int
     total_tokens: int
 
+    @property
+    def finish_reason(self) -> FinishReason | None:
+        return None
+
 
 StatusChunk = PrefillProgressChunk
 GenerationChunk = TokenChunk | ImageChunk | ToolCallChunk | ErrorChunk
+TextGenerationChunk = TokenChunk | ToolCallChunk | ErrorChunk
+ImageGenerationChunk = ImageChunk | ErrorChunk
 Chunk = StatusChunk | GenerationChunk

src/exo/shared/types/commands.py

@@ -67,12 +67,6 @@ class RequestEventLog(BaseCommand):
     since_idx: int
 
 
-class RequestSnapshot(BaseCommand):
-    """Ask the current master to send a State snapshot to this node."""
-
-    requester_node_id: NodeId
-
-
 class StartDownload(BaseCommand):
     target_node_id: NodeId
     shard_metadata: ShardMetadata
@@ -112,7 +106,6 @@ DownloadCommand = StartDownload | DeleteDownload | CancelDownload
 Command = (
     TestCommand
     | RequestEventLog
-    | RequestSnapshot
     | TextGeneration
     | ImageGeneration
     | ImageEdits

src/exo/shared/types/events.py

@@ -159,10 +159,10 @@ Event = (
     | NodeTimedOut
     | NodeGatheredInfo
     | NodeDownloadProgress
-    | ChunkGenerated
-    | InputChunkReceived
     | TopologyEdgeCreated
     | TopologyEdgeDeleted
+    | ChunkGenerated
+    | InputChunkReceived
     | TracesCollected
     | TracesMerged
     | CustomModelCardAdded

src/exo/shared/types/snapshots.py

@@ -1,54 +0,0 @@
-"""Wire types for snapshot transfer between master and a joining node.
-
-Snapshots can be tens of MB; the gossipsub message ceiling is around 1 MB.
-We slice the compressed snapshot body into chunks and publish each chunk on
-the SNAPSHOT_RESPONSES topic. The receiver collects chunks for its own
-`requester_node_id`, validates the SHA-256 of the reassembled body, and
-materialises the State.
-"""
-
-import base64
-
-from exo.shared.types.common import Id, NodeId, SessionId
-from exo.utils.pydantic_ext import FrozenModel
-
-
-class SnapshotTransferId(Id):
-    """Identifies a single snapshot transfer (one master response to one
-    `RequestSnapshot`). Distinct transfers may interleave; the id lets
-    receivers keep them apart."""
-
-
-class SnapshotChunk(FrozenModel):
-    """One slice of a snapshot in flight.
-
-    `data_b64` carries a base64-encoded slice of the zstd-compressed JSON
-    dump of State. Concatenating the *decoded* bytes of all chunks for a
-    `transfer_id` in order of `chunk_index` yields the full compressed
-    body; `sha256_hex` is the SHA-256 of that decoded blob.
-
-    We use base64 explicitly because the topic layer JSON-encodes messages,
-    and JSON can't carry raw binary. Helpers `from_data` / `data` keep the
-    base64 detail at the boundaries.
-    """
-
-    transfer_id: SnapshotTransferId
-    requester_node_id: NodeId
-    session_id: SessionId
-    schema_version: int
-    last_event_applied_idx: int
-    chunk_index: int
-    total_chunks: int
-    sha256_hex: str
-    data_b64: str
-
-    @classmethod
-    def from_data(cls, *, data: bytes, **kwargs: object) -> "SnapshotChunk":
-        return cls(data_b64=base64.b64encode(data).decode("ascii"), **kwargs)  # pyright: ignore[reportArgumentType]
-
-    @property
-    def data(self) -> bytes:
-        return base64.b64decode(self.data_b64)
-
-
-__all__ = ["SnapshotChunk", "SnapshotTransferId"]

src/exo/shared/types/state.py

@@ -5,9 +5,9 @@ from typing import Any, cast
 from pydantic import ConfigDict, Field, field_serializer, field_validator
 from pydantic.alias_generators import to_camel
 
+from exo.shared.models.model_cards import ModelCard
 from exo.shared.topology import Topology, TopologySnapshot
-from exo.shared.types.chunks import InputImageChunk
-from exo.shared.types.common import CommandId, NodeId
+from exo.shared.types.common import ModelId, NodeId
 from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.profiling import (
     DiskUsage,
@@ -42,16 +42,10 @@ class State(FrozenModel):
         strict=True,
         arbitrary_types_allowed=True,
     )
-    # Bump when a State change makes older snapshots unsafe to restore.
-    schema_version: int = Field(default=1, ge=1)
-
     instances: Mapping[InstanceId, Instance] = {}
     runners: Mapping[RunnerId, RunnerStatus] = {}
     downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {}
     tasks: Mapping[TaskId, Task] = {}
-    # Durable request input chunks for active image requests. Workers rebuild
-    # local image caches from this state instead of reading events directly.
-    input_chunks: Mapping[CommandId, Mapping[int, InputImageChunk]] = {}
     last_seen: Mapping[NodeId, datetime] = {}
     topology: Topology = Field(default_factory=Topology)
     last_event_applied_idx: int = Field(default=-1, ge=-1)
@@ -72,6 +66,9 @@ class State(FrozenModel):
     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] = {}
+
     @field_serializer("topology", mode="plain")
     def _encode_topology(self, value: Topology) -> TopologySnapshot:
         return value.to_snapshot()

src/exo/shared/types/text_generation.py

@@ -135,9 +135,9 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
     prefill_endpoint: str | None = None
 
     def with_card_sampling_defaults(self) -> "TextGenerationTaskParams":
-        from exo.shared.models.model_cards import get_card
+        from exo.shared.models import model_cards
 
-        card = get_card(self.model)
+        card = model_cards.card_cache.get(self.model)
         if card is None:
             return self
 

src/exo/utils/banner.py

@@ -38,7 +38,7 @@ def print_startup_banner(port: int) -> None:
 
 ╔═══════════════════════════════════════════════════════════════════════╗
 ║                                                                       ║
-║  🌐 Dashboard & API Ready                                             ║
+║  Dashboard & API Ready                                                ║
 ║                                                                       ║
 ║  {dashboard_url}{" " * (69 - len(dashboard_url))}║
 ║                                                                       ║

src/exo/utils/event_buffer.py

@@ -47,18 +47,6 @@ class OrderedBuffer[T]:
         logger.trace(f"Releasing event {ret}")
         return ret
 
-    def fast_forward_to(self, idx: int) -> None:
-        """Skip every event before idx.
-
-        Snapshot restore uses this after applying state that already includes
-        events before idx. Any buffered or future event below idx is stale.
-        """
-        if idx <= self.next_idx_to_release:
-            return
-        self.next_idx_to_release = idx
-        for stale_idx in [i for i in self.store if i < idx]:
-            del self.store[stale_idx]
-
 
 class MultiSourceBuffer[SourceId, T]:
     """

src/exo/utils/power_sampler.py

@@ -19,19 +19,21 @@ class PowerSampler:
     ):
         self._get_node_system = get_node_system
         self._interval = interval
-        self._samples: defaultdict[NodeId, list[SystemPerformanceProfile]] = (
-            defaultdict(list)
-        )
+        self._samples: defaultdict[
+            NodeId, list[tuple[float, SystemPerformanceProfile]]
+        ] = defaultdict(list)
         self._start_time: float | None = None
         self._stopped = False
 
-    def _take_sample(self) -> None:
+    def _take_sample(self, t_rel: float | None = None) -> None:
+        assert self._start_time is not None
+        ts = t_rel if t_rel is not None else time.perf_counter() - self._start_time
         for node_id, profile in self._get_node_system().items():
-            self._samples[node_id].append(profile)
+            self._samples[node_id].append((ts, profile))
 
     async def run(self) -> None:
         self._start_time = time.perf_counter()
-        self._take_sample()
+        self._take_sample(t_rel=0.0)
         while not self._stopped:
             await anyio.sleep(self._interval)
             self._take_sample()
@@ -39,26 +41,51 @@ class PowerSampler:
     def result(self) -> PowerUsage:
         self._stopped = True
         assert self._start_time is not None, "result() called before run()"
-        self._take_sample()
         elapsed = time.perf_counter() - self._start_time
+        self._take_sample(t_rel=elapsed)
 
         node_stats: list[NodePowerStats] = []
-        for node_id, profiles in self._samples.items():
-            n = len(profiles)
+        total_energy_j = 0.0
+        for node_id, ts_profiles in self._samples.items():
+            n = len(ts_profiles)
             if n == 0:
                 continue
+            node_energy_j = trapezoidal_energy(ts_profiles, elapsed)
+            avg_power_w = node_energy_j / elapsed if elapsed > 0 else 0.0
+            total_energy_j += node_energy_j
             node_stats.append(
                 NodePowerStats(
                     node_id=node_id,
                     samples=n,
-                    avg_sys_power=sum(p.sys_power for p in profiles) / n,
+                    avg_sys_power=avg_power_w,
                 )
             )
 
-        total_avg_sys = sum(ns.avg_sys_power for ns in node_stats)
+        total_avg_sys_w = total_energy_j / elapsed if elapsed > 0 else 0.0
         return PowerUsage(
             elapsed_seconds=elapsed,
             nodes=node_stats,
-            total_avg_sys_power_watts=total_avg_sys,
-            total_energy_joules=total_avg_sys * elapsed,
+            total_avg_sys_power_watts=total_avg_sys_w,
+            total_energy_joules=total_energy_j,
         )
+
+
+def trapezoidal_energy(
+    ts_profiles: list[tuple[float, SystemPerformanceProfile]],
+    elapsed: float,
+) -> float:
+    """Integrate sys_power(t) over the sample window using the trapezoidal rule.
+    First sample is anchored at t=0 and last at t=elapsed (set by `run` /
+    `result`), so the integral spans the full request interval. Falls back to
+    power * elapsed when only one sample exists (constant-power assumption)."""
+    if len(ts_profiles) == 1:
+        return ts_profiles[0][1].sys_power * elapsed
+    energy_j = 0.0
+    for i in range(1, len(ts_profiles)):
+        t_prev, p_prev = ts_profiles[i - 1]
+        t_cur, p_cur = ts_profiles[i]
+        dt = t_cur - t_prev
+        if dt <= 0:
+            continue
+        energy_j += (p_prev.sys_power + p_cur.sys_power) / 2.0 * dt
+    return energy_j

src/exo/utils/tests/conftest.py

@@ -0,0 +1,8 @@
+import multiprocessing as mp
+
+import pytest
+
+
+@pytest.fixture(scope="session", autouse=True)
+def mp_force_spawn():
+    mp.set_start_method("spawn", force=True)

src/exo/utils/tests/test_pidfile.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+import gc
+import os
+import subprocess
+import sys
+import textwrap
+from pathlib import Path
+from typing import Final
+
+import exo.utils.pidfile as pidfile
+import pytest
+from exo.utils.pidfile import acquire_exo_pidfile
+
+_CHILD_ACQUIRE_PIDFILE_SCRIPT: Final = textwrap.dedent(
+    """
+    import sys
+    from pathlib import Path
+    from unittest.mock import patch
+
+    import exo.utils.pidfile as pidfile
+    from exo.utils.pidfile import PidfileLockError, acquire_exo_pidfile
+
+    with patch.object(pidfile, "EXO_PID_FILE", Path(sys.argv[1])):
+        try:
+            handle = acquire_exo_pidfile()
+        except PidfileLockError as exception:
+            print(str(exception))
+            raise SystemExit(73) from exception
+
+        del handle
+    """
+)
+
+
+def _use_pidfile_path(monkeypatch: pytest.MonkeyPatch, path: Path) -> None:
+    monkeypatch.setattr(pidfile, "EXO_PID_FILE", path)
+
+
+def _run_child_acquire_pidfile(path: Path) -> subprocess.CompletedProcess[str]:
+    return subprocess.run(
+        [sys.executable, "-c", _CHILD_ACQUIRE_PIDFILE_SCRIPT, str(path)],
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+
+
+def test_acquire_exo_pidfile_writes_current_pid_and_removes_on_drop(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    path = tmp_path / "exo.pid"
+    _use_pidfile_path(monkeypatch, path)
+
+    handle = acquire_exo_pidfile()
+    assert path.read_text() == str(os.getpid())
+
+    del handle
+    gc.collect()
+
+    assert not path.exists()
+
+
+def test_acquire_exo_pidfile_rejects_second_process(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    path = tmp_path / "exo.pid"
+    _use_pidfile_path(monkeypatch, path)
+
+    handle = acquire_exo_pidfile()
+    try:
+        blocked_child = _run_child_acquire_pidfile(path)
+        assert blocked_child.returncode == 73
+        assert "Failed to acquire EXO pidfile" in blocked_child.stdout
+    finally:
+        del handle
+        gc.collect()
+
+    unblocked_child = _run_child_acquire_pidfile(path)
+    assert unblocked_child.returncode == 0
+    assert unblocked_child.stdout == ""

src/exo/utils/tests/test_power_sampler.py

@@ -111,6 +111,36 @@ async def test_empty_state() -> None:
     assert result.total_energy_joules == 0.0
 
 
+def test_trapezoidal_unit_dt_weighting() -> None:
+    """Pure unit test on the integration helper. Crafted samples where the
+    arithmetic mean is wildly wrong vs the time-weighted result."""
+    from exo.utils.power_sampler import trapezoidal_energy
+
+    # 5 s window. Power = 10 W for the first 4.9 s, then 100 W for the last 0.1 s.
+    # Three samples: t=0 W=10, t=4.9 W=10, t=5.0 W=100.
+    samples = [
+        (0.0, _make_profile(10.0)),
+        (4.9, _make_profile(10.0)),
+        (5.0, _make_profile(100.0)),
+    ]
+    energy = trapezoidal_energy(samples, elapsed=5.0)
+    # (10+10)/2 * 4.9 + (10+100)/2 * 0.1 = 49 + 5.5 = 54.5 J
+    assert abs(energy - 54.5) < 1e-9
+    avg = energy / 5.0  # 10.9 W
+    # Arithmetic mean of the three samples would be (10+10+100)/3 ≈ 40 W.
+    # Trapezoidal correctly weights each segment by its dt.
+    assert abs(avg - 10.9) < 1e-9
+
+
+def test_trapezoidal_unit_single_sample() -> None:
+    """One sample: no window to integrate over, so fall back to constant power
+    over the elapsed duration."""
+    from exo.utils.power_sampler import trapezoidal_energy
+
+    samples = [(0.0, _make_profile(42.0))]
+    assert trapezoidal_energy(samples, elapsed=3.0) == 42.0 * 3.0
+
+
 async def test_result_stops_sampling() -> None:
     """Calling result() should stop the sampler's run loop."""
     state: dict[NodeId, SystemPerformanceProfile] = {

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

@@ -143,6 +143,7 @@ class ImageEngine(Engine):
         Generator[tuple[TaskId, Chunk | FinishedResponse | CancelledResponse]] | None
     ) = field(init=False, default=None)
     queue: deque[ImageTask] = field(init=False, default_factory=deque)
+    _cancelled_tasks: set[TaskId] = field(init=False, default_factory=set)
 
     def warmup(self) -> None:
         image = warmup_image_generator(model=self.image_model)
@@ -168,7 +169,11 @@ class ImageEngine(Engine):
             task = self.queue.popleft()
             self.current_gen = self._run_image_task(task.task_id, task.task_params)
             resp = next(self.current_gen, None)
-        return (resp,) if resp is not None else ()
+        return (
+            (resp,)
+            if resp is not None and _is_primary_output_node(self.shard_metadata)
+            else ()
+        )
 
     def close(self) -> None:
         with contextlib.suppress(NameError, AttributeError):

src/exo/worker/main.py

@@ -4,29 +4,25 @@ from datetime import datetime, timezone
 
 import anyio
 from anyio import fail_after, to_thread
+from exo_net import PySession
 from loguru import logger
 
 from exo.api.types import ImageEditsTaskParams
 from exo.download.download_utils import is_read_only_model_dir, resolve_existing_model
-from exo.routing.event_router import EventRouter
-from exo.routing.snapshot_receiver import SnapshotReceiver
 from exo.shared.apply import apply
 from exo.shared.constants import EXO_MAX_INSTANCE_RETRIES
-from exo.shared.models.model_cards import ModelId, add_to_card_cache, delete_custom_card
+from exo.shared.models.model_cards import ModelId, card_cache
 from exo.shared.types.chunks import InputImageChunk
 from exo.shared.types.commands import (
     DeleteInstance,
-    ForwarderCommand,
     ForwarderDownloadCommand,
-    RequestSnapshot,
     StartDownload,
 )
-from exo.shared.types.common import CommandId, NodeId, SessionId, SystemId
+from exo.shared.types.common import CommandId, NodeId, SystemId
 from exo.shared.types.events import (
-    CustomModelCardAdded,
-    CustomModelCardDeleted,
     Event,
     IndexedEvent,
+    InputChunkReceived,
     InstanceDeleted,
     NodeDownloadProgress,
     NodeGatheredInfo,
@@ -36,7 +32,6 @@ from exo.shared.types.events import (
     TopologyEdgeDeleted,
 )
 from exo.shared.types.multiaddr import Multiaddr
-from exo.shared.types.snapshots import SnapshotChunk
 from exo.shared.types.state import State
 from exo.shared.types.tasks import (
     CancelTask,
@@ -62,32 +57,25 @@ from exo.utils.task_group import TaskGroup
 from exo.worker.plan import plan
 from exo.worker.runner.supervisor import RunnerSupervisor
 
-_SNAPSHOT_FETCH_TIMEOUT_SECONDS = 30
-
 
 class Worker:
     def __init__(
         self,
         node_id: NodeId,
-        session_id: SessionId,
         *,
-        event_router: EventRouter,
         event_receiver: Receiver[IndexedEvent],
         event_sender: Sender[Event],
-        snapshot_chunk_receiver: Receiver[SnapshotChunk],
+        session: PySession,
         # This is for requesting updates. It doesn't need to be a general command sender right now,
         # but I think it's the correct way to be thinking about commands
-        command_sender: Sender[ForwarderCommand],
         download_command_sender: Sender[ForwarderDownloadCommand],
         api_port: int,
     ):
         self.node_id: NodeId = node_id
-        self.session_id: SessionId = session_id
-        self.event_router = event_router
         self.event_receiver = event_receiver
         self.event_sender = event_sender
-        self.snapshot_chunk_receiver = snapshot_chunk_receiver
-        self.command_sender = command_sender
+        self.session = session
+        self.command_sender = session.net_sender("orchestrator")
         self.download_command_sender = download_command_sender
         self.api_port = api_port
 
@@ -116,57 +104,22 @@ class Worker:
 
         try:
             async with self._tg as tg:
-                tg.start_soon(self._bootstrap_then_run, info_gatherer, info_recv)
+                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)
+                tg.start_soon(self._reconcile_custom_cards)
+
         finally:
             # Actual shutdown code - waits for all tasks to complete before executing.
             logger.info("Stopping Worker")
             self.event_sender.close()
-            self.snapshot_chunk_receiver.close()
-            self.command_sender.close()
             self.download_command_sender.close()
             for runner in self.runners.values():
                 runner.shutdown()
             self._stopped.set()
 
-    async def _bootstrap_then_run(
-        self, info_gatherer: InfoGatherer, info_recv: Receiver[GatheredInfo]
-    ) -> None:
-        await self._fetch_snapshot()
-        self._sync_input_views_from_state()
-        self._tg.start_soon(info_gatherer.run)
-        self._tg.start_soon(self._forward_info, info_recv)
-        self._tg.start_soon(self.plan_step)
-        self._tg.start_soon(self._event_applier)
-        self._tg.start_soon(self._poll_connection_updates)
-
-    async def _fetch_snapshot(self) -> None:
-        receiver = SnapshotReceiver(self.node_id, self.session_id)
-        await self.command_sender.send(
-            ForwarderCommand(
-                origin=self._system_id,
-                command=RequestSnapshot(requester_node_id=self.node_id),
-            )
-        )
-
-        with anyio.move_on_after(_SNAPSHOT_FETCH_TIMEOUT_SECONDS):
-            with self.snapshot_chunk_receiver as chunks:
-                async for chunk in chunks:
-                    received = receiver.ingest(chunk)
-                    if received is None:
-                        continue
-                    self.state = received.state
-                    self.event_router.set_buffer_start(
-                        received.last_event_applied_idx + 1
-                    )
-                    logger.info(
-                        f"Worker bootstrapped from snapshot at idx "
-                        f"{received.last_event_applied_idx}"
-                    )
-                    return
-        logger.info(
-            "No snapshot received before timeout; falling back to full event-log replay"
-        )
-
     async def _forward_info(self, recv: Receiver[GatheredInfo]):
         with recv as info_stream:
             async for info in info_stream:
@@ -181,8 +134,6 @@ class Worker:
     async def _event_applier(self):
         with self.event_receiver as events:
             async for event in events:
-                if event.idx <= self.state.last_event_applied_idx:
-                    continue
                 # 2. for each event, apply it to the state
                 self.state = apply(self.state, event=event)
                 event = event.event
@@ -190,41 +141,48 @@ class Worker:
                 if isinstance(event, InstanceDeleted):
                     self._instance_backoff.reset(event.instance_id)
 
-                if isinstance(event, CustomModelCardAdded):
-                    await event.model_card.save_to_custom_dir()
-                    add_to_card_cache(event.model_card)
+                # Buffer input image chunks for image editing
+                if isinstance(event, InputChunkReceived):
+                    cmd_id = event.command_id
+                    if cmd_id not in self.input_chunk_buffer:
+                        self.input_chunk_buffer[cmd_id] = {}
+                        self.input_chunk_counts[cmd_id] = event.chunk.total_chunks
 
-                if isinstance(event, CustomModelCardDeleted):
-                    await delete_custom_card(event.model_id)
+                    self.input_chunk_buffer[cmd_id][event.chunk.chunk_index] = (
+                        event.chunk
+                    )
+                    if (
+                        len(self.input_chunk_buffer[cmd_id])
+                        == self.input_chunk_counts[cmd_id]
+                    ):
+                        per_image: defaultdict[int, list[InputImageChunk]] = (
+                            defaultdict(list)
+                        )
+                        for chunk in self.input_chunk_buffer[cmd_id].values():
+                            per_image[chunk.image_index].append(chunk)
+                        for chunks_for_image in per_image.values():
+                            sorted_chunks = sorted(
+                                chunks_for_image, key=lambda c: c.chunk_index
+                            )
+                            img = Base64Image("".join(c.data for c in sorted_chunks))
+                            self.image_cache[
+                                Base64ImageHash(
+                                    hashlib.sha256(img.encode("ascii")).hexdigest()
+                                )
+                            ] = img
 
-                self._sync_input_views_from_state()
+    async def _reconcile_custom_cards(self) -> None:
+        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:
+                    continue
+                await card_cache.save(card)
 
-    def _sync_input_views_from_state(self) -> None:
-        self.input_chunk_buffer = {
-            command_id: dict(chunks)
-            for command_id, chunks in self.state.input_chunks.items()
-        }
-        self.input_chunk_counts = {
-            command_id: next(iter(chunks.values())).total_chunks
-            for command_id, chunks in self.input_chunk_buffer.items()
-            if chunks
-        }
-
-        self.image_cache = {}
-        for command_id, chunks in self.input_chunk_buffer.items():
-            expected_chunks = self.input_chunk_counts.get(command_id)
-            if expected_chunks is None or len(chunks) != expected_chunks:
-                continue
-
-            per_image: defaultdict[int, list[InputImageChunk]] = defaultdict(list)
-            for chunk in chunks.values():
-                per_image[chunk.image_index].append(chunk)
-            for chunks_for_image in per_image.values():
-                sorted_chunks = sorted(chunks_for_image, key=lambda c: c.chunk_index)
-                image = Base64Image("".join(chunk.data for chunk in sorted_chunks))
-                self.image_cache[
-                    Base64ImageHash(hashlib.sha256(image.encode("ascii")).hexdigest())
-                ] = image
+            for card in await card_cache.list_all():
+                if card.model_id not in target:
+                    await card_cache.pop(card.model_id)
 
     async def plan_step(self):
         while True:
@@ -236,7 +194,7 @@ class Worker:
                 self.state.instances,
                 self.state.runners,
                 self.state.tasks,
-                self.state.input_chunks,
+                self.input_chunk_buffer,
                 self.image_cache,
                 self._instance_backoff,
                 self._download_backoff,
@@ -251,10 +209,9 @@ class Worker:
                         f"Instance {iid} exceeded {EXO_MAX_INSTANCE_RETRIES} retries, requesting deletion"
                     )
                     await self.command_sender.send(
-                        ForwarderCommand(
-                            origin=self._system_id,
-                            command=DeleteInstance(instance_id=iid),
-                        )
+                        DeleteInstance(instance_id=iid)
+                        .model_dump_json()
+                        .encode("utf-8")
                     )
                     continue
 
@@ -368,9 +325,15 @@ class Worker:
                             advanced_params=task.task_params.advanced_params,
                         ),
                     )
+                    # Cleanup buffers
+                    if cmd_id in self.input_chunk_buffer:
+                        del self.input_chunk_buffer[cmd_id]
+                    if cmd_id in self.input_chunk_counts:
+                        del self.input_chunk_counts[cmd_id]
                     await self._start_runner_task(modified_task)
 
                 case TextGeneration() if task.task_params.image_hashes:
+                    cmd_id = task.command_id
                     resolved_images = [
                         self.image_cache[h]
                         for _, h in sorted(task.task_params.image_hashes.items())
@@ -382,6 +345,10 @@ class Worker:
                             )
                         }
                     )
+                    if cmd_id in self.input_chunk_buffer:
+                        del self.input_chunk_buffer[cmd_id]
+                    if cmd_id in self.input_chunk_counts:
+                        del self.input_chunk_counts[cmd_id]
                     await self._start_runner_task(modified_task)
                 case LoadModel(instance_id=instance_id):
                     if (instance := self.state.instances.get(instance_id)) is not None:
@@ -407,6 +374,7 @@ class Worker:
         runner = RunnerSupervisor.create(
             bound_instance=task.bound_instance,
             event_sender=self.event_sender.clone(),
+            session=self.session,
         )
         self.runners[task.bound_instance.bound_runner_id] = runner
         self._tg.start_soon(runner.run)

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

@@ -138,8 +138,10 @@ class SequentialGenerator(Engine):
     def agree_on_tasks(self) -> None:
         """Agree between all ranks about the task ordering (some may have received in different order or not at all)."""
         agreed, different = mx_all_gather_tasks(self._maybe_queue, self.group)
-        self._queue.extend(task for task in self._maybe_queue if task in agreed)
-        self._maybe_queue = [task for task in self._maybe_queue if task in different]
+        # Extend from `agreed` (sorted by task_id on all ranks) to guarantee every
+        # rank enqueues tasks in the same order, preventing TP collective deadlocks.
+        self._queue.extend(agreed)
+        self._maybe_queue = list(different)
 
     def agree_on_cancellations(self) -> None:
         """Agree between all ranks about which tasks to cancel."""
@@ -197,9 +199,14 @@ class SequentialGenerator(Engine):
             self._active = None
             raise
 
-        return itertools.chain(
-            output,
-            map(lambda task: (task, CancelledResponse()), self._cancelled_tasks),
+        return filter(
+            lambda chunk: (
+                not isinstance(chunk[1], GenerationChunk) or self.device_rank == 0
+            ),
+            itertools.chain(
+                output,
+                map(lambda task: (task, CancelledResponse()), self._cancelled_tasks),
+            ),
         )
 
     def _start_next(self) -> None:
@@ -368,8 +375,10 @@ class BatchGenerator(Engine):
     def agree_on_tasks(self) -> None:
         """Agree between all ranks about the task ordering (some may have received in different order or not at all)."""
         agreed, different = mx_all_gather_tasks(self._maybe_queue, self.group)
-        self._queue.extend(task for task in self._maybe_queue if task in agreed)
-        self._maybe_queue = [task for task in self._maybe_queue if task in different]
+        # Extend from `agreed` (sorted by task_id on all ranks) to guarantee every
+        # rank enqueues tasks in the same order, preventing TP collective deadlocks.
+        self._queue.extend(agreed)
+        self._maybe_queue = list(different)
 
     def agree_on_cancellations(self) -> None:
         """Agree between all ranks about which tasks to cancel."""
@@ -449,7 +458,12 @@ class BatchGenerator(Engine):
                 output.append((task.task_id, FinishedResponse()))
                 del self._active_tasks[uid]
 
-        return itertools.chain(output, self._apply_cancellations())
+        return filter(
+            lambda chunk: (
+                not isinstance(chunk[1], GenerationChunk) or self.device_rank == 0
+            ),
+            itertools.chain(output, self._apply_cancellations()),
+        )
 
     def _apply_cancellations(
         self,

src/exo/worker/runner/runner.py

@@ -390,5 +390,5 @@ class Runner:
         chunk: Chunk,
         command_id: CommandId,
     ):
-        if self.device_rank == 0:
-            self.event_sender.send(ChunkGenerated(command_id=command_id, chunk=chunk))
+        assert isinstance(self.generator, Engine)
+        self.event_sender.send(ChunkGenerated(command_id=command_id, chunk=chunk))

src/exo/worker/runner/supervisor.py

@@ -10,9 +10,11 @@ from anyio import (
     ClosedResourceError,
     to_thread,
 )
+from exo_net import NetSender, PySession
 from loguru import logger
 
 from exo.shared.types.chunks import ErrorChunk
+from exo.shared.types.commands import CommandId
 from exo.shared.types.events import (
     ChunkGenerated,
     Event,
@@ -45,20 +47,18 @@ from exo.utils.channels import MpReceiver, MpSender, Sender, mp_channel
 from exo.utils.task_group import TaskGroup
 from exo.worker.runner.bootstrap import entrypoint
 
-PREFILL_TIMEOUT_SECONDS = 60
-DECODE_TIMEOUT_SECONDS = 5
-
 
 @dataclass(eq=False)
 class RunnerSupervisor:
     shard_metadata: ShardMetadata
     bound_instance: BoundInstance
     runner_process: mp.Process
-    initialize_timeout: float
     _ev_recv: MpReceiver[Event]
     _task_sender: MpSender[Task]
     _event_sender: Sender[Event]
     _cancel_sender: MpSender[TaskId]
+    session: PySession
+
     _tg: TaskGroup = field(default_factory=TaskGroup, init=False)
     status: RunnerStatus = field(default_factory=RunnerIdle, init=False)
     pending: dict[TaskId, anyio.Event] = field(default_factory=dict, init=False)
@@ -75,7 +75,7 @@ class RunnerSupervisor:
         *,
         bound_instance: BoundInstance,
         event_sender: Sender[Event],
-        initialize_timeout: float = 400,
+        session: PySession,
     ) -> Self:
         ev_send, ev_recv = mp_channel[Event]()
         task_sender, task_recv = mp_channel[Task]()
@@ -99,11 +99,11 @@ class RunnerSupervisor:
             bound_instance=bound_instance,
             shard_metadata=shard_metadata,
             runner_process=runner_process,
-            initialize_timeout=initialize_timeout,
             _ev_recv=ev_recv,
             _task_sender=task_sender,
             _cancel_sender=cancel_sender,
             _event_sender=event_sender,
+            session=session,
         )
 
         return self
@@ -210,9 +210,26 @@ class RunnerSupervisor:
             await self._check_runner(TimeoutError("cancel pipe blocked"))
 
     async def _forward_events(self):
+        pubs: dict[CommandId, NetSender] = {}
         try:
             with self._ev_recv as events:
                 async for event in events:
+                    if isinstance(event, ChunkGenerated):
+                        if (pub := pubs.get(event.command_id, None)) is None:
+                            pub = pubs[event.command_id] = self.session.net_sender(
+                                f"runners/{self.bound_instance.bound_runner_id}/active_tasks/{event.command_id}/chunks"
+                            )
+                        sent = await pub.send(
+                            event.chunk.model_dump_json().encode("utf-8")
+                        )
+                        if not sent:
+                            logger.warning(
+                                "api node closed communication, dropping chunk"
+                            )
+                        if event.chunk.finish_reason is not None:
+                            pubs.pop(event.command_id, None)
+                        continue
+
                     if isinstance(event, RunnerStatusUpdated):
                         self.status = event.runner_status
                     if isinstance(event, TaskAcknowledged):
@@ -275,17 +292,18 @@ class RunnerSupervisor:
         for task in self.in_progress.values():
             if isinstance(task, (TextGeneration, ImageGeneration, ImageEdits)):
                 with anyio.CancelScope(shield=True):
-                    await self._event_sender.send(
-                        ChunkGenerated(
-                            command_id=task.command_id,
-                            chunk=ErrorChunk(
-                                model=self.shard_metadata.model_card.model_id,
-                                error_message=(
-                                    "Runner shutdown before completing command "
-                                    f"({cause})"
-                                ),
+                    send = self.session.net_sender(
+                        f"runners/{self.bound_instance.bound_runner_id}/active_tasks/{task.command_id}/chunks"
+                    )
+                    await send.send(
+                        ErrorChunk(
+                            model=self.shard_metadata.model_card.model_id,
+                            error_message=(
+                                f"Runner shutdown before completing command ({cause})"
                             ),
                         )
+                        .model_dump_json()
+                        .encode("utf-8")
                     )
 
         try:

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

@@ -16,7 +16,7 @@ from exo.download.download_utils import (
     fetch_file_list_with_cache,
     resolve_model_dir,
 )
-from exo.shared.models.model_cards import ModelCard, ModelId, get_model_cards
+from exo.shared.models.model_cards import ModelCard, ModelId, card_cache
 from exo.worker.engines.mlx.utils_mlx import (
     get_eos_token_ids_for_model,
     load_tokenizer_for_model_id,
@@ -76,7 +76,7 @@ def get_test_models() -> list[ModelCard]:
     """Get a representative sample of models to test."""
     # Pick one model from each family to test
     families: dict[str, ModelCard] = {}
-    for card in asyncio.run(get_model_cards()):
+    for card in asyncio.run(card_cache.list_all()):
         # Extract family name (e.g., "llama-3.1" from "llama-3.1-8b")
         parts = card.model_id.short().split("-")
         family = "-".join(parts[:2]) if len(parts) >= 2 else parts[0]
@@ -298,7 +298,7 @@ async def test_tokenizer_special_tokens(model_card: ModelCard) -> None:
 async def test_kimi_tokenizer_specifically():
     """Test Kimi tokenizer with its specific patches and quirks."""
     kimi_models = [
-        card for card in await get_model_cards() if "kimi" in card.model_id.lower()
+        card for card in await card_cache.list_all() if "kimi" in card.model_id.lower()
     ]
 
     if not kimi_models:
@@ -350,7 +350,7 @@ async def test_glm_tokenizer_specifically():
 
     glm_model_cards = [
         card
-        for card in await get_model_cards()
+        for card in await card_cache.list_all()
         if contains(card, "glm")
         and not contains(card, "-5")
         and not contains(card, "4.7")

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

@@ -1,126 +0,0 @@
-# pyright: reportPrivateUsage=false
-
-import hashlib
-
-import anyio
-import pytest
-import zstandard
-
-from exo.routing.event_router import EventRouter
-from exo.shared.types.commands import (
-    ForwarderCommand,
-    ForwarderDownloadCommand,
-    RequestSnapshot,
-)
-from exo.shared.types.common import NodeId, SessionId
-from exo.shared.types.events import (
-    Event,
-    GlobalForwarderEvent,
-    IndexedEvent,
-    LocalForwarderEvent,
-    TestEvent,
-)
-from exo.shared.types.snapshots import SnapshotChunk, SnapshotTransferId
-from exo.shared.types.state import State
-from exo.utils.channels import Receiver, Sender, channel
-from exo.worker.main import Worker
-
-
-def _snapshot_chunk(
-    state: State, *, requester_node_id: NodeId, session_id: SessionId
-) -> SnapshotChunk:
-    body = zstandard.ZstdCompressor().compress(state.model_dump_json().encode("utf-8"))
-    return SnapshotChunk.from_data(
-        data=body,
-        transfer_id=SnapshotTransferId("transfer-1"),
-        requester_node_id=requester_node_id,
-        session_id=session_id,
-        schema_version=state.schema_version,
-        last_event_applied_idx=state.last_event_applied_idx,
-        chunk_index=0,
-        total_chunks=1,
-        sha256_hex=hashlib.sha256(body).hexdigest(),
-    )
-
-
-def _worker(
-    node_id: NodeId, session_id: SessionId
-) -> tuple[
-    Worker,
-    EventRouter,
-    Receiver[ForwarderCommand],
-    Sender[SnapshotChunk],
-    Sender[IndexedEvent],
-]:
-    router_command_sender, _router_command_receiver = channel[ForwarderCommand]()
-    _global_event_sender, global_event_receiver = channel[GlobalForwarderEvent]()
-    local_event_sender, _local_event_receiver = channel[LocalForwarderEvent]()
-    event_router = EventRouter(
-        session_id=session_id,
-        command_sender=router_command_sender,
-        external_inbound=global_event_receiver,
-        external_outbound=local_event_sender,
-    )
-
-    event_sender, event_receiver = channel[IndexedEvent]()
-    local_event_output_sender, _local_event_output_receiver = channel[Event]()
-    command_sender, command_receiver = channel[ForwarderCommand]()
-    download_command_sender, _download_command_receiver = channel[
-        ForwarderDownloadCommand
-    ]()
-    snapshot_sender, snapshot_receiver = channel[SnapshotChunk]()
-
-    worker = Worker(
-        node_id,
-        session_id,
-        event_router=event_router,
-        event_receiver=event_receiver,
-        event_sender=local_event_output_sender,
-        snapshot_chunk_receiver=snapshot_receiver,
-        command_sender=command_sender,
-        download_command_sender=download_command_sender,
-        api_port=52415,
-    )
-    return worker, event_router, command_receiver, snapshot_sender, event_sender
-
-
-@pytest.mark.asyncio
-async def test_worker_fetch_snapshot_applies_state_and_fast_forwards_router() -> None:
-    node_id = NodeId("worker")
-    session_id = SessionId(master_node_id=NodeId("master"), election_clock=1)
-    worker, event_router, command_receiver, snapshot_sender, _event_sender = _worker(
-        node_id, session_id
-    )
-    state = State(last_event_applied_idx=7)
-
-    async with anyio.create_task_group() as tg:
-        tg.start_soon(worker._fetch_snapshot)
-        command = await command_receiver.receive()
-        assert isinstance(command.command, RequestSnapshot)
-        assert command.command.requester_node_id == node_id
-
-        await snapshot_sender.send(
-            _snapshot_chunk(state, requester_node_id=node_id, session_id=session_id)
-        )
-
-    assert worker.state.last_event_applied_idx == 7
-    assert event_router.event_buffer.next_idx_to_release == 8
-
-
-@pytest.mark.asyncio
-async def test_worker_event_applier_ignores_events_covered_by_snapshot() -> None:
-    node_id = NodeId("worker")
-    session_id = SessionId(master_node_id=NodeId("master"), election_clock=1)
-    worker, _event_router, _command_receiver, _snapshot_sender, event_sender = _worker(
-        node_id, session_id
-    )
-    worker.state = State(last_event_applied_idx=7)
-
-    async with anyio.create_task_group() as tg:
-        tg.start_soon(worker._event_applier)
-        await event_sender.send(IndexedEvent(idx=7, event=TestEvent()))
-        await event_sender.send(IndexedEvent(idx=8, event=TestEvent()))
-
-        while worker.state.last_event_applied_idx != 8:
-            await anyio.sleep(0.001)
-        tg.cancel_scope.cancel()

tests/conftest.py

@@ -0,0 +1,181 @@
+# type: ignore
+"""Pytest configuration for marker-driven exo integration tests.
+
+Test authors declare requirements via markers:
+
+    @pytest.mark.cluster(count=2, thunderbolt='a2a')
+    @pytest.mark.instance('mlx-community/Llama-3.2-1B-Instruct-4bit',
+                          sharding='tensor', comm='jaccl')
+    def test_jaccl_inference(session):
+        resp = session.chat('What is 2+2?')
+        assert '4' in resp
+
+Clusters are cached by `ClusterSpec`; tests with the same cluster_spec
+share a deployment. Each test places its own instance (matching its
+`@pytest.mark.instance`), and instances are cleaned up after the test.
+
+Run with:
+    uv run pytest tests/ -v
+    uv run pytest tests/ -v --hosts s2,s4,s9,s10
+"""
+
+from __future__ import annotations
+
+import contextlib
+import json
+
+import pytest
+from exo_tools.cluster import ClusterInfo, EcoSession
+from exo_tools.harness import cleanup_all_instances, place_instance
+
+from .framework import (
+    ClusterSpec,
+    Session,
+    parse_cluster_marker,
+    parse_instance_marker,
+)
+
+# Single eco session for the entire test process.
+eco = EcoSession(user_prefix="test")
+
+# Cluster cache keyed by ClusterSpec — tests with the same spec share a deployment.
+# Cleared at session teardown.
+_cluster_cache: dict[ClusterSpec, ClusterInfo] = {}
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--hosts",
+        default=None,
+        help="Comma-separated list of hosts (e.g. s2,s4,s9,s10). "
+        "Overrides constraint-based reservation.",
+    )
+
+
+def pytest_configure(config):
+    """Register custom markers."""
+    config.addinivalue_line(
+        "markers",
+        "cluster(count=N, thunderbolt=Thunderbolt|None, min_memory=GB, chip=PATTERN): "
+        "declare cluster requirements for a test",
+    )
+    config.addinivalue_line(
+        "markers",
+        "instance(model_id, sharding=Sharding, comm=Comm, min_nodes=N): "
+        "declare instance placement for a test",
+    )
+
+
+def pytest_report_header(config):
+    """Show the eco user and hosts for this test session."""
+    hosts = config.getoption("--hosts")
+    lines = [f"eco user: {eco.user}"]
+    if hosts:
+        lines.append(f"hosts override: {hosts}")
+    return lines
+
+
+@pytest.fixture(scope="session")
+def _host_pool(request) -> list[str] | None:
+    raw = request.config.getoption("--hosts")
+    if raw:
+        return [h.strip() for h in raw.split(",") if h.strip()]
+    return None
+
+
+@pytest.fixture
+def session(request, _host_pool) -> Session:
+    """Per-test fixture providing a Session matching the test's markers.
+
+    Reads @pytest.mark.cluster and @pytest.mark.instance from the test, deploys
+    a matching cluster (cached across tests with the same spec), places the
+    model, and yields a Session for the test to interact with. Cleans up the
+    instance after the test, and invalidates the cluster cache if the test
+    left nodes disconnected.
+    """
+    cluster_marker = request.node.get_closest_marker("cluster")
+    instance_marker = request.node.get_closest_marker("instance")
+
+    cluster_spec = parse_cluster_marker(cluster_marker)
+    instance_spec = parse_instance_marker(instance_marker)
+
+    # Deploy or reuse a cluster matching the spec
+    cluster = _cluster_cache.get(cluster_spec)
+    if cluster is None:
+        if _host_pool:
+            cluster = eco.start_deploy(
+                hosts=_host_pool[: cluster_spec.count], wait=True
+            )
+        else:
+            cluster = eco.start_deploy(
+                count=cluster_spec.count,
+                thunderbolt=cluster_spec.thunderbolt,
+                chip=cluster_spec.chip,
+                min_memory_gb=cluster_spec.min_memory_gb,
+                wait=True,
+            )
+        _cluster_cache[cluster_spec] = cluster
+
+    # Place an instance for this test if the test specified one
+    instance_id = None
+    if instance_spec is not None:
+        client = cluster.make_client()
+        instance_id = place_instance(
+            client,
+            instance_spec.model_id,
+            sharding=instance_spec.sharding,
+            comm=instance_spec.comm,
+            min_nodes=instance_spec.min_nodes,
+        )
+
+    sess = Session(
+        cluster=cluster,
+        eco=eco,
+        instance_spec=instance_spec,
+        instance_id=instance_id,
+    )
+
+    yield sess
+
+    # ---- Teardown ----
+
+    # If the test left nodes disconnected, invalidate the cluster cache and
+    # stop the cluster so the next test deploys fresh.
+    if sess._stopped_hosts:
+        _cluster_cache.pop(cluster_spec, None)
+        with contextlib.suppress(Exception):
+            eco.stop(sess.cluster.hosts)
+        return
+
+    # Otherwise, clean up any instances created during the test
+    with contextlib.suppress(Exception):
+        cleanup_all_instances(sess.client)
+
+
+# ---------------------------------------------------------------------------
+# Session-level teardown — stop all cached clusters
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(scope="session", autouse=True)
+def _teardown_clusters():
+    yield
+    for cluster in _cluster_cache.values():
+        with contextlib.suppress(Exception):
+            eco.stop(cluster.hosts)
+    _cluster_cache.clear()
+
+
+def pytest_runtest_makereport(item, call):
+    """Attach cluster logs to the test report when a test fails."""
+    if call.when != "call" or call.excinfo is None:
+        return
+
+    sess = item.funcargs.get("session")
+    if sess is None:
+        return
+    try:
+        logs = eco.logs(sess.cluster.hosts, lines=200)
+        item.add_report_section("call", "Cluster Logs", json.dumps(logs, indent=2))
+    except Exception:
+        pass

tests/framework.py

@@ -0,0 +1,199 @@
+"""Marker-driven test framework for exo integration tests.
+
+Test authors declare requirements via markers:
+
+    @pytest.mark.cluster(count=2, thunderbolt='a2a')
+    @pytest.mark.instance('mlx-community/Llama-3.2-1B-Instruct-4bit',
+                          sharding='tensor', comm='jaccl')
+    def test_jaccl_inference(session):
+        resp = session.chat('What is 2+2?')
+        assert '4' in resp
+
+The `session` fixture reads the markers, deploys the cluster, places the
+instance, and provides a `Session` object. All cluster/instance orchestration
+lives in `exo_tools.harness`; this module is purely the pytest-facing layer.
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field
+from typing import Any
+
+from exo_tools.client import ExoClient
+from exo_tools.cluster import (
+    Chip,
+    ClusterInfo,
+    EcoSession,
+    Thunderbolt,
+    make_client_from_url,
+)
+from exo_tools.harness import Comm, Sharding
+
+from exo.api.types.api import (
+    ChatCompletionChoice,
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+)
+
+DEFAULT_MODEL = "mlx-community/Llama-3.2-1B-Instruct-4bit"
+
+
+def _extract_content(resp: ChatCompletionResponse) -> str:
+    """Extract plain-text content from a non-streaming chat completion."""
+    choice = resp.choices[0]
+    if not isinstance(choice, ChatCompletionChoice):
+        raise RuntimeError(
+            f"Expected non-streaming choice, got {type(choice).__name__}"
+        )
+    content = choice.message.content
+    if not isinstance(content, str):
+        raise RuntimeError(f"Expected string content, got {type(content).__name__}")
+    return content
+
+
+@dataclass(frozen=True)
+class ClusterSpec:
+    count: int = 1
+    thunderbolt: Thunderbolt | None = None
+    min_memory_gb: float | None = None
+    chip: Chip | None = None
+
+
+@dataclass(frozen=True)
+class InstanceSpec:
+    model_id: str
+    sharding: Sharding = Sharding.PIPELINE
+    comm: Comm = Comm.RING
+    min_nodes: int = 1
+
+
+def parse_cluster_marker(marker) -> ClusterSpec:
+    if marker is None:
+        return ClusterSpec()
+    return ClusterSpec(
+        count=marker.kwargs.get("count", 1),
+        thunderbolt=marker.kwargs.get("thunderbolt"),
+        min_memory_gb=marker.kwargs.get("min_memory"),
+        chip=marker.kwargs.get("chip"),
+    )
+
+
+def parse_instance_marker(marker) -> InstanceSpec | None:
+    if marker is None:
+        return None
+    if not marker.args:
+        raise ValueError(
+            "@pytest.mark.instance requires a positional model_id argument"
+        )
+    return InstanceSpec(
+        model_id=marker.args[0],
+        sharding=marker.kwargs.get("sharding", Sharding.PIPELINE),
+        comm=marker.kwargs.get("comm", Comm.RING),
+        min_nodes=marker.kwargs.get("min_nodes", 1),
+    )
+
+
+@dataclass
+class Session:
+    cluster: ClusterInfo
+    eco: EcoSession
+    instance_spec: InstanceSpec | None = None
+    instance_id: str | None = None
+    _stopped_hosts: set[str] = field(default_factory=set)
+
+    @property
+    def client(self) -> ExoClient:
+        for host in self.cluster.hosts:
+            if host not in self._stopped_hosts:
+                return make_client_from_url(self.cluster.api_endpoints[host])
+        return self.cluster.make_client()
+
+    @property
+    def state(self) -> dict[str, Any]:
+        return self.client.request_json("GET", "/state") or {}
+
+    @property
+    def instances(self) -> dict[str, Any]:
+        return self.state.get("instances", {})
+
+    # ---- Inference ----
+
+    def chat(self, prompt: str, max_tokens: int = 100) -> str:
+        resp = self.chat_raw(prompt, max_tokens=max_tokens)
+        return _extract_content(resp)
+
+    def chat_raw(self, prompt: str, **kwargs: Any) -> ChatCompletionResponse:
+        if not self.instance_spec:
+            raise RuntimeError(
+                "No instance placed; add @pytest.mark.instance to the test"
+            )
+        max_tokens = kwargs.pop("max_tokens", 100)
+        request = ChatCompletionRequest.model_validate(
+            {
+                "model": self.instance_spec.model_id,
+                "messages": [{"role": "user", "content": prompt}],
+                "max_tokens": max_tokens,
+                **kwargs,
+            }
+        )
+        return self._post_chat(request)
+
+    def multi_turn(self, messages: list[dict[str, str]], max_tokens: int = 100) -> str:
+        if not self.instance_spec:
+            raise RuntimeError(
+                "No instance placed; add @pytest.mark.instance to the test"
+            )
+        request = ChatCompletionRequest.model_validate(
+            {
+                "model": self.instance_spec.model_id,
+                "messages": messages,
+                "max_tokens": max_tokens,
+            }
+        )
+        return _extract_content(self._post_chat(request))
+
+    def _post_chat(self, request: ChatCompletionRequest) -> ChatCompletionResponse:
+        raw = self.client.request_json(
+            "POST",
+            "/v1/chat/completions",
+            body=request.model_dump(exclude_none=True),
+        )
+        return ChatCompletionResponse.model_validate(raw)
+
+    def disconnect_node(self, index: int) -> None:
+        """Stop exo on a node and wait for the cluster to observe the disconnect."""
+        host = self.cluster.hosts[index]
+        self.eco.stop([host], keep=True)
+        self._stopped_hosts.add(host)
+
+    def reconnect_node(self, index: int) -> None:
+        """Restart a previously disconnected node into the existing namespace."""
+        host = self.cluster.hosts[index]
+        self.eco.start_hosts([host], namespace=self.cluster.namespace)
+        self._stopped_hosts.discard(host)
+
+    def wait_ready(
+        self, expected_nodes: int | None = None, timeout: float = 60
+    ) -> None:
+        """Wait until the cluster has exactly `expected_nodes` visible and reporting memory.
+
+        Defaults to the count of non-stopped hosts. Use this after
+        `disconnect_node` / `reconnect_node` to wait for the cluster to settle.
+        """
+        if expected_nodes is None:
+            expected_nodes = len(self.cluster.hosts) - len(self._stopped_hosts)
+        start = time.time()
+        while time.time() - start < timeout:
+            try:
+                state = self.state
+                identities = len(state.get("nodeIdentities", {}))
+                memory = len(state.get("nodeMemory", {}))
+                if identities == expected_nodes and memory == expected_nodes:
+                    return
+            except Exception:
+                pass
+            time.sleep(2.0)
+        raise TimeoutError(
+            f"Cluster did not reach exactly {expected_nodes} ready nodes within {timeout}s"
+        )

tests/headless_runner.py

@@ -1,264 +0,0 @@
-import socket
-from typing import Literal
-
-import anyio
-from fastapi import FastAPI
-from fastapi.responses import Response, StreamingResponse
-from hypercorn import Config
-from hypercorn.asyncio import serve  # pyright: ignore[reportUnknownVariableType]
-from loguru import logger
-from pydantic import BaseModel
-
-from exo.shared.constants import EXO_DEFAULT_MODELS_DIR
-from exo.shared.models.model_cards import ModelCard, ModelId
-from exo.shared.types.chunks import TokenChunk
-from exo.shared.types.commands import CommandId
-from exo.shared.types.common import Host, NodeId
-from exo.shared.types.events import ChunkGenerated, Event, RunnerStatusUpdated
-from exo.shared.types.tasks import (
-    ConnectToGroup,
-    LoadModel,
-    Shutdown,
-    StartWarmup,
-    Task,
-    TextGeneration,
-)
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
-from exo.shared.types.worker.instances import (
-    BoundInstance,
-    Instance,
-    InstanceId,
-    MlxJacclInstance,
-    MlxRingInstance,
-)
-from exo.shared.types.worker.runners import (
-    RunnerFailed,
-    RunnerId,
-    RunnerShutdown,
-    ShardAssignments,
-)
-from exo.shared.types.worker.shards import PipelineShardMetadata, TensorShardMetadata
-from exo.utils.channels import channel, mp_channel
-from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer
-from exo.worker.runner.bootstrap import entrypoint
-
-
-class Tests(BaseModel):
-    # list[hostname, ip addr]
-    devs: list[list[str]]
-    ibv_devs: list[list[str | None]] | None
-    model_id: ModelId
-    kind: Literal["ring", "jaccl", "both"]
-
-
-iid = InstanceId("im testing here")
-
-
-async def main():
-    logger.info("starting cool server majig")
-    cfg = Config()
-    cfg.bind = "0.0.0.0:52414"
-    # nb: shared.logging needs updating if any of this changes
-    cfg.accesslog = "-"
-    cfg.errorlog = "-"
-    ev = anyio.Event()
-    app = FastAPI()
-    app.post("/run_test")(run_test)
-    app.post("/kill")(lambda: kill(ev))
-    app.get("/tb_detection")(tb_detection)
-    app.get("/models")(list_models)
-    await serve(
-        app,  # type: ignore
-        cfg,
-        shutdown_trigger=lambda: ev.wait(),
-    )
-
-
-def kill(ev: anyio.Event):
-    ev.set()
-    return Response(status_code=204)
-
-
-async def tb_detection():
-    send, recv = channel[GatheredInfo]()
-    ig = InfoGatherer(send)
-    with anyio.move_on_after(1):
-        await ig._monitor_system_profiler_thunderbolt_data()  # pyright: ignore[reportPrivateUsage]
-    with recv:
-        return recv.collect()
-
-
-def list_models():
-    sent = set[str]()
-    for path in EXO_DEFAULT_MODELS_DIR.rglob("model-*.safetensors"):
-        if "--" not in path.parent.name:
-            continue
-        name = path.parent.name.replace("--", "/")
-        if name in sent:
-            continue
-        sent.add(name)
-        yield ModelId(path.parent.name.replace("--", "/"))
-
-
-async def run_test(test: Tests):
-    weird_hn = socket.gethostname()
-    for dev in test.devs:
-        if weird_hn.startswith(dev[0]) or dev[0].startswith(weird_hn):
-            hn = dev[0]
-            break
-    else:
-        raise ValueError(f"{weird_hn} not in {test.devs}")
-
-    async def run():
-        logger.info(f"testing {test.model_id}")
-
-        instances: list[Instance] = []
-        if test.kind in ["ring", "both"]:
-            i = await ring_instance(test, hn)
-            if i is None:
-                yield "no model found"
-                return
-            instances.append(i)
-        if test.kind in ["jaccl", "both"]:
-            i = await jaccl_instance(test)
-            if i is None:
-                yield "no model found"
-                return
-            instances.append(i)
-
-        for instance in instances:
-            recv = await execute_test(test, instance, hn)
-
-            str_out = ""
-
-            for item in recv:
-                if isinstance(item, ChunkGenerated):
-                    assert isinstance(item.chunk, TokenChunk)
-                    str_out += item.chunk.text
-
-                if isinstance(item, RunnerStatusUpdated) and isinstance(
-                    item.runner_status, (RunnerFailed, RunnerShutdown)
-                ):
-                    yield str_out + "\n"
-                    yield item.model_dump_json() + "\n"
-
-    return StreamingResponse(run())
-
-
-async def ring_instance(test: Tests, hn: str) -> Instance | None:
-    hbn = [Host(ip="198.51.100.0", port=52417) for _ in test.devs]
-    world_size = len(test.devs)
-    for i in range(world_size):
-        if test.devs[i][0] == hn:
-            hn = test.devs[i][0]
-        hbn[(i - 1) % world_size] = Host(ip=test.devs[i - 1][1], port=52417)
-        hbn[(i + 1) % world_size] = Host(ip=test.devs[i + 1][1], port=52417)
-        hbn[i] = Host(ip="0.0.0.0", port=52417)
-        break
-    else:
-        raise ValueError(f"{hn} not in {test.devs}")
-
-    card = await ModelCard.load(test.model_id)
-    instance = MlxRingInstance(
-        instance_id=iid,
-        ephemeral_port=52417,
-        hosts_by_node={NodeId(hn): hbn},
-        shard_assignments=ShardAssignments(
-            model_id=test.model_id,
-            node_to_runner={NodeId(host[0]): RunnerId(host[0]) for host in test.devs},
-            runner_to_shard={
-                RunnerId(test.devs[i][0]): PipelineShardMetadata(
-                    model_card=card,
-                    device_rank=i,
-                    world_size=world_size,
-                    start_layer=(card.n_layers // world_size) * i,
-                    end_layer=min(
-                        card.n_layers, (card.n_layers // world_size) * (i + 1)
-                    ),
-                    n_layers=min(card.n_layers, (card.n_layers // world_size) * (i + 1))
-                    - (card.n_layers // world_size) * i,
-                )
-                for i in range(world_size)
-            },
-        ),
-    )
-
-    return instance
-
-
-async def execute_test(test: Tests, instance: Instance, hn: str) -> list[Event]:
-    world_size = len(test.devs)
-    commands: list[Task] = [
-        (LoadModel(instance_id=iid)),
-        (StartWarmup(instance_id=iid)),
-        (
-            TextGeneration(
-                task_params=TextGenerationTaskParams(
-                    model=test.model_id,
-                    instructions="You are a helpful assistant",
-                    input=[
-                        InputMessage(
-                            role="user", content="What is the capital of France?"
-                        )
-                    ],
-                ),
-                command_id=CommandId("yo"),
-                instance_id=iid,
-            )
-        ),
-        (Shutdown(runner_id=RunnerId(hn), instance_id=iid)),
-    ]
-    if world_size > 1:
-        commands.insert(0, ConnectToGroup(instance_id=iid))
-    bound_instance = BoundInstance(
-        instance=instance, bound_runner_id=RunnerId(hn), bound_node_id=NodeId(hn)
-    )
-    ev_send, _ev_recv = mp_channel[Event]()
-    task_send, task_recv = mp_channel[Task]()
-
-    for command in commands:
-        task_send.send(command)
-
-    entrypoint(
-        bound_instance,
-        ev_send,
-        task_recv,
-        logger,
-    )
-
-    # TODO(evan): return ev_recv.collect()
-    return []
-
-
-async def jaccl_instance(test: Tests) -> MlxJacclInstance | None:
-    card = await ModelCard.load(test.model_id)
-    world_size = len(test.devs)
-    assert test.ibv_devs
-
-    return MlxJacclInstance(
-        instance_id=iid,
-        jaccl_devices=test.ibv_devs,
-        # rank 0 is always coordinator
-        jaccl_coordinators={
-            NodeId(host[0]): test.devs[0][1] + ":52417" for host in test.devs
-        },
-        shard_assignments=ShardAssignments(
-            model_id=test.model_id,
-            node_to_runner={NodeId(host[0]): RunnerId(host[0]) for host in test.devs},
-            runner_to_shard={
-                RunnerId(host[0]): TensorShardMetadata(
-                    model_card=card,
-                    device_rank=i,
-                    world_size=world_size,
-                    start_layer=0,
-                    end_layer=card.n_layers,
-                    n_layers=card.n_layers,
-                )
-                for i, host in enumerate(test.devs)
-            },
-        ),
-    )
-
-
-if __name__ == "__main__":
-    anyio.run(main)

tests/start_distributed_test.py

@@ -1,85 +0,0 @@
-#!/usr/bin/env python3
-import itertools
-import json
-import subprocess
-import sys
-from concurrent.futures import ThreadPoolExecutor
-from typing import Any, cast
-from urllib.request import Request, urlopen
-
-if not (args := sys.argv[1:]):
-    sys.exit(
-        f"USAGE: {sys.argv[0]} <kind> [host1] [host2] ...\nkind is optional, and should be jaccl or ring"
-    )
-
-kind = args[0] if args[0] in ("jaccl", "ring") else "both"
-hosts = args[1:] if kind != "both" else args
-ts = subprocess.run(
-    ["tailscale", "status"], check=True, text=True, capture_output=True
-).stdout.splitlines()
-ip = {sl[1]: sl[0] for line in ts if len(sl := line.split()) >= 2}
-ips = [ip[h] for h in hosts]
-devs = [[h, ip[h]] for h in hosts]
-n = len(hosts)
-
-
-def get_tb(a: str) -> list[dict[str, Any]]:
-    with urlopen(f"http://{a}:52414/tb_detection", timeout=5) as r:  # pyright: ignore[reportAny]
-        return json.loads(r.read())  # pyright: ignore[reportAny]
-
-
-def get_models(a: str) -> set[str]:
-    with urlopen(f"http://{a}:52414/models", timeout=5) as r:  # pyright: ignore[reportAny]
-        return set(json.loads(r.read()))  # pyright: ignore[reportAny]
-
-
-def run(h: str, a: str, body: bytes) -> None:
-    with urlopen(
-        Request(
-            f"http://{a}:52414/run_test",
-            data=body,
-            method="POST",
-            headers={"Content-Type": "application/json"},
-        ),
-        timeout=300,
-    ) as r:  # pyright: ignore[reportAny]
-        for line in r.read().decode(errors="replace").splitlines():  # pyright: ignore[reportAny]
-            print(f"\n{h}@{a}: {line}", flush=True)
-
-
-with ThreadPoolExecutor(n) as exctr:
-    if kind in ("jaccl", "both"):
-        payloads = list(exctr.map(get_tb, ips))
-
-        u2e = {
-            ident["domainUuid"]: (i, ident["rdmaInterface"])
-            for i, p in enumerate(payloads)
-            for d in p
-            for ident in cast(
-                list[dict[str, str]],
-                d.get("MacThunderboltIdentifiers", {}).get("idents", []),  # pyright: ignore[reportAny]
-            )
-        }
-        edges = {
-            (u2e[s][0], u2e[t][0]): u2e[t][1]
-            for p in payloads
-            for d in p
-            for c in d.get("MacThunderboltConnections", {}).get("conns", [])  # pyright: ignore[reportAny]
-            if (s := c["sourceUuid"]) in u2e and (t := c["sinkUuid"]) in u2e  # pyright: ignore[reportAny]
-        }
-        ibv_devs = [[edges.get((i, j)) for j in range(n)] for i in range(n)]
-    else:
-        ibv_devs = None
-
-    models = set[str].intersection(*exctr.map(get_models, ips))
-
-    print("\n")
-    print("=" * 70)
-    print(f"Starting test with {models}")
-    print("=" * 70)
-    print("\n")
-    for model in models:
-        body = json.dumps(
-            {"devs": devs, "model_id": model, "ibv_devs": ibv_devs, "kind": kind}
-        ).encode()
-        list(exctr.map(run, hosts, ips, itertools.repeat(body)))

tests/test_1node.py

@@ -0,0 +1,75 @@
+# type: ignore
+"""Single-node integration tests.
+
+Run with:
+    uv run pytest tests/test_1node.py -v
+"""
+
+from __future__ import annotations
+
+import time
+
+import pytest
+from exo_tools.harness import is_model_downloaded, place_instance
+
+from .framework import DEFAULT_MODEL, InstanceSpec
+
+
+@pytest.mark.cluster(count=1)
+@pytest.mark.instance(DEFAULT_MODEL)
+def test_place_instance_and_chat(session):
+    resp = session.chat("Say hello in one sentence.")
+    assert len(resp) > 0
+
+
+@pytest.mark.cluster(count=1)
+@pytest.mark.instance(DEFAULT_MODEL)
+def test_chat_multiple_turns(session):
+    first_reply = session.chat("What is 2 + 2?")
+    assert len(first_reply) > 0
+
+    second_reply = session.multi_turn(
+        [
+            {"role": "user", "content": "What is 2 + 2?"},
+            {"role": "assistant", "content": first_reply},
+            {"role": "user", "content": "Now multiply that by 3."},
+        ]
+    )
+    assert len(second_reply) > 0
+
+
+@pytest.mark.cluster(count=1)
+@pytest.mark.instance(DEFAULT_MODEL)
+def test_delete_instance(session):
+    from exo_tools.harness import wait_for_instance_gone
+
+    session.client.request_json("DELETE", f"/instance/{session.instance_id}")
+    wait_for_instance_gone(session.client, session.instance_id, timeout=30.0)
+    assert len(session.instances) == 0, (
+        f"Expected no instances, found {len(session.instances)}"
+    )
+
+
+@pytest.mark.cluster(count=1)
+def test_download_from_scratch(session):
+    """Ensure the model is not on the cluster, then place an instance to
+    trigger a fresh download and verify inference.
+    """
+    node_id = next(iter(session.state.get("nodeIdentities", {})))
+
+    # Delete any existing download — the API call is idempotent
+    session.client.request_json("DELETE", f"/download/{node_id}/{DEFAULT_MODEL}")
+
+    # Poll until the model is gone (it may already be gone)
+    deadline = time.time() + 60.0
+    while time.time() < deadline:
+        if not is_model_downloaded(session.client, DEFAULT_MODEL):
+            break
+        time.sleep(2.0)
+    else:
+        raise AssertionError(f"Expected {DEFAULT_MODEL} to be deleted from cluster")
+
+    place_instance(session.client, DEFAULT_MODEL, timeout=900.0)
+    session.instance_spec = InstanceSpec(model_id=DEFAULT_MODEL)
+    resp = session.chat("Say hello in one sentence.")
+    assert len(resp) > 0