docs: request the mlx extra in the documented setup commands (#2245)

Since #2087 moved mlx, mlx-lm, mlx-vlm and mflux out of `dependencies`
into the `mlx` extra, the documented setup path never installs them, and
`[tool.uv]` sets no default extras. `uv run exo` starts the API, then
every runner crashes with `ModuleNotFoundError: No module named 'mlx'`
(#2156).

Adds the extra to the commands documented as the way to set up:

- README macOS: `uv sync --extra mlx`
- README Linux: `uv sync --extra mlx-cpu`, with the mlx-cuda12 /
mlx-cuda13 alternatives
- CONTRIBUTING.md quick start
- `just sync` / `just sync-clean`, which `just build-app` runs before
pyinstaller — the spec aborts when the mlx Metal libraries are missing

Docs plus two justfile recipes; no dependency or lock changes. Backend
choice stays explicit, since mlx-cpu / mlx-cuda12 / mlx-cuda13 /
mlx-none are declared as conflicting extras.

#2234 overlaps on the macOS README block only, via a setup script that
runs `uv sync --extra mlx`; the Linux block, CONTRIBUTING.md and the
justfile are not covered there.

CONTRIBUTING.md

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

Cargo.toml

@@ -22,10 +22,10 @@ opt-level = 3
 networking = { path = "rust/networking" }
 
 # pyo3
-pyo3 = "0.27.2"
-pyo3-async-runtimes = "0.27.0"
-pyo3-log = "0.13.2"
-pyo3-stub-gen = "0.22.2"
+pyo3 = "0.28.3"
+pyo3-async-runtimes = "0.28.0"
+pyo3-log = "0.13.3"
+pyo3-stub-gen = "0.22.3"
 
 # util
 extend = "1.2"
@@ -42,7 +42,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"
@@ -54,6 +53,11 @@ zenoh-plugin-storage-manager = { version = "=1.9.0", default-features = false }
 zenoh-plugin-trait = "=1.9.0"
 netwatcher = "0.6.0"
 bytemuck = "1.25.0"
+blake3 = "1.8.5"
+smol = "2.0.2"
+socket2 = "0.6.4"
+tracing = "0.1.44"
+pidfile-rs = { git = "https://github.com/AndreiCravtov/pidfile-rs" }
 
 [patch.crates-io]
 zenoh = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }

README.md

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

TODO.md

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

dashboard/package-lock.json

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

dashboard/package.json

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

docs/architecture.md

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

flake.nix

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

justfile

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

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

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

rust/exo_rs/Cargo.toml

@@ -35,7 +35,7 @@ pyo3-async-runtimes = { workspace = true, features = [
 ] }
 pyo3-log.workspace = true
 
-pidfile-rs = { git = "https://github.com/AndreiCravtov/pidfile-rs" }
+pidfile-rs = { workspace = true }
 
 # async runtime
 tokio = { workspace = true, features = ["full"] }

rust/exo_rs/exo_rs.pyi

@@ -6,16 +6,36 @@ import os
 import pathlib
 import typing
 __all__ = [
+    "FromSwarm",
     "NetworkingHandle",
     "Pidfile",
     "PidfileError",
-    "PyFromSwarm",
 ]
 
+class FromSwarm:
+    @typing.final
+    class Connection(FromSwarm):
+        __match_args__ = ("connected",)
+        @property
+        def connected(self) -> builtins.bool: ...
+        def __new__(cls, connected: builtins.bool) -> FromSwarm.Connection: ...
+    
+    @typing.final
+    class Message(FromSwarm):
+        __match_args__ = ("topic", "data",)
+        @property
+        def topic(self) -> builtins.str: ...
+        @property
+        def data(self) -> bytes: ...
+        def __new__(cls, topic: builtins.str, data: bytes) -> FromSwarm.Message: ...
+    
+    ...
+
 @typing.final
 class NetworkingHandle:
     @staticmethod
-    def new(identity: bytes, listen_port: builtins.int, discovery_service_port: builtins.int) -> NetworkingHandle: ...
+    def new(identity: builtins.str, namespace: builtins.str, listen_port: builtins.int, discovery_service_port: builtins.int) -> NetworkingHandle: ...
+    def recv(self) -> typing.Awaitable[FromSwarm]: ...
     async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
         r"""
         Subscribe to a `GossipSub` topic.
@@ -34,7 +54,6 @@ class NetworkingHandle:
         
         If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
         """
-    async def recv(self) -> PyFromSwarm: ...
 
 @typing.final
 class Pidfile:
@@ -42,7 +61,7 @@ class Pidfile:
     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`][]
+    lock it, detect already running daemons. It is backed by [`pidfile`]
     functions of `libbsd`/`libutil` which use `flopen` to lock the PID
     file.
     
@@ -61,7 +80,6 @@ class Pidfile:
     def __new__(cls, path: builtins.str | os.PathLike | pathlib.Path, mode: builtins.int) -> Pidfile:
         r"""
         Creates a new PID file and locks it.
-        Writes the current process ID to the PID file.
         
         If the PID file cannot be locked, returns `PidfileError::AlreadyRunning` with
         a PID of the already running process, or `None` if no PID has been written to
@@ -93,22 +111,3 @@ 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: ...
-    
-    ...
-

rust/exo_rs/pyproject.toml

@@ -21,7 +21,13 @@ dev = ["exo_rs", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"]
 module-name = "exo_rs"
 features = ["pyo3/extension-module", "pyo3/experimental-async"]
 
+[tool.pyo3-stub-gen]
+generate-init-py = true
+
 [tool.pytest.ini_options]
 log_cli = true
 log_cli_level = "INFO"
 asyncio_mode = "auto"
+
+[tool.uv]
+cache-keys = [{ file = "src/**/*.rs" }]

rust/exo_rs/src/ident.rs

@@ -1,44 +0,0 @@
-use crate::ext::ResultExt as _;
-use pyo3::types::{PyBytes, PyBytesMethods as _};
-use pyo3::{Bound, PyResult, Python, pyclass, pymethods};
-use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
-
-/// Identity keypair of a node.
-#[gen_stub_pyclass]
-#[pyclass(name = "Keypair", frozen)]
-#[repr(transparent)]
-pub struct PyKeypair(pub u128);
-
-#[gen_stub_pymethods]
-#[pymethods]
-#[allow(clippy::needless_pass_by_value)]
-impl PyKeypair {
-    /// Generate a new Ed25519 keypair.
-    #[staticmethod]
-    fn generate() -> Self {
-        Self(rand::random())
-    }
-
-    /// Construct an Ed25519 keypair from secret key bytes
-    #[staticmethod]
-    fn from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
-        let bytes = Vec::from(bytes.as_bytes());
-        Ok(Self(u128::from_le_bytes(
-            bytes
-                .try_into()
-                .map_err(|_| "passed too many bytes to from_bytes")
-                .pyerr()?,
-        )))
-    }
-
-    /// Get the secret key bytes underlying the keypair
-    fn to_bytes<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
-        let bytes = self.0.to_le_bytes();
-        Ok(PyBytes::new(py, &bytes))
-    }
-
-    /// Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
-    fn to_node_id(&self) -> String {
-        format!("{:x}", self.0)
-    }
-}

rust/exo_rs/src/lib.rs

@@ -5,9 +5,9 @@
 //!
 
 mod allow_threading;
-mod pidfile;
 // mod ident;
 mod networking;
+mod pidfile;
 
 use crate::networking::networking_submodule;
 use crate::pidfile::pidfile_submodule;
@@ -146,7 +146,7 @@ pub(crate) mod ext {
 /// A Python module implemented in Rust. The name of this function must match
 /// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to
 /// import the module.
-#[pymodule(name = "exo_rs")]
+#[pymodule(name = "exo_rs", gil_used = true)]
 fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     // install logger
     pyo3_log::init();
@@ -159,7 +159,6 @@ fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     //       too many importing issues...
     pidfile_submodule(m)?;
     // m.add_class::<PyKeypair>()?;
-    // networking_submodule(m)?;
     networking_submodule(m)?;
 
     // top-level constructs

rust/exo_rs/src/networking.rs

@@ -4,27 +4,26 @@ use std::sync::Arc;
 use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
 use crate::ext::{ResultExt as _, TokioMpscSenderExt as _};
 use futures_lite::{Stream, StreamExt as _};
-use networking::swarm::{FromSwarm, ToSwarm, create_swarm};
+use networking::swarm::{FromSwarm, Swarm, ToSwarm, create_swarm};
+use networking::{Session, is_valid_zid};
 use pyo3::exceptions::{PyRuntimeError, PyValueError};
 use pyo3::prelude::*;
 use pyo3::types::PyBytes;
 use pyo3::{Bound, Py, PyAny, PyErr, PyResult, Python, pymethods};
-use pyo3_stub_gen::derive::{
-    gen_methods_from_python, gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods,
-};
+use pyo3_stub_gen::derive::{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 {
+pub 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 {
+#[pyclass(name = "FromSwarm")]
+pub enum PyFromSwarm {
     Connection { connected: bool },
     Message { topic: String, data: Py<PyBytes> },
 }
@@ -41,6 +40,20 @@ impl From<FromSwarm> for PyFromSwarm {
     }
 }
 
+impl PyNetworkingHandle {
+    pub fn from_session(session: Session) -> Self {
+        let (to_swarm, from_client) = mpsc::channel(1024);
+        let swarm = Swarm {
+            from_client,
+            session,
+        };
+        PyNetworkingHandle {
+            swarm: Arc::new(Mutex::new(swarm.into_stream())),
+            to_swarm,
+        }
+    }
+}
+
 #[gen_stub_pymethods]
 #[pymethods]
 impl PyNetworkingHandle {
@@ -51,30 +64,31 @@ impl PyNetworkingHandle {
     // ---- Lifecycle management methods ----
 
     #[staticmethod]
-    fn new<'py>(
-        identity: Bound<'py, PyBytes>,
+    pub fn new(
+        identity: &str,
+        namespace: &str,
         listen_port: u16,
         discovery_service_port: u16,
     ) -> PyResult<PyNetworkingHandle> {
         // todo: zenoh self assigned peers
         if listen_port == 0 {
-            todo!();
+            todo!("cannot listen on port 0 yet");
         }
         // 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"))?,
-        );
+        if !is_valid_zid(identity) {
+            return Err(PyValueError::new_err(format!(
+                "{identity} is not a valid zenoh identity"
+            )));
+        }
 
         // create networking swarm (within tokio context!! or it crashes)
         let swarm = pyo3_async_runtimes::tokio::get_runtime()
             .block_on(create_swarm(
                 identity,
+                namespace,
                 from_client,
                 listen_port,
                 discovery_service_port,
@@ -87,8 +101,10 @@ impl PyNetworkingHandle {
         })
     }
 
-    #[gen_stub(skip)]
-    fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+    #[gen_stub(override_return_type(
+        type_repr="typing.Awaitable[FromSwarm]", imports=("typing")
+    ))]
+    pub fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         let swarm = Arc::clone(&self.swarm);
         pyo3_async_runtimes::tokio::future_into_py(py, async move {
             swarm
@@ -106,7 +122,7 @@ impl PyNetworkingHandle {
     /// 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> {
+    pub async fn gossipsub_subscribe(&self, topic: String) -> PyResult<bool> {
         let (tx, rx) = oneshot::channel();
 
         // send off request to subscribe
@@ -128,7 +144,7 @@ impl PyNetworkingHandle {
     /// 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> {
+    pub async fn gossipsub_unsubscribe(&self, topic: String) -> PyResult<bool> {
         let (tx, rx) = oneshot::channel();
 
         // send off request to unsubscribe
@@ -149,7 +165,7 @@ impl PyNetworkingHandle {
     /// 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<()> {
+    pub async fn gossipsub_publish(&self, topic: String, data: Py<PyBytes>) -> PyResult<()> {
         let (tx, rx) = oneshot::channel();
 
         // send off request to subscribe
@@ -173,15 +189,6 @@ impl PyNetworkingHandle {
     }
 }
 
-pyo3_stub_gen::inventory::submit! {
-    gen_methods_from_python! {
-        r#"
-            class PyNetworkingHandle:
-                async def recv() -> PyFromSwarm: ...
-        "#
-    }
-}
-
 pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<PyNetworkingHandle>()?;
     m.add_class::<PyFromSwarm>()?;

rust/exo_rs/src/pidfile.rs

@@ -38,7 +38,7 @@ impl PyPidfileError {
 /// 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`][]
+/// lock it, detect already running daemons. It is backed by [`pidfile`]
 /// functions of `libbsd`/`libutil` which use `flopen` to lock the PID
 /// file.
 ///
@@ -77,7 +77,6 @@ impl PyPidfile {
 #[pymethods]
 impl PyPidfile {
     /// Creates a new PID file and locks it.
-    /// Writes the current process ID to the PID file.
     ///
     /// If the PID file cannot be locked, returns `PidfileError::AlreadyRunning` with
     /// a PID of the already running process, or `None` if no PID has been written to

rust/exo_rs/tests/test_python.py

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

rust/networking/Cargo.toml

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

rust/networking/examples/do_nothing.rs

@@ -6,8 +6,8 @@ use zenoh::{Result, Wait};
 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, 52413).await?;
+    let cfg = networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414)?;
+    let session = networking::open(cfg, "exo", 52414, 52413).await?;
     let _tok = session
         .z
         .liveliness()

rust/networking/examples/heartbeat.rs

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

rust/networking/examples/put_string.rs

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

rust/networking/examples/serve_storage.rs

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

rust/networking/examples/z_get.rs

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

rust/networking/src/discovery.rs

@@ -1,6 +1,4 @@
 use std::{
-    env,
-    hash::{DefaultHasher, Hash, Hasher},
     io,
     net::{Ipv6Addr, SocketAddr, SocketAddrV6},
     sync::Arc,
@@ -39,16 +37,24 @@ pub struct Discovered {
 }
 
 impl Discovery {
-    pub async fn new(zid: ZenohId, listen_port: u16, discovery_port: u16) -> io::Result<Self> {
-        let namespace = {
-            let mut hasher = DefaultHasher::new();
-            env::var("EXO_ZENOH_NAMESPACE")
-                .unwrap_or_else(|_| "exo".to_string())
-                .hash(&mut hasher);
-            hasher.finish().to_le_bytes()
-        };
-        let sock = Arc::new(UdpSocket::bind(format!("[::]:{discovery_port}")).await?);
-        //sock.set_multicast_loop_v6(false)?;
+    pub async fn new(
+        zid: ZenohId,
+        namespace: [u8; 8],
+        listen_port: u16,
+        discovery_port: u16,
+    ) -> io::Result<Self> {
+        let sock = socket2::Socket::new(
+            socket2::Domain::IPV6,
+            socket2::Type::DGRAM,
+            Some(socket2::Protocol::UDP),
+        )?;
+        sock.set_reuse_address(true)?;
+        #[cfg(unix)]
+        sock.set_reuse_port(true)?;
+        sock.bind(&SocketAddrV6::new(Ipv6Addr::UNSPECIFIED, discovery_port, 0, 0).into())?;
+        sock.set_nonblocking(true)?;
+        sock.set_multicast_loop_v6(true)?;
+        let sock = Arc::new(UdpSocket::from_std(sock.into())?);
         let ifaces: Arc<Mutex<Vec<SocketAddrV6>>> = Default::default();
         let _sync = Mutex::new(
             netwatcher::watch_interfaces_with_callback({
@@ -64,9 +70,12 @@ impl Discovery {
                         }
 
                         match sock.join_multicast_v6(&GROUP, *iface_idx) {
-                            Ok(()) => ifaces
-                                .lock()
-                                .push(SocketAddrV6::new(GROUP, 52413, 0, *iface_idx)),
+                            Ok(()) => ifaces.lock().push(SocketAddrV6::new(
+                                GROUP,
+                                discovery_port,
+                                0,
+                                *iface_idx,
+                            )),
                             Err(e) if e.kind() != io::ErrorKind::AddrInUse => {
                                 // skip AddrInUse - just means we've already joined the mv6
                                 if let Some(iface) = update.interfaces.get(&iface_idx) {

rust/networking/src/lib.rs

@@ -12,11 +12,21 @@ use crate::discovery::Discovery;
 pub mod discovery;
 pub mod swarm;
 
-pub fn cfg(identity: u128, listen_port: u16) -> Result<zenoh::Config> {
-    assert!(listen_port != 0, "must used defined listen port port");
+pub fn is_valid_zid(identity: &str) -> bool {
+    let mut iter = identity.chars();
+    iter.next()
+        .is_some_and(|c| ('1'..='9').contains(&c) || ('a'..='f').contains(&c))
+        && iter.all(|c| ('0'..='9').contains(&c) || ('a'..='f').contains(&c))
+        && identity.len() <= 32
+}
+
+pub fn cfg(identity: &str, listen_port: u16) -> Result<zenoh::Config> {
+    assert!(is_valid_zid(identity));
+    assert!(identity.len() <= 32);
+    assert!(listen_port != 0, "must used defined listen port");
     let mut cfg = zenoh::Config::default();
     // todo: cleanup
-    cfg.insert_json5("id", &format!("\"{identity:x}\""))?;
+    cfg.insert_json5("id", &format!("\"{identity}\""))?;
     cfg.insert_json5("mode", "\"router\"")?;
     cfg.insert_json5("listen/endpoints", &format!("[\"tcp/[::]:{listen_port}\"]"))?;
     cfg.insert_json5("scouting/multicast/enabled", "false")?;
@@ -25,7 +35,7 @@ pub fn cfg(identity: u128, listen_port: u16) -> Result<zenoh::Config> {
     cfg.insert_json5("adminspace/enabled", "true")?;
     //cfg.insert_json5("transport/link/tx/batch_size", "9216")?;
     cfg.insert_json5("transport/link/rx/buffer_size", "16777216")?;
-    cfg.insert_json5("timestamping/enabled", "true")?;
+    //cfg.insert_json5("timestamping/enabled", "true")?;
     cfg.insert_json5("plugins/storage_manager/__required__", "true")?;
     cfg.insert_json5(
         "plugins/storage_manager/storages/mem1",
@@ -43,10 +53,16 @@ pub fn cfg(identity: u128, listen_port: u16) -> Result<zenoh::Config> {
 
 pub async fn open(
     cfg: zenoh::Config,
+    namespace: &str,
     listen_port: u16,
     discovery_service_port: u16,
 ) -> Result<Session> {
     assert!(listen_port != 0, "must used defined listen port");
+    let namespace: [u8; 8] = {
+        blake3::hash(namespace.as_bytes()).as_bytes()[..8]
+            .try_into()
+            .expect("8 is equal to 8")
+    };
     let mut plugins = PluginsManager::static_plugins_only();
     plugins.declare_static_plugin::<StoragesPlugin, _>("storage_manager", true);
     let mut runtime = zenoh::internal::runtime::RuntimeBuilder::new(cfg)
@@ -55,8 +71,9 @@ pub async fn open(
         .await?;
     let z = zenoh::session::init(runtime.clone().into()).await?;
     runtime.start().await?;
-    let mut discovery = Discovery::new(z.zid(), listen_port, discovery_service_port).await?;
-    let _jh = Arc::new(tokio::task::spawn(async move {
+    let mut discovery =
+        Discovery::new(z.zid(), namespace, listen_port, discovery_service_port).await?;
+    let _jh = Arc::new(AbortOnDrop(tokio::task::spawn(async move {
         loop {
             let Ok(discovered) = discovery.next().await.inspect_err(|e| {
                 log::warn!("discovery error {e}");
@@ -81,17 +98,19 @@ pub async fn open(
                 .connect_peer(&discovered.zid.into(), &[locator])
                 .await;
         }
-    }));
+    })));
     Ok(Session { z, _jh })
 }
 
+struct AbortOnDrop(JoinHandle<()>);
+impl Drop for AbortOnDrop {
+    fn drop(&mut self) {
+        self.0.abort();
+    }
+}
+
 #[derive(Clone)]
 pub struct Session {
     pub z: ZSession,
-    _jh: Arc<JoinHandle<()>>,
-}
-impl Drop for Session {
-    fn drop(&mut self) {
-        self._jh.abort();
-    }
+    _jh: Arc<AbortOnDrop>,
 }

rust/networking/src/swarm.rs

@@ -70,14 +70,14 @@ impl Swarm {
                     token = discovery.recv_async() => {
                         if let Ok(token) = token {
                             let key_expr = token.key_expr().as_str().to_owned();
-                            let nid = key_expr.strip_prefix("nodes/").and_then(|s| s.strip_suffix("/live"));
+                            let zid = key_expr.strip_prefix("live/");
                             yield match token.kind() {
                                 SampleKind::Put => {
-                                    log::info!("discovered: {nid:?}");
+                                    log::info!("discovered: {zid:?}");
                                     FromSwarm::Discovered {}
                                 }
                                 SampleKind::Delete => {
-                                    log::info!("expired: {nid:?}");
+                                    log::info!("expired: {zid:?}");
                                     FromSwarm::Expired {}
                                 }
                             }
@@ -96,11 +96,11 @@ async fn register_liveness(
 ) -> Result<(LivelinessToken, Subscriber<FifoChannelHandler<Sample>>)> {
     let token = session
         .liveliness()
-        .declare_token(format!("nodes/{}/live", session.zid()))
+        .declare_token(format!("live/{}", session.zid()))
         .await?;
     let sub = session
         .liveliness()
-        .declare_subscriber("nodes/*/live")
+        .declare_subscriber("live/*")
         .history(true)
         .await?;
     Ok((token, sub))
@@ -193,13 +193,14 @@ async fn on_message(
 }
 
 pub async fn create_swarm(
-    identity: u128,
+    identity: &str,
+    namespace: &str,
     from_client: mpsc::Receiver<ToSwarm>,
     listen_port: u16,
     discovery_service_port: u16,
 ) -> Result<Swarm> {
     let cfg = crate::cfg(identity, listen_port)?;
-    let session = crate::open(cfg, listen_port, discovery_service_port).await?;
+    let session = crate::open(cfg, namespace, listen_port, discovery_service_port).await?;
     Ok(Swarm {
         session,
         from_client,

src/exo/__init__.py

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

src/exo/main.py

@@ -15,12 +15,13 @@ from loguru import logger
 from pydantic import PositiveInt
 
 import exo.routing.topics as topics
+from exo import __version__
 from exo.api.main import API
 from exo.download.coordinator import DownloadCoordinator
 from exo.download.impl_shard_downloader import exo_shard_downloader
 from exo.master.main import Master
 from exo.routing.event_router import EventRouter
-from exo.routing.router import Router
+from exo.routing.router import Router, get_node_zid
 from exo.shared.constants import EXO_DEFAULT_MODELS_DIR, EXO_LOG, EXO_PID_FILE
 from exo.shared.election import Election, ElectionResult
 from exo.shared.logging import logger_cleanup, logger_setup
@@ -50,12 +51,13 @@ class Node:
 
     @classmethod
     async def create(cls, args: "Args") -> Self:
-        keypair = os.urandom(16)
-        node_id = NodeId(keypair.hex())
+        node_id = get_node_zid()
         session_id = SessionId(master_node_id=node_id, election_clock=0)
         router = Router.create(
-            keypair,
+            node_id,
+            namespace=args.namespace,
             listen_port=args.zenoh_port,
+            discovery_service_port=args.discovery_port,
         )
         await router.register_topic(topics.GLOBAL_EVENTS)
         await router.register_topic(topics.LOCAL_EVENTS)
@@ -388,7 +390,9 @@ class Args(FrozenModel):
     fast_synch: bool | None = None  # None = auto, True = force on, False = force off
     legacy_daemon: bool = False
     bootstrap_peers: list[str] = []
+    namespace: str
     zenoh_port: int
+    discovery_port: int
 
     @classmethod
     def parse(cls) -> Self:
@@ -460,12 +464,26 @@ class Args(FrozenModel):
             dest="bootstrap_peers",
             help="Comma-separated libp2p multiaddrs to dial on startup (env: EXO_BOOTSTRAP_PEERS)",
         )
+        parser.add_argument(
+            "--namespace",
+            type=str,
+            default=__version__,
+            dest="namespace",
+            help="Discovery namespace, nodes with different namespaces will not connect.",
+        )
         parser.add_argument(
             "--zenoh-port",
             type=int,
-            default=0,
+            default=52414,
             dest="zenoh_port",
-            help="Fixed TCP port for zenoh to listen on (0 = OS-assigned).",
+            help="Fixed TCP port for zenoh to listen.",
+        )
+        parser.add_argument(
+            "--discovery-port",
+            type=int,
+            default=52413,
+            dest="discovery_port",
+            help="Fixed UDP port for the discovery service.",
         )
         fast_synch_group = parser.add_mutually_exclusive_group()
         fast_synch_group.add_argument(

src/exo/routing/connection_message.py

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

src/exo/routing/router.py

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

tmp/run_exo_on.sh

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