Compare commits

..
13 Commits
Author SHA1 Message Date
mlpy0 21a54c5ea0 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.
2026-08-25 18:59:53 +00:00
b5375f8cee Add Kimi K2.7-Code model card (official INT4 weights + vision) (#2167)
Adds a model card for
[moonshotai/Kimi-K2.7-Code](https://huggingface.co/moonshotai/Kimi-K2.7-Code),
released 2026-06-12.

Same architecture as Kimi K2.6 (`kimi_k25`, 61 layers, official INT4),
so the card mirrors the existing `moonshotai--Kimi-K2.6.toml`. Sampling
defaults per the model card (temperature 1.0 / top_p 0.95 for thinking
mode).

**Vision:** the official repo ships MoonViT weights inline, so I
extracted the 335 `vision_tower.*` / `mm_projector.*` tensors
(unmodified bf16) into
[aidiffuser/Kimi-K2.7-Code-vision](https://huggingface.co/aidiffuser/Kimi-K2.7-Code-vision),
following the `exolabs/Kimi-K2.6-vision` format. The vision config is
byte-identical to K2.6's; the extraction script is included in the repo
for verification. Happy to have this re-hosted under the exolabs org if
you prefer — it's a one-line change to the card.

**Tested:** distributed serving on 2× Mac Studio M3 Ultra (512 GB),
tensor parallelism, text + thinking + image understanding all confirmed
working.

Co-authored-by: aidiffuser <your-noreply-email@users.noreply.github.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
2026-06-22 13:38:54 +00:00
OrbisAI Security cdf1add867 fix: upgrade devalue to 5.6.2 (CVE-2026-22774) (#2150)
## Summary
Upgrade devalue from 5.5.0 to 5.6.2 to fix CVE-2026-22774.

## Vulnerability
| Field | Value |
|-------|-------|
| **ID** | CVE-2026-22774 |
| **Severity** | HIGH |
| **Scanner** | trivy |
| **Rule** | `CVE-2026-22774` |
| **File** | `dashboard/package-lock.json` |
| **Assessment** | Likely exploitable |

**Description**: devalue: devalue: Denial of Service due to excessive
resource consumption from untrusted input

## Evidence

**Scanner confirmation**: trivy rule `CVE-2026-22774` flagged this
pattern.

**Production code**: This file is in the production codebase, not
test-only code.

## Threat Model Context

This is a web service - vulnerabilities in request handlers are directly
exploitable by remote attackers.

## Changes
- `dashboard/package.json`
- `dashboard/package-lock.json`

## Verification
- [x] Build passes
- [x] Scanner re-scan confirms fix
- [x] LLM code review passed

---
*This change addresses a pattern flagged by static analysis. The code
path handles user-influenced input and the fix reduces the attack
surface against both manual and automated exploitation.*

---
*Automated security fix by [OrbisAI Security](https://orbisappsec.com)*
2026-06-22 13:29:15 +00:00
Evan QuineyandAndrei Cravtov 09f9ea313f libp2p -> zenoh (#2132)
supercedes #2076 and #2073

---------

Co-authored-by: Andrei Cravtov <the.andrei.cravtov@gmail.com>
2026-06-03 16:31:56 +01:00
Sakutaro 81d7cb0fcd docs: add Homebrew cask install instructions (#2140)
## Motivation

exo is now available as a Homebrew cask, so the README should show the
simplest macOS installation path alongside the existing DMG download.

Fixes https://github.com/exo-explore/exo/issues/2105
https://github.com/exo-explore/exo/issues/176

## Changes

- Added `brew install --cask exo` to the macOS App section of
`README.md`
- Kept the existing DMG download link as the first installation option

## Why It Works

Adding the Homebrew cask command gives macOS users a
package-manager-managed installation path while preserving the existing
DMG download option.

## Test Plan

### Manual Testing

- Reviewed the rendered Markdown structure in `README.md`

### Automated Testing

- Not run. Documentation-only change.

## Related

- https://github.com/Homebrew/homebrew-cask/pull/265956
2026-06-02 15:35:18 +00:00
Andrei Cravtov 629c55d6ba Rename exo_pyo3_bindings to exo_rs (#2131)
## Motivation

(I think it) Makes Evan's massive PR easier to merge later on

## Changes

- Renamed exo_pyo3_bindings to exo_rs
- Upgraded versions of pyo3-based dependencies
- Renamed PyFromSwarm to just FromSwarm, and PyNetworkingHandle to just
NetworkingHandle
2026-05-31 19:23:41 +01:00
Andrei Cravtov f9f8cbb3c3 fix: make app builds work again (#2127)
## Motivation

They didn't

## Changes

They now do

## Why It Works

I changed an env flag, and added a keyword
2026-05-29 18:37:47 +01:00
ciaranbor 051a64e3b4 Capture energy in prefill and ageneration separately (#2124)
## Motivation

Energy was reported as a single aggregate. Split into prefill vs.
generation so each phase can be analysed independently.

## Changes

- `PowerSampler`: `mark_prefill_done()` + `trapezoidal_energy_range()`
helper; `result()` now emits per-phase splits.
- `PowerUsage` / `NodePowerStats`: optional `prefill_*` / `generation_*`
fields (back-compat: `None` if unmarked).
- API marks the boundary on the first non-`PrefillProgressChunk`.
- `bench/exo_bench.py` surfaces the split in the log line and persists
`power_usage` to JSON.
- METHODOLOGY: one sentence + one bullet.

## Why It Works

First non-prefill chunk *is* the boundary. Anchoring a sample there and
interpolating power at the boundary makes phase energies sum exactly to
the unsplit total.

## Test Plan

### Manual Testing

`eco`-reserved nodes:
- M3 Ultra, Qwen3-VL-4B, pp=8192/tg=1024: server 1940 J vs client 1931 J
(+0.5 %)
- M4 Pro, Qwen3.6-27B, pp=16384/tg=2048: server 20,292 J vs client
20,221 J (+0.35 %)

### Automated Testing

5 new tests in `test_power_sampler.py` (range integrator,
splits-sum-to-total, `None`-when-unmarked, idempotency). 14/14 pass.
2026-05-28 14:42:36 -07:00
Andrei Cravtov a8602ea6d5 fix(bug): no longer repeated _trigger_notify_user_to_download_model (#2114)
## Motivation

Partially fixes [this](https://github.com/exo-explore/exo/issues/2098)
issue. Removed erroneous logic for telling user to download when they
already downloaded.

Could not figure out about the "spontaneous crashes" in that issue,
author should consolidate more logs and open a new issue dedicated to
that. I believe
[this](https://github.com/exo-explore/exo/commit/74e9fe15e62fe189dc7e019db86e75c83eca2721)
commit solved some EventRouter-related crashes, which was mentioned in
[this](https://github.com/exo-explore/exo/issues/2098) issue, so it may
have already been solved. If not, should be re-submitted as a new issue.

## Changes

- Consolidated _resolve_and_validate_text_model and
_validate_image_model into one function: _validate_model_has_instance;
- + They already had virtually identical logic, it being different seems
to be an artifact of history
- + Added logic to ensure that _trigger_notify_user_to_download_model is
only called when no such model is downloaded, not just if there is no
instance of it
- Added a new `/instance/await` SSE streaming endpoint to wait for when
a model has an instance available. Complements instance-placement API,
so we can wait till that is done without client-side polling.
- Updated docs and a /tmp script to reflect some of the changes
- Updated dashboard `getModelForRequest` to only return model ID if an
instance exists for it, and updated bits to use `handleChatSend` instead
of `sendMessage` because that checks for if a model instance exists
first.

## Why It Works

The problem was that there was erroneous logging for model not
downloaded. I fixed that logic. The rest is extra.
2026-05-26 14:42:39 +01:00
Andrei Cravtov a1a22b5f38 feat: added background/daemon support (#2106)
## Motivation

Addresses [this](https://github.com/exo-explore/exo/issues/1931) issue.

## Changes

You can now launch Exo as a legacy SysV-style daemin (in the background)
with `--legacy-daemon` flag.
NOTE: don't use it if you're managing Exo with systemd or launchd

SIDE FIX: the macmon process not found trace is no longer displayed on
process shutdown via ctrl+c, that error is supressed.

## Why It Works

Because I used a daemonization library and tweaked it not to break
multiprocessing.

## Test Plan

I ran it in daemon mode, non daemon mode, etc., and pid locking +
inference + everything else works just fine.

Also ran it `ssh user@host -t 'cd exo && nohup nix run .#exo --
--legacy-daemon'` on a 4-node TB mac-mini cluster and the mDNS didn't
die
2026-05-25 20:42:47 +01:00
Andrei Cravtov 74e9fe15e6 fix(bug): EventRouter lifetime-handling fixed, no more process crashes (#2102)
## Motivation

Trying to (partially) fix
[this](https://github.com/exo-explore/exo/issues/2101) issue.

## Changes

Changed channels (in channels.py) to support exception overriding.

Made EventRouter channels throw a subclass of the resource closed/broken
errors.

The current lifetime logic of EventRouter in event loop no longer blows
up because components that use channels from EventRouter now catch the
subclass exceptions in the run method: Worker, Master,
DownloadCoordinator, RunnerSupervisor.

Added logic to throw when API server exits without being asked to shut
down - this kill the sleep-forever in the task-group.
2026-05-22 14:20:04 +01:00
Evan Quiney 90f24bef30 fix model cards not validating properly after #2071 (#2096) 2026-05-15 15:17:35 +00:00
Andrei Cravtov 5097b2665d Tweaked workspace settings (#2095)
workspace settings
2026-05-15 13:04:50 +00:00
96 changed files with 5124 additions and 5069 deletions

No files matched your search

View File
Whitespace-only changes.
-12
View File
@@ -1,13 +1 @@
# installs nix-direnv if doesn't exist (speeds up evaluation)
if ! has nix_direnv_version || ! nix_direnv_version 3.1.1; then
source_url "https://raw.githubusercontent.com/nix-community/nix-direnv/3.1.1/direnvrc" "sha256-p+fzQdrms/hDa7g+soShAybJNo4bN4SIAeSfqNKgD5I="
fi
use flake
# creates .venv if doesn't exist and loads its environment
export VIRTUAL_ENV=".venv"
if ! [ -d "./$VIRTUAL_ENV" ]; then
uv venv
fi
layout python
+2
View File
@@ -38,6 +38,8 @@ bench/**/*.json
# tmp
tmp/models
/build/exo
/.agents
/.claude/skills
/.claude
/.codex
skills-lock.json
+5 -5
View File
@@ -4,7 +4,7 @@ This file provides guidance to AI coding agents when working with code in this r
## Project Overview
exo is a distributed AI inference system that connects multiple devices into a cluster. It enables running large language models across multiple machines using MLX as the inference backend and libp2p for peer-to-peer networking.
exo is a distributed AI inference system that connects multiple devices into a cluster. It enables running large language models across multiple machines using MLX as the inference backend and zenoh for peer-to-peer networking.
## Build & Run Commands
@@ -69,7 +69,7 @@ If `nix fmt` changes any files, stage them before committing. The CI runs `nix f
### Node Composition
A single exo `Node` (src/exo/main.py) runs multiple components:
- **Router**: libp2p-based pub/sub messaging via Rust bindings (exo_pyo3_bindings)
- **Router**: zenoh-based pub/sub messaging via Rust bindings (exo_rs)
- **Worker**: Handles inference tasks, downloads models, manages runner processes
- **Master**: Coordinates cluster state, places model instances across nodes
- **Election**: Bully algorithm for master election
@@ -81,7 +81,7 @@ Components communicate via typed pub/sub topics (src/exo/routing/topics.py):
- `LOCAL_EVENTS`: Workers send events to master for indexing
- `COMMANDS`: Workers/API send commands to master
- `ELECTION_MESSAGES`: Election protocol messages
- `CONNECTION_MESSAGES`: libp2p connection updates
- `CONNECTION_MESSAGES`: zenoh connection updates
### Event Sourcing
The system uses event sourcing for state management:
@@ -98,8 +98,8 @@ The system uses event sourcing for state management:
### Rust Components
Rust code in `rust/` provides:
- `networking`: libp2p networking (gossipsub, peer discovery)
- `exo_pyo3_bindings`: PyO3 bindings exposing Rust to Python
- `networking`: zenoh networking (gossipsub, peer discovery)
- `exo_rs`: PyO3 bindings exposing Rust to Python
- `system_custodian`: System-level operations
### Dashboard
+1
View File
@@ -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
```
Generated
+2697 -2400
View File
File diff suppressed because it is too large. Load diff
+54 -11
View File
@@ -1,6 +1,6 @@
[workspace]
resolver = "3"
members = ["rust/networking", "rust/exo_pyo3_bindings", "rust/util"]
members = ["rust/exo_rs", "rust/networking"]
[workspace.package]
version = "0.0.1"
@@ -20,30 +20,73 @@ opt-level = 3
[workspace.dependencies]
## Crate members as common dependencies
networking = { path = "rust/networking" }
util = { path = "rust/util" }
# Macro dependecies
# pyo3
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"
delegate = "0.13"
# Utility dependencies
keccak-const = "0.2"
nix = "0.31"
# 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"
# 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"
blake3 = "1.8.5"
smol = "2.0.2"
socket2 = "0.6.4"
tracing = "0.1.44"
pidfile-rs = { git = "https://github.com/AndreiCravtov/pidfile-rs" }
[patch.crates-io]
zenoh = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-buffers = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-codec = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-collections = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-config = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-core = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-crypto = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-keyexpr = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-link = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-link-commons = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-link-quic = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-link-quic_datagram = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-link-tcp = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-link-tls = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-link-udp = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-link-unixsock_stream = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-link-ws = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-macros = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-plugin-trait = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-protocol = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-result = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-runtime = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-sync = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-task = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-transport = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh-util = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
zenoh_backend_traits = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
[workspace.lints.rust]
static_mut_refs = "warn" # Or use "warn" instead of deny
-41
View File
@@ -1,41 +0,0 @@
# Missed things
[X] Log EXO_LIBP2P_NAMESPACE on start in exo/main.py
[X] Ordering of warmup was changed, which is wrong. It was changed to rank < n-1, then rank=n-1. It should be rank!=0 then rank=0 (this matches the auto_parallel implementation. NOTE: we use a different convention to mlx-lm, our terminal rank is rank=n-1 whereas mlx-lm is rank=0 hence i can see why this was changed wrongly).
[X] Downloads keying by model_id not shard_metadata (worker/plan.py, worker/main.py).
[X] Fetching download status of all models on start
[X] Deduplication of tasks in plan_step.
[X] resolve_allow_patterns should just be wildcard now.
[X] no mx_barrier in genreate.py mlx_generate at the end.
[] cache assertion not needed in auto_parallel.py PipelineLastLayer.
[X] GPTOSS support dropped in auto_parallel.py.
[X] sharding changed "all-to-sharded" became _all_to_sharded in auto_parallel.py.
[X] same as above with "sharded-to-all" became _sharded_to_all in auto_parallel.py.
[X] Dropped support for Ministral3Model, DeepseekV32Model, Glm4MoeModel, Qwen3NextModel, GptOssMode in auto_parallel.py.
[] Dropped prefill/decode code in auto_parallel.py and utils_mlx.py.
[X] KV_CACHE_BITS should be None to disable quantized KV cache.
[X] Dropped _set_nofile_limit in utils_mlx.py.
[X] We have group optional in load_mlx_items in utils_mlx.py.
[X] Dropped add_missing_chat_templates for GptOss in load_mlx_items in utils_mlx.py.
[X] Dropped model.make_cache in make_kv_cache in utils_mlx.py.
[X] We put cache limit back in utils_mlx.py.
[X] topology.py remove_node removes the connections after checking if node is is in self._node_id_to_rx_id_map. on beta_1 it checks after, so would remove stale connections I guess?
[X] Missing Glm 4.7 model cards (this isn't ready yet but should be picked up, probably create an issue... the blocker is transforemrs version doesn't support the tokenizer for Glm 4.7. rc-1 does but we can't upgrade as it breaks other things.)
[] try-except in _command_processor only excepts ValueError. This was silently failing leading to un-debuggable errors (we had a KeyError that was happening ). Changed this to catch Exception instead of ValueError. See exo-v2 89ae38405e0052e3c22405daf094b065878aa873 and fb99fea69b5a39017efc90c5dad0072e677455f0.
[X] In placement.py, place_instance no longer looks at model_meta.supports_tensor and check if this tensor parallel number of nodes is supported by the model's tensor dimensions.
[X] In placement.py, place_instanec, we no longer have the special case to exclude DeepSeek v3.1 pipeline parallel (it doesn't work).
[] logger.warning("You have likely selected ibv for a single node instance; falling back to MlxRing") was changed to debug. That will spam this warning since it happens every time we query instance previews.
[X] In placement_utils.py, get_mlx_jaccl_coordinators, We no longer prioritise Jaccl Coordinator IP. Now it picks the first one, which is unstable (Jaccl coordinator over TB5 is unstable).
[X] Downloads keying by model_id not shard_metadata (worker/plan.py, worker/main.py).
[X] Fetching download status of all models on start
[X] Deduplication of tasks in plan_step.
[X] resolve_allow_patterns should just be wildcard now.
[X] KV_CACHE_BITS should be None to disable quantized KV cache.
[X] We put cache limit back in utils_mlx.py.
[X] In placement.py, place_instance no longer looks at model_meta.supports_tensor and check if this tensor parallel number of nodes is supported by the model's tensor dimensions.
[X] In placement.py, place_instanec, we no longer have the special case to exclude DeepSeek v3.1 pipeline parallel (it doesn't work).
[X] In placement_utils.py, get_mlx_jaccl_coordinators, We no longer prioritise Jaccl Coordinator IP. Now it picks the first one, which is unstable (Jaccl coordinator over TB5 is unstable).
+33 -2
View File
@@ -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
```
@@ -201,6 +208,12 @@ This starts the exo dashboard and API at http://localhost:52415/
uv run exo --no-worker
```
- `--legacy-daemon`: Run exo as a legacy SysV-style background daemon using double-fork daemonization. This is intended for legacy init scripts; systemd and launchd should run exo in the foreground without this flag.
```bash
uv run exo --legacy-daemon
```
**File Locations (Linux):**
exo follows the [XDG Base Directory Specification](https://specifications.freedesktop.org/basedir-spec/basedir-spec-latest.html) on Linux:
@@ -223,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:**
@@ -395,6 +414,18 @@ Sample response:
}
```
This command is asynchronous. Before sending inference requests, wait until the
API sees the new instance for this model:
```bash
curl -N "http://localhost:52415/instance/await?model_id=mlx-community/Llama-3.2-1B-Instruct-4bit"
```
The endpoint returns an SSE stream. A successful wait emits a message with
`"type": "ready"` and the matching instance; a timeout emits `"type": "timeout"`.
By default it waits indefinitely. Set `timeout_seconds` to a positive value to
bound the wait.
---
**3. Send a chat completion**
+2 -3
View File
@@ -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
+1 -1
View File
@@ -352,7 +352,7 @@ final class ExoProcessController: ObservableObject {
private func makeEnvironment(for runtimeURL: URL) -> [String: String] {
var environment = ProcessInfo.processInfo.environment
environment["EXO_RUNTIME_DIR"] = runtimeURL.path
environment["EXO_LIBP2P_NAMESPACE"] = computeNamespace()
environment["EXO_ZENOH_NAMESPACE"] = computeNamespace()
if !hfToken.isEmpty {
environment["HF_TOKEN"] = hfToken
}
@@ -37,7 +37,7 @@ final class ClusterStateService: ObservableObject {
/// gain nothing from being cached on disk. Use an ephemeral session
/// with `urlCache = nil` so neither response bodies nor metadata
/// touch disk.
private static func makeNonCachingSession() -> URLSession {
nonisolated private static func makeNonCachingSession() -> URLSession {
let config = URLSessionConfiguration.ephemeral
config.urlCache = nil
config.requestCachePolicy = .reloadIgnoringLocalCacheData
+2 -1
View File
@@ -125,7 +125,7 @@ A background thread polls each node at 1 Hz, collecting:
- System power draw (W)
- CPU cluster usage (performance and efficiency cores)
**Energy** is computed via trapezoidal integration of the power samples over each inference window (the wall-clock span of each benchmark request or concurrent batch). Average power is `total_joules / total_inference_seconds`.
**Energy** is computed via trapezoidal integration of the power samples over each inference window (the wall-clock span of each benchmark request or concurrent batch). Average power is `total_joules / total_inference_seconds`. The server additionally returns a `power_usage` block in each non-stream `/bench/chat/completions` response that splits energy into prefill and generation phases, with the boundary anchored to the first non-`PrefillProgressChunk` from the runner.
---
@@ -136,6 +136,7 @@ Results are written as JSON with three top-level keys:
- **`runs`**: Array of per-request result objects, each containing:
- `elapsed_s`, `output_text_preview` (first 200 chars)
- `stats`: `{ prompt_tps, generation_tps, prompt_tokens, generation_tokens, peak_memory_usage }`
- `power_usage`: server-side total + prefill/generation split, per-node breakdown (non-stream requests only)
- Placement metadata: `model_id`, `placement_sharding`, `placement_instance_meta`, `placement_nodes`
- Run metadata: `pp_tokens`, `tg`, `repeat_index`, `concurrency`, `concurrent_index`
- `download_duration_s` (if model was freshly downloaded)
+26
View File
@@ -295,6 +295,7 @@ def run_one_completion(
elapsed = time.perf_counter() - t0
stats = out.get("generation_stats")
power_usage = out.get("power_usage")
choices = out.get("choices") or [{}]
message = choices[0].get("message", {}) if choices else {}
content = message.get("content") or ""
@@ -330,6 +331,7 @@ def run_one_completion(
elapsed = time.perf_counter() - t0
preview = "".join(text_parts)[:200]
power_usage = None
if not stats:
ttft = (first_token_time - t0) if first_token_time else elapsed
@@ -348,6 +350,7 @@ def run_one_completion(
"elapsed_s": elapsed,
"output_text_preview": preview,
"stats": stats,
"power_usage": power_usage,
}, pp_tokens
@@ -764,6 +767,7 @@ def main() -> int:
out = c.post_bench_chat_completions(_payload)
elapsed = time.perf_counter() - t0
stats = out.get("generation_stats")
power_usage = out.get("power_usage")
choices = out.get("choices") or [{}]
message = (
choices[0].get("message", {}) if choices else {}
@@ -773,6 +777,7 @@ def main() -> int:
"elapsed_s": elapsed,
"output_text_preview": text[:200],
"stats": stats,
"power_usage": power_usage,
}, _actual_pp
inf_t0 = time.monotonic()
@@ -868,6 +873,27 @@ def main() -> int:
inf_seconds = sum(t1 - t0 for t0, t1 in inference_windows)
avg_watts = joules / inf_seconds if inf_seconds > 0 else 0
summary += f" energy={joules:.1f}J ({avg_watts:.1f}W avg over {inf_seconds:.1f}s inference)"
# mean() not sum() across concurrent runs: each
# request's PowerSampler observes the same shared
# cluster state, so they all report the same figure.
prefill_energies = [
(x.get("power_usage") or {}).get("prefill_energy_joules")
for x in runs
]
gen_energies = [
(x.get("power_usage") or {}).get("generation_energy_joules")
for x in runs
]
prefill_vals = [e for e in prefill_energies if e is not None]
gen_vals = [e for e in gen_energies if e is not None]
if prefill_vals and gen_vals:
avg_pref = mean(prefill_vals)
avg_gen = mean(gen_vals)
summary += (
f" prefill_energy={avg_pref:.1f}J "
f"gen_energy={avg_gen:.1f}J"
)
logger.info(f"{summary}\n")
time.sleep(2)
finally:
+4 -3
View File
@@ -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": {
+3 -2
View File
@@ -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",
+11 -5
View File
@@ -2253,10 +2253,9 @@ class AppStore {
* @returns The model ID to use, or null if none available
*/
private getModelForRequest(modelId?: string): string | null {
if (modelId) return modelId;
if (this.selectedChatModel) return this.selectedChatModel;
const requestedModelId = modelId || this.selectedChatModel;
// Try to get model from first running instance
// Only models with a placed instance can receive requests; disk downloads alone are not enough.
for (const [, instanceWrapper] of Object.entries(this.instances)) {
if (instanceWrapper && typeof instanceWrapper === "object") {
const keys = Object.keys(instanceWrapper as Record<string, unknown>);
@@ -2264,8 +2263,15 @@ class AppStore {
const instance = (instanceWrapper as Record<string, unknown>)[
keys[0]
] as { shardAssignments?: { modelId?: string } };
if (instance?.shardAssignments?.modelId) {
return instance.shardAssignments.modelId;
const instanceModelId = instance?.shardAssignments?.modelId;
// ensure to only return requestedModelId that matches an instance
// or fall back to first instance
if (
instanceModelId &&
(!requestedModelId || requestedModelId === instanceModelId)
) {
return instanceModelId;
}
}
}
+18 -8
View File
@@ -1461,6 +1461,9 @@
addToast({ type: "info", message: `Launching model...` });
// Always auto-select the newly launched model so the user chats to what they just launched
setSelectedChatModel(modelId);
userForcedIdle = false;
pendingChatModelId = modelId;
chatLaunchState = "launching";
// Record the launch in recent models history
recordRecentLaunch(modelId);
@@ -2547,12 +2550,10 @@
];
// ── Seamless chat: launch models from chat view ──
type ChatLaunchState =
| "idle"
| "launching"
| "downloading"
| "loading"
| "ready";
type InFlightChatLaunchState = "launching" | "downloading" | "loading";
type ReadyLikeChatLaunchState = "idle" | "ready";
type ChatLaunchState = InFlightChatLaunchState | ReadyLikeChatLaunchState;
let chatLaunchState = $state<ChatLaunchState>("idle");
let pendingChatModelId = $state<string | null>(null);
let selectedChatCategory = $state<string | null>(null);
@@ -3129,6 +3130,15 @@
if (model) {
pendingAutoMessage = { content, files };
userForcedIdle = false;
// The selected model is already being placed or loaded; keep the queued
// message and let the existing launch state effects send it once ready.
if (
pendingChatModelId === model &&
chatLaunchState !== "idle" &&
chatLaunchState !== "ready"
) {
return;
}
launchModelForChat(model, "picker", messages().length > 0);
return;
}
@@ -4603,7 +4613,7 @@
type="button"
onclick={() => {
completeOnboarding();
sendMessage(chip, undefined, thinkingEnabled());
handleChatSend(chip);
}}
class="px-4 py-2 rounded-full border border-white/10 bg-white/5 text-sm text-white/60 hover:bg-white/10 hover:text-white/80 hover:border-white/20 transition-all duration-200 cursor-pointer"
>
@@ -6100,7 +6110,7 @@
onclick={() => {
chatLaunchState = "idle";
selectedChatCategory = null;
sendMessage(prompt, undefined, thinkingEnabled());
handleChatSend(prompt);
}}
class="text-left px-3 py-2.5 text-xs text-exo-light-gray hover:text-white font-mono rounded-lg border border-exo-medium-gray/30 hover:border-exo-yellow/30 bg-exo-dark-gray/30 hover:bg-exo-dark-gray/60 transition-all duration-200 cursor-pointer"
>
+35 -6
View File
@@ -66,7 +66,9 @@ Creates a new model instance in the cluster.
```
**Response:**
JSON description of the created instance.
Command acknowledgement. Instance creation is asynchronous; clients should wait
for the model to appear through `/instance/await` before sending inference
requests for that model.
### Delete Instance
@@ -94,6 +96,31 @@ Returns details of a specific instance.
**Response:**
JSON description of the instance.
### Await Instance
**GET** `/instance/await?model_id=...&timeout_seconds=0`
Waits until API state contains an instance for the requested model. The response
is an SSE stream so clients receive keep-alive comments while waiting.
**Query parameters:**
* `model_id`: string, required
* `timeout_seconds`: float, optional, default `0`. `0` waits indefinitely;
positive values time out after that many seconds. Maximum positive value:
`300`.
**Stream messages:**
```text
data: {"type": "ready", "instance": {...}}
data: {"type": "timeout", "message": "No instance found for model ..."}
```
The HTTP status is `200` for both messages because the stream starts before the
final result is known. The `type` field disambiguates the terminal message.
### Preview Placements
**GET** `/instance/previews?model_id=...`
@@ -123,17 +150,18 @@ Computes a placement for a potential instance without creating it.
**Response:**
JSON object describing the proposed placement / instance configuration.
### Place Instance (Dry Operation)
### Place Instance
**POST** `/place_instance`
Performs a placement operation for an instance (planning step), without necessarily creating it.
Places an instance for a model using the server's placement logic.
**Request body:**
JSON describing the instance to be placed.
**Response:**
Placement result.
Command acknowledgement. The instance may not be ready immediately; wait for it
to appear through `/instance/await` before sending inference requests.
## 3. Models
@@ -639,10 +667,11 @@ GET /events
# Instance Management
POST /instance
GET /instance/{instance_id}
DELETE /instance/{instance_id}
GET /instance/await
GET /instance/previews
GET /instance/placement
GET /instance/{instance_id}
DELETE /instance/{instance_id}
POST /place_instance
# Models
-84
View File
@@ -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.
Generated
+6 -6
View File
@@ -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": {
+1 -1
View File
@@ -110,7 +110,7 @@
nixpkgs-fmt.enable = true;
ruff-format = {
enable = true;
excludes = [ "rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi" ];
excludes = [ "rust/exo_rs/exo_rs.pyi" ];
};
rustfmt = {
enable = true;
+4 -4
View File
@@ -16,14 +16,14 @@ 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
uv sync --reinstall-package exo_pyo3_bindings
uv sync --reinstall-package exo_rs
build-dashboard:
#!/usr/bin/env bash
@@ -37,7 +37,7 @@ package: build-dashboard
rm -rf build
build-app: rust-rebuild sync-clean package
xcodebuild build -project app/EXO/EXO.xcodeproj -scheme EXO -configuration Debug -derivedDataPath app/EXO/build
env -u LD xcodebuild build -project app/EXO/EXO.xcodeproj -scheme EXO -configuration Debug -derivedDataPath app/EXO/build
@echo "\nBuild complete. Run with:\n open {{justfile_directory()}}/app/EXO/build/Build/Products/Debug/EXO.app"
clean:
+12 -12
View File
@@ -15,7 +15,7 @@ dependencies = [
"huggingface-hub>=1.8.0",
"psutil>=7.0.0",
"loguru>=0.7.3",
"exo-pyo3-bindings", # rust bindings
"exo-rs", # rust bindings
"anyio==4.11.0",
"tiktoken>=0.12.0", # required for kimi k2 tokenizer
"hypercorn>=0.18.0",
@@ -26,6 +26,7 @@ dependencies = [
"msgspec>=0.19.0",
"zstandard>=0.23.0",
"transformers>=5.6.2",
"python-daemon>=3.1.2",
]
[project.scripts]
@@ -75,22 +76,14 @@ mlx-cuda13 = [
###
[tool.uv.workspace]
members = ["rust/exo_pyo3_bindings", "bench", "tools"]
members = ["rust/exo_rs", "bench", "tools"]
[tool.uv.sources]
exo-pyo3-bindings = { workspace = true }
exo-rs = { workspace = true }
mlx = [
{ git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" },
{ url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
{ url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine != 'aarch64'" },
]
mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/deepseek-v4" }
mflux = { git = "https://github.com/evanev7/mflux", branch = "exo2" }
torch = [
{ index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'mlx-cpu' and extra != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
{ index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13' " },
{ index = "pytorch-cu130", marker = "sys_platform == 'linux' and extra == 'mlx-cuda13'" },
]
mlx-cuda-12 = [
{ url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx_cuda_12-0.32.0-py3-none-manylinux_2_35_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
@@ -100,6 +93,13 @@ mlx-cuda-13 = [
{ url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx_cuda_13-0.32.0-py3-none-manylinux_2_35_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
{ url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx_cuda_13-0.32.0-py3-none-manylinux_2_35_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine != 'aarch64'" },
]
mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/deepseek-v4" }
mflux = { git = "https://github.com/evanev7/mflux", branch = "exo2" }
torch = [
{ index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'mlx-cpu' and extra != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
{ index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13' " },
{ index = "pytorch-cu130", marker = "sys_platform == 'linux' and extra == 'mlx-cuda13'" },
]
torchvision = [
{ index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'mlx-cpu' and extra != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
{ index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13'" },
@@ -240,7 +240,7 @@ torchaudio = ["torch"]
###
[tool.ruff]
extend-exclude = [".typings/**", "rust/exo_pyo3_bindings/**", "bench/vendor/**"]
extend-exclude = [".typings/**", "rust/exo_rs/**", "bench/vendor/**"]
[tool.ruff.lint]
extend-select = ["I", "N", "B", "A", "PIE", "SIM"]
+9 -8
View File
@@ -44,20 +44,21 @@ let
paths = builtins.concatMap (p: [ (lib.getBin p) (lib.getLib p) (lib.getDev p) ]) (cudaLibs ++ [ cudaPackages.cuda_nvcc cuda_cccl_compat ]);
};
exoOverlay = final: prev: {
# Replace workspace exo_pyo3_bindings with Nix-built wheel.
# Replace workspace exo_rs 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-rs = pkgs.stdenv.mkDerivation {
pname = "exo-rs";
version = "0.1.0";
src = self'.packages.exo_pyo3_bindings;
src = self'.packages.exo-rs;
# Install from pre-built wheel
nativeBuildInputs = [ final.pyprojectWheelHook ];
dontStrip = true;
passthru = prev.exo-pyo3-bindings.passthru or { };
passthru = prev.exo-rs.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_rs
cp ${inputs.self}/rust/exo_rs/exo_rs.pyi $siteDir/
touch $siteDir/py.typed
'';
};
@@ -223,7 +224,7 @@ let
++ lib.optionals isDarwin [ pkgs.macmon ];
passthru = {
venv = venv name;
evenv = ((pythonSet.overrideScope editableOverlay).mkVirtualEnv "${name}-evenv" (members // { exo = (members.exo or [ ]) ++ [ "dev" ]; })).overrideAttrs (_: {
evenv = ((pythonSet.overrideScope editableOverlay).mkVirtualEnv "${name}-evenv" (members // { exo = (members.exo or [ ]) ++ [ "dev" ]; exo-rs = [ ]; })).overrideAttrs (_: {
venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" "lib/python${python.pythonVersion}/site-packages/build_backend.py" ];
});
} // lib.optionalAttrs cudaSupport {
@@ -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
-69
View File
@@ -1,69 +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"
pidfile-rs = "0.3"
# macro dependencies
extend = { workspace = true }
delegate = { workspace = true }
thiserror = "2.0"
# 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"
-47
View File
@@ -1,47 +0,0 @@
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};
/// Identity keypair of a node.
#[gen_stub_pyclass]
#[pyclass(name = "Keypair", frozen)]
#[repr(transparent)]
pub struct PyKeypair(pub Keypair);
#[gen_stub_pymethods]
#[pymethods]
#[allow(clippy::needless_pass_by_value)]
impl PyKeypair {
/// Generate a new Ed25519 keypair.
#[staticmethod]
fn generate() -> Self {
Self(Keypair::generate_ed25519())
}
/// Construct an Ed25519 keypair from secret key bytes
#[staticmethod]
fn from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
let mut bytes = Vec::from(bytes.as_bytes());
Ok(Self(Keypair::ed25519_from_bytes(&mut bytes).pyerr()?))
}
/// Get the secret key bytes underlying the keypair
fn to_bytes<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
let bytes = self
.0
.clone()
.try_into_ed25519()
.pyerr()?
.secret()
.as_ref()
.to_vec();
Ok(PyBytes::new(py, &bytes))
}
/// Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
fn to_node_id(&self) -> String {
self.0.public().to_peer_id().to_base58()
}
}
-318
View File
@@ -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(())
}
+53
View File
@@ -0,0 +1,53 @@
[package]
name = "exo_rs"
version = { workspace = true }
edition = { workspace = true }
publish = false
[lib]
doctest = false
path = "src/lib.rs"
name = "exo_rs"
# "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
pidfile-rs = { 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
File renamed without changes.
@@ -1,50 +1,41 @@
# This file is automatically generated by pyo3_stub_gen
# ruff: noqa: E501, F401
# ruff: noqa: E501, F401, F403, F405
import builtins
import os
import pathlib
import typing
__all__ = [
"FromSwarm",
"NetworkingHandle",
"Pidfile",
"PidfileError",
]
@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: ...
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:
def __new__(cls, identity: Keypair, bootstrap_peers: typing.Sequence[builtins.str], listen_port: builtins.int) -> NetworkingHandle: ...
@staticmethod
def new(identity: builtins.str, namespace: builtins.str, listen_port: builtins.int, discovery_service_port: builtins.int) -> NetworkingHandle: ...
def recv(self) -> typing.Awaitable[FromSwarm]: ...
async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
r"""
Subscribe to a `GossipSub` topic.
@@ -63,13 +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 NoPeersSubscribedToTopicError(builtins.Exception):
def __new__(cls, *args: typing.Any) -> NoPeersSubscribedToTopicError: ...
def __repr__(self) -> builtins.str: ...
def __str__(self) -> builtins.str: ...
@typing.final
class Pidfile:
@@ -77,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.
@@ -107,32 +91,23 @@ class Pidfile:
The file is truncated before writing.
"""
def as_raw_fd(self) -> builtins.int:
r"""
Extracts the raw file descriptor.
This function is typically used to **borrow** an owned file descriptor.
When used in this way, this method does **not** pass ownership of the
raw file descriptor to the caller, and the file descriptor is only
guaranteed to be valid while the original object has not yet been
destroyed.
"""
def close(self) -> None:
r"""
Closes the PID file and releases associated resources.
"""
@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__ = ("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: ...
...
@@ -3,27 +3,31 @@ requires = ["maturin>=1.0,<2.0"]
build-backend = "maturin"
[project]
name = "exo_pyo3_bindings"
version = "0.2.2"
name = "exo_rs"
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_rs", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"]
[tool.maturin]
#purelib = true
#python-source = "python"
module-name = "exo_pyo3_bindings"
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" }]
@@ -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_rs::stub_info()?;
stub.generate()?;
Ok(())
}
@@ -5,23 +5,16 @@
//!
mod allow_threading;
mod ident;
// mod ident;
mod networking;
mod pidfile;
use crate::ident::PyKeypair;
use crate::networking::networking_submodule;
use crate::pidfile::pidfile_submodule;
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;
@@ -153,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_pyo3_bindings")]
#[pymodule(name = "exo_rs", gil_used = true)]
fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
// install logger
pyo3_log::init();
@@ -164,9 +157,9 @@ 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>()?;
networking_submodule(m)?;
pidfile_submodule(m)?;
// m.add_class::<PyKeypair>()?;
networking_submodule(m)?;
// top-level constructs
// TODO: ...
+197
View File
@@ -0,0 +1,197 @@
use std::pin::Pin;
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, Swarm, ToSwarm, create_swarm};
use networking::{Session, is_valid_zid};
use pyo3::exceptions::{PyRuntimeError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::PyBytes;
use pyo3::{Bound, Py, PyAny, PyErr, PyResult, Python, pymethods};
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods};
use tokio::sync::{Mutex, mpsc, oneshot};
#[gen_stub_pyclass]
#[pyclass(name = "NetworkingHandle")]
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(name = "FromSwarm")]
pub 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(),
},
}
}
}
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 {
// NOTE: `async fn`s here that use `.await` will wrap the future in `.allow_threads_py()`
// immediately beforehand to release the interpreter.
// SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
// ---- Lifecycle management methods ----
#[staticmethod]
pub fn new(
identity: &str,
namespace: &str,
listen_port: u16,
discovery_service_port: u16,
) -> PyResult<PyNetworkingHandle> {
// todo: zenoh self assigned peers
if listen_port == 0 {
todo!("cannot listen on port 0 yet");
}
// create communication channels
let (to_swarm, from_client) = mpsc::channel(1024);
// get identity
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,
))
.pyerr()?;
Ok(PyNetworkingHandle {
swarm: Arc::new(Mutex::new(swarm.into_stream())),
to_swarm,
})
}
#[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
.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.
pub 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.
pub 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.
pub 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(())
}
}
pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PyNetworkingHandle>()?;
m.add_class::<PyFromSwarm>()?;
Ok(())
}
@@ -3,7 +3,9 @@ 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;
use std::fs::Permissions;
use std::os::fd::{AsRawFd, RawFd};
use std::os::unix::prelude::PermissionsExt;
use std::path::PathBuf;
@@ -36,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.
///
@@ -53,7 +55,23 @@ impl PyPidfileError {
/// [`daemon`(3)]: https://linux.die.net/man/3/daemon
#[gen_stub_pyclass]
#[pyclass(name = "Pidfile")]
pub struct PyPidfile(Pidfile);
pub struct PyPidfile(Option<Pidfile>);
impl PyPidfile {
#[inline(always)]
fn get(&self) -> &Pidfile {
self.0
.as_ref()
.expect("cannot use resource after exiting context")
}
#[inline(always)]
fn get_mut(&mut self) -> &mut Pidfile {
self.0
.as_mut()
.expect("cannot use resource after exiting context")
}
}
#[gen_stub_pymethods]
#[pymethods]
@@ -65,17 +83,40 @@ impl PyPidfile {
/// 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))?,
))
// create all parent directories if don't exist
if let Some(parent) = path.parent() {
fs::create_dir_all(parent)
.map_err(|e| PyPidfileError(PidfileError::Io(e)).into_pyerr(py))?;
}
let pidfile = Pidfile::new(&path, Permissions::from_mode(mode))
.map_err(|e| PyPidfileError(e).into_pyerr(py))?;
Ok(Self(Some(pidfile)))
}
/// 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))
self.get_mut()
.write()
.map_err(|e| PyPidfileError(e).into_pyerr(py))
}
/// Extracts the raw file descriptor.
///
/// This function is typically used to **borrow** an owned file descriptor.
/// When used in this way, this method does **not** pass ownership of the
/// raw file descriptor to the caller, and the file descriptor is only
/// guaranteed to be valid while the original object has not yet been
/// destroyed.
fn as_raw_fd(&self) -> RawFd {
self.get().as_raw_fd()
}
/// Closes the PID file and releases associated resources.
fn close(&mut self) {
self.0 = None;
}
}
File renamed without changes.
@@ -1,31 +1,28 @@
import asyncio
import os
import pytest
from _pytest.capture import CaptureFixture
from exo_pyo3_bindings import (
Keypair,
from exo_rs import (
NetworkingHandle,
NoPeersSubscribedToTopicError,
Pidfile,
PyFromSwarm,
FromSwarm,
)
@pytest.mark.asyncio
async def test_sleep_on_multiple_items() -> None:
print("PYTHON: starting handle")
h = NetworkingHandle(Keypair.generate(), [], 0)
h = NetworkingHandle.new(os.urandom(16).hex().lstrip("0"), 52414, 52413)
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]):
@@ -39,11 +36,15 @@ 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}")
def scoped_lock_file():
a = Pidfile("/tmp/lock.pid", 0o0600)
if __name__ == "__main__":
asyncio.run(test_sleep_on_multiple_items())
+20 -35
View File
@@ -1,42 +1,27 @@
[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"] }
socket2.workspace = true
blake3.workspace = true
[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"
[dev-dependencies]
env_logger.workspace = true
smol.workspace = true
tracing.workspace = true
-86
View File
@@ -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 => {}
}
}
}
+34
View File
@@ -0,0 +1,34 @@
use networking;
use tracing::{info, warn};
use zenoh::{Result, Wait};
#[tokio::main]
async fn main() -> Result<()> {
zenoh::init_log_from_env_or("info");
info!("Opening session...");
let cfg = networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414)?;
let session = networking::open(cfg, "exo", 52414, 52413).await?;
let _tok = session
.z
.liveliness()
.declare_token(format!("nodes/{}/live", session.z.zid()))
.wait()?;
let subs = session
.z
.liveliness()
.declare_subscriber("**")
.history(true)
.wait()?;
loop {
tokio::select! {
_ = tokio::signal::ctrl_c() => break,
s = subs.recv_async() => {
match s {
Err(e) => warn!("{e}"),
Ok(s) => info!("{}: {}", s.kind(), s.key_expr().to_string().split("/").nth(1).unwrap()),
}
}
}
}
Ok(())
}
+32
View File
@@ -0,0 +1,32 @@
use env_logger::Env;
use log::info;
use networking;
use zenoh::{Result, Wait};
#[tokio::main]
async fn main() -> Result<()> {
env_logger::try_init_from_env(Env::new().default_filter_or("info")).expect("logger failed");
info!("Opening session...");
let cfg = networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414)?;
let session = networking::open(cfg, "exo", 52414, 52413).await?;
let _tok = session
.z
.liveliness()
.declare_token(format!("nodes/{}/live", session.z.zid()))
.wait()?;
session
.z
.liveliness()
.declare_subscriber("**")
.history(true)
.callback(|tok| info!("{}: {}", tok.kind(), tok.key_expr().to_string()))
.background()
.wait()?;
loop {
tokio::select! {
_ = tokio::signal::ctrl_c() => break,
_ = session.z.put("hello", "world") => {},
}
}
Ok(())
}
+48
View File
@@ -0,0 +1,48 @@
use std::{env, time::Duration};
use env_logger::Env;
use log::info;
use networking;
use zenoh::Result;
#[tokio::main]
async fn main() -> Result<()> {
env_logger::try_init_from_env(Env::new().default_filter_or("info")).expect("logger failed");
let n_bytes = env::args()
.nth(1)
.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(&format!("{:x}", rand::random::<u128>()), 52414)?;
let session = networking::open(cfg, "exo", 52414, 52413).await?;
let _tok = session
.z
.liveliness()
.declare_token(format!("nodes/{}/live", session.z.zid()))
.await?;
let key_expr = "stream/data";
let payload = "n".repeat(n_bytes);
let pubs = session
.z
.declare_publisher(key_expr)
.congestion_control(zenoh::qos::CongestionControl::Block)
.await?;
let pubs_l = pubs.matching_listener().await?;
if !pubs.matching_status().await?.matching() {
while !pubs_l.recv_async().await?.matching() {}
}
tokio::time::sleep(Duration::from_secs(1)).await;
info!("Putting Data ('{key_expr}': '{}')...", payload.len());
for _ in 0..10 {
let t = tokio::time::Instant::now();
for _ in 0..5000 {
pubs.put(payload.clone()).await?;
}
info!("{:?}", t.elapsed());
tokio::time::sleep(Duration::from_secs(1)).await;
}
tokio::signal::ctrl_c().await?;
Ok(())
}
+74
View File
@@ -0,0 +1,74 @@
use std::time::Duration;
use env_logger::Env;
use log::info;
use networking;
use zenoh::Result;
#[tokio::main]
async fn main() -> Result<()> {
env_logger::try_init_from_env(Env::new().default_filter_or("info")).expect("logger failed");
info!("Opening session...");
let cfg = networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414)?;
let session = networking::open(cfg, "exo", 52414, 52413).await?;
let _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| {
info!(
"{}: {}",
tok.kind(),
tok.key_expr()
.to_string()
.strip_prefix("nodes/")
.and_then(|it| it.strip_suffix("/live"))
.unwrap()
)
})
.await?;
let watch = async {
for _ in 0..1000 {
tokio::time::sleep(Duration::from_secs(1)).await;
session
.z
.get("**")
.callback(|reply| {
let sample = reply.into_result().expect("no errs");
info!(
"got {} bytes on {}",
sample.payload().len(),
sample.key_expr()
)
})
.await?;
}
Result::<()>::Ok(())
};
let subs = session.z.declare_subscriber("**").await?;
let mut i = 0;
let _a = async {
while let Ok(sample) = subs.recv_async().await {
i += 1;
info!(
"[{i}] received {} bytes on {}",
sample.payload().len(),
sample.key_expr()
)
}
};
tokio::select! {
_ = watch => {},
_ = _a => {},
_ = tokio::signal::ctrl_c() => {},
}
Ok(())
}
+41
View File
@@ -0,0 +1,41 @@
use std::{borrow::Cow, env};
use env_logger::Env;
use log::{info, warn};
use networking;
use zenoh::{Result, Wait};
#[tokio::main]
async fn main() -> Result<()> {
env_logger::try_init_from_env(Env::new().default_filter_or("info")).expect("logger failed");
info!("Opening session...");
let cfg = networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414)?;
let session = networking::open(cfg, "exo", 52414, 52413).await?;
let other_live = session
.z
.liveliness()
.declare_subscriber("**")
.history(true)
.wait()?;
_ = other_live.recv_async().await?;
let other_live = session.z.liveliness().get("**").wait()?;
while let Ok(s) = other_live.recv_async().await {
info!("{s:?}");
}
let query = env::args().nth(1).expect("USAGE: z_get [query]");
info!("Querying {query}");
let subs = session.z.liveliness().get(query).await?;
while let Ok(r) = subs.recv_async().await {
match r.into_result() {
Ok(s) => info!(
"{}: {}",
s.key_expr(),
s.payload()
.try_to_string()
.unwrap_or_else(|_| Cow::Borrowed("-bytes-"))
),
Err(e) => warn!("{e}"),
}
}
Ok(())
}
-44
View File
@@ -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
+316 -367
View File
@@ -1,390 +1,339 @@
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,
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 MAGIC: [u8; 3] = *b"EXO";
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<SocketAddrV6>>>,
namespace: [u8; 8],
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,
},
#[derive(Debug, Clone, Copy)]
pub struct Discovered {
pub zid: ZenohId,
pub addr: SocketAddrV6,
}
/// 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,
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({
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);
}
},
// handle ping events => if error then disconnect
managed::BehaviourEvent::Ping(e) => {
if let Err(_) = e.result {
self.close_connection(e.peer, e.connection.clone())
match sock.join_multicast_v6(&GROUP, *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) {
warn!(
"failed to join multicast v6 for interface {}: {e}",
iface.name
)
}
}
_ => {}
}
}
for iface_idx in update.diff.removed {
ifaces.lock().retain(|addr| addr.scope_id() != iface_idx);
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,
namespace,
ifaces,
last_nonce: Mutex::new(rand::random()),
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, &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: &[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 != MAGIC {
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);
}
if hello.namespace != self.namespace {
trace!("dropped: different namespace");
return Ok(None);
}
// reply
trace!("replying to Hello({:?})", hello.nonce);
let reply = WhatsUp {
nonce: hello.nonce,
zid: self.zid.to_le_bytes(),
port_le: self.listen_port.to_le_bytes(),
}
.alloc();
for i in 1..6 {
if self
.sock
.send_to(&reply, 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
);
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 != *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);
}
// discovery success!
// the incoming port is our listen port;
// overwrite it with the whats_up port corresponding to the remote zenoh service
let addr = {
let mut x = v6;
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 buf = Hello {
nonce,
namespace: self.namespace,
}
.alloc();
// wait for pending events
Poll::Pending
let addrs = self.ifaces.lock().clone();
debug!("announcing Hello({nonce:?}) 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() == io::ErrorKind::HostUnreachable => {
debug!("disabling discovery address {addr}: {e}");
_ = self.ifaces.lock().swap_remove(i);
}
Err(e) => debug!("failed to reach {addr}: {e}"),
}
}
Ok(())
}
}
#[repr(u8)]
#[derive(Debug, Clone, Copy)]
// packet & version
pub enum Kind {
Hello = 0,
WhatsUp = 1,
}
pub struct UnknownKind;
impl TryFrom<u8> for Kind {
type Error = UnknownKind;
fn try_from(value: u8) -> Result<Self, Self::Error> {
match value {
0 => Ok(Self::Hello),
1 => Ok(Self::WhatsUp),
_ => Err(UnknownKind),
}
}
}
pub trait Message: Pod {
const KIND: Kind;
}
// should be part of the Message trait, but const in traits isnt stabilized. this lets alloc :: Self -> [u8; Self::buf_size()]
macro_rules! impl_alloc {
($a:ident) => {
impl $a {
const fn buf_size() -> usize {
size_of::<Header>() + size_of::<Self>()
}
pub fn alloc(self) -> [u8; Self::buf_size()] {
let mut buf = [0u8; Self::buf_size()];
buf[0..size_of::<Header>()].copy_from_slice(bytemuck::bytes_of(&Header {
magic: MAGIC,
kind: Self::KIND as u8,
}));
buf[size_of::<Header>()..Self::buf_size()]
.copy_from_slice(bytemuck::bytes_of(&self));
buf
}
}
};
}
#[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],
pub namespace: [u8; 8],
}
impl Message for Hello {
const KIND: Kind = Kind::Hello;
}
impl_alloc!(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 Message for WhatsUp {
const KIND: Kind = Kind::WhatsUp;
}
impl_alloc!(WhatsUp);
+105 -33
View File
@@ -1,44 +1,116 @@
//! TODO: crate documentation
//!
//! this is here as a placeholder documentation
//!
//!
use std::sync::Arc;
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 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
}
/// 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 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}\""))?;
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("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("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)
}
#[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;
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)
.plugins_manager(plugins)
.build()
.await?;
let z = zenoh::session::init(runtime.clone().into()).await?;
runtime.start().await?;
let mut discovery =
Discovery::new(z.zid(), namespace, listen_port, discovery_service_port).await?;
let _jh = Arc::new(AbortOnDrop(tokio::task::spawn(async move {
loop {
let Ok(discovered) = discovery.next().await.inspect_err(|e| {
log::warn!("discovery error {e}");
}) else {
continue;
};
let Some(Protocol::Tcp(port)) = ps.next() 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 parse locator from addr: {e}");
})
else {
continue;
};
Some((ip, port))
runtime
.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<AbortOnDrop>,
}
+154 -232
View File
@@ -1,24 +1,22 @@
//! 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::Publisher;
use zenoh::pubsub::Subscriber;
use zenoh::qos::CongestionControl;
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 +24,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<()>, Publisher<'static>)>;
pub struct Swarm {
swarm: libp2p::Swarm<Behaviour>,
from_client: mpsc::Receiver<ToSwarm>,
pub session: crate::Session,
pub 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 zid = key_expr.strip_prefix("live/");
yield match token.kind() {
SampleKind::Put => {
log::info!("discovered: {zid:?}");
FromSwarm::Discovered {}
}
SampleKind::Delete => {
log::info!("expired: {zid:?}");
FromSwarm::Expired {}
}
}
}
}
}
}
};
@@ -79,208 +91,118 @@ 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!("live/{}", session.zid()))
.await?;
let sub = session
.liveliness()
.declare_subscriber("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 = match topics.get(&topic) {
Some(topic) => topic.1.put(data).await,
None => {
// TODO: this should be an error but the python FromSwarm is somewhat nondeterministic
Ok(()) //Err("not subscribed to topic!".into()),
}
};
_ = result_sender.send(res);
}
ToSwarm::Unsubscribe {
topic,
result_sender,
} => {
let Some((_, (subscriber, publisher))) = topics.remove_entry(&topic) else {
_ = result_sender.send(false);
return;
};
_ = publisher.undeclare().await;
_ = 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 publisher_res = session
.declare_publisher(format!("topics/{topic}"))
.congestion_control(CongestionControl::Block)
.await;
let publisher = match publisher_res {
Ok(p) => p,
Err(e) => {
_ = result_sender.send(Err(e));
return;
}
};
let subscriber_res = 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;
let subscriber = match subscriber_res {
Ok(s) => s,
Err(e) => {
_ = result_sender.send(Err(e));
return;
}
};
assert!(topics.insert(topic, (subscriber, publisher)).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: &str,
namespace: &str,
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")
}
discovery_service_port: u16,
) -> Result<Swarm> {
let cfg = crate::cfg(identity, listen_port)?;
let session = crate::open(cfg, namespace, listen_port, discovery_service_port).await?;
Ok(Swarm {
session,
from_client,
})
}
-107
View File
@@ -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");
}
-7
View File
@@ -1,7 +0,0 @@
// maybe this will hold test in the future...??
#[cfg(test)]
mod tests {
#[test]
fn does_nothing() {}
}
+4 -3
View File
@@ -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-rs = craneLib.buildPackage (
commonArgs
// {
inherit cargoArtifacts;
pname = "exo_pyo3_bindings";
pname = "exo-rs";
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_rs/Cargo.toml \
--features "pyo3/extension-module,pyo3/experimental-async" \
--interpreter ${pkgs.python313}/bin/python \
--out dist
-15
View File
@@ -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]
-1
View File
@@ -1 +0,0 @@
pub mod wakerdeque;
-55
View File
@@ -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);
}
}
+3
View File
@@ -0,0 +1,3 @@
from importlib.metadata import version
__version__ = version("exo")
+88 -42
View File
@@ -20,7 +20,7 @@ from fastapi.staticfiles import StaticFiles
from hypercorn.asyncio import serve # pyright: ignore[reportUnknownVariableType]
from hypercorn.config import Config
from hypercorn.typing import ASGIFramework
from hypercorn.utils import LifespanTimeoutError
from hypercorn.utils import LifespanTimeoutError, ShutdownError
from loguru import logger
from exo.api.adapters.chat_completions import (
@@ -50,6 +50,8 @@ from exo.api.keepalive import with_sse_keepalive
from exo.api.types import (
AddCustomModelParams,
AdvancedImageParams,
AwaitInstanceReadyMessage,
AwaitInstanceTimeoutMessage,
BenchChatCompletionRequest,
BenchChatCompletionResponse,
BenchImageGenerationResponse,
@@ -344,6 +346,7 @@ class API:
self.app.post("/place_instance")(self.place_instance)
self.app.get("/instance/placement")(self.get_placement)
self.app.get("/instance/previews")(self.get_placement_previews)
self.app.get("/instance/await", response_model=None)(self.await_instance)
self.app.get("/instance/{instance_id}")(self.get_instance)
self.app.delete("/instance/{instance_id}")(self.delete_instance)
self.app.get("/v1/instance-links")(self.list_instance_links)
@@ -633,6 +636,48 @@ class API:
raise HTTPException(status_code=404, detail="Instance not found")
return self.state.instances[instance_id]
async def await_instance(
self,
model_id: ModelId,
timeout_seconds: float = Query(default=0.0, ge=0.0, le=300.0),
) -> StreamingResponse:
_sleep = 0.1
async def _stream() -> AsyncGenerator[str, None]:
deadline = (
None if timeout_seconds == 0 else anyio.current_time() + timeout_seconds
)
while True:
for instance in self.state.instances.values():
if instance.shard_assignments.model_id == model_id:
payload = AwaitInstanceReadyMessage(instance=instance)
yield f"data: {payload.model_dump_json()}\n\n"
return
if deadline is None:
await anyio.sleep(_sleep)
else:
remaining = deadline - anyio.current_time()
if remaining <= 0:
payload = AwaitInstanceTimeoutMessage(
message=f"No instance found for model {model_id}"
)
yield f"data: {payload.model_dump_json()}\n\n"
return
await anyio.sleep(min(_sleep, remaining))
return StreamingResponse(
with_sse_keepalive(_stream()),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "close",
"X-Accel-Buffering": "no",
},
)
async def delete_instance(self, instance_id: InstanceId) -> DeleteInstanceResponse:
if instance_id not in self.state.instances:
raise HTTPException(status_code=404, detail="Instance not found")
@@ -761,6 +806,8 @@ class API:
if isinstance(chunk, PrefillProgressChunk):
continue
sampler.mark_prefill_done()
if chunk.finish_reason == "error":
raise HTTPException(
status_code=500,
@@ -871,10 +918,8 @@ class API:
) -> ChatCompletionResponse | StreamingResponse:
"""OpenAI Chat Completions API - adapter."""
task_params = await chat_request_to_text_generation(payload)
resolved_model = await self._resolve_and_validate_text_model(
ModelId(task_params.model)
)
task_params = task_params.model_copy(update={"model": resolved_model})
validated_model = await self._validate_model_has_instance(task_params.model)
task_params = task_params.model_copy(update={"model": validated_model})
command = await self._send_text_generation_with_images(task_params)
@@ -906,10 +951,10 @@ class API:
self, payload: BenchChatCompletionRequest
) -> BenchChatCompletionResponse | StreamingResponse:
task_params = await chat_request_to_text_generation(payload)
resolved_model = await self._resolve_and_validate_text_model(
validated_model = await self._validate_model_has_instance(
ModelId(task_params.model)
)
task_params = task_params.model_copy(update={"model": resolved_model})
task_params = task_params.model_copy(update={"model": validated_model})
task_params = task_params.model_copy(
update={
@@ -939,8 +984,10 @@ class API:
return await self._collect_text_generation_with_stats(command.command_id)
async def _resolve_and_validate_text_model(self, model_id: ModelId) -> ModelId:
"""Validate a text model exists and return the resolved model ID.
async def _validate_model_has_instance(self, model_id: ModelId) -> ModelId:
"""Validate a model has an active instance.
If the model isn't even downloaded, triggers notification to user to download model.
Raises HTTPException 404 if no instance is found for the model.
"""
@@ -948,30 +995,21 @@ class API:
instance.shard_assignments.model_id == model_id
for instance in self.state.instances.values()
):
await self._trigger_notify_user_to_download_model(model_id)
# Check if model is actually downloaded
model_is_downloaded = any(
isinstance(download, DownloadCompleted)
and download.shard_metadata.model_card.model_id == model_id
for node_downloads in self.state.downloads.values()
for download in node_downloads
)
if not model_is_downloaded:
await self._trigger_notify_user_to_download_model(model_id)
raise HTTPException(
status_code=404,
detail=f"No instance found for model {model_id}",
status_code=404, detail=f"No instance found for model {model_id}"
)
return model_id
async def _validate_image_model(self, model: ModelId) -> ModelId:
"""Validate model exists and return resolved model ID.
Raises HTTPException 404 if no instance is found for the model.
"""
model_card = await ModelCard.load(model)
resolved_model = model_card.model_id
if not any(
instance.shard_assignments.model_id == resolved_model
for instance in self.state.instances.values()
):
await self._trigger_notify_user_to_download_model(resolved_model)
raise HTTPException(
status_code=404, detail=f"No instance found for model {resolved_model}"
)
return resolved_model
def stream_events(self) -> StreamingResponse:
def _generate_json_array(events: Iterable[Event]) -> Iterable[str]:
yield "["
@@ -1024,7 +1062,9 @@ class API:
"""
payload = payload.model_copy(
update={
"model": await self._validate_image_model(ModelId(payload.model)),
"model": await self._validate_model_has_instance(
ModelId(payload.model)
),
"advanced_params": _ensure_seed(payload.advanced_params),
}
)
@@ -1292,7 +1332,9 @@ class API:
) -> BenchImageGenerationResponse:
payload = payload.model_copy(
update={
"model": await self._validate_image_model(ModelId(payload.model)),
"model": await self._validate_model_has_instance(
ModelId(payload.model)
),
"stream": False,
"partial_images": 0,
"advanced_params": _ensure_seed(payload.advanced_params),
@@ -1328,7 +1370,7 @@ class API:
advanced_params: AdvancedImageParams | None,
) -> ImageEdits:
"""Prepare and send an image edits command with chunked image upload."""
resolved_model = await self._validate_image_model(model)
validated_model = await self._validate_model_has_instance(model)
advanced_params = _ensure_seed(advanced_params)
image_content = await image.read()
@@ -1347,7 +1389,7 @@ class API:
image_data="",
total_input_chunks=total_chunks,
prompt=prompt,
model=resolved_model,
model=validated_model,
n=n,
size=size,
response_format=response_format,
@@ -1368,7 +1410,7 @@ class API:
await self._send(
SendInputChunk(
chunk=InputImageChunk(
model=resolved_model,
model=validated_model,
command_id=command.command_id,
data=chunk_data,
chunk_index=chunk_index,
@@ -1492,10 +1534,10 @@ class API:
) -> ClaudeMessagesResponse | StreamingResponse:
"""Claude Messages API - adapter."""
task_params = await claude_request_to_text_generation(payload)
resolved_model = await self._resolve_and_validate_text_model(
validated_model = await self._validate_model_has_instance(
ModelId(task_params.model)
)
task_params = task_params.model_copy(update={"model": resolved_model})
task_params = task_params.model_copy(update={"model": validated_model})
command = await self._send_text_generation_with_images(task_params)
@@ -1530,8 +1572,8 @@ class API:
) -> ResponsesResponse | StreamingResponse:
"""OpenAI Responses API."""
task_params = await responses_request_to_text_generation(payload)
resolved_model = await self._resolve_and_validate_text_model(task_params.model)
task_params = task_params.model_copy(update={"model": resolved_model})
validated_model = await self._validate_model_has_instance(task_params.model)
task_params = task_params.model_copy(update={"model": validated_model})
command = await self._send_text_generation_with_images(task_params)
@@ -1573,10 +1615,10 @@ class API:
body = await request.body()
payload = OllamaChatRequest.model_validate_json(body)
task_params = ollama_request_to_text_generation(payload)
resolved_model = await self._resolve_and_validate_text_model(
validated_model = await self._validate_model_has_instance(
ModelId(task_params.model)
)
task_params = task_params.model_copy(update={"model": resolved_model})
task_params = task_params.model_copy(update={"model": validated_model})
command = await self._send_text_generation_with_images(task_params)
@@ -1609,10 +1651,10 @@ class API:
body = await request.body()
payload = OllamaGenerateRequest.model_validate_json(body)
task_params = ollama_generate_request_to_text_generation(payload)
resolved_model = await self._resolve_and_validate_text_model(
validated_model = await self._validate_model_has_instance(
ModelId(task_params.model)
)
task_params = task_params.model_copy(update={"model": resolved_model})
task_params = task_params.model_copy(update={"model": validated_model})
command = await self._send_text_generation_with_images(task_params)
@@ -1914,6 +1956,10 @@ class API:
cfg,
shutdown_trigger=ev.wait,
)
if not ev.is_set():
raise ShutdownError(
"Server exited without shutdown trigger - exiting abnormally"
)
except LifespanTimeoutError as e:
logger.warning(
"Graceful server shutdown timed out, some connections forcebly closed"
+2
View File
@@ -1,5 +1,7 @@
from .api import AddCustomModelParams as AddCustomModelParams
from .api import AdvancedImageParams as AdvancedImageParams
from .api import AwaitInstanceReadyMessage as AwaitInstanceReadyMessage
from .api import AwaitInstanceTimeoutMessage as AwaitInstanceTimeoutMessage
from .api import BenchChatCompletionRequest as BenchChatCompletionRequest
from .api import BenchChatCompletionResponse as BenchChatCompletionResponse
from .api import BenchImageGenerationResponse as BenchImageGenerationResponse
+26
View File
@@ -186,6 +186,12 @@ class NodePowerStats(BaseModel, frozen=True):
node_id: NodeId
samples: int
avg_sys_power: float
# Per-phase breakdown. Populated only when the caller marks a phase
# boundary (e.g. prefill -> generation); None otherwise.
prefill_avg_sys_power: float | None = None
generation_avg_sys_power: float | None = None
prefill_energy_joules: float | None = None
generation_energy_joules: float | None = None
class PowerUsage(BaseModel, frozen=True):
@@ -193,6 +199,16 @@ class PowerUsage(BaseModel, frozen=True):
nodes: list[NodePowerStats]
total_avg_sys_power_watts: float
total_energy_joules: float
# Split between the prefill (prompt-processing) phase and the
# generation/decode phase. Populated only when the caller marks a phase
# boundary; None otherwise. The two phase energies should sum to
# approximately `total_energy_joules` (modulo interpolation rounding).
prefill_seconds: float | None = None
generation_seconds: float | None = None
prefill_energy_joules: float | None = None
generation_energy_joules: float | None = None
prefill_avg_sys_power_watts: float | None = None
generation_avg_sys_power_watts: float | None = None
class BenchChatCompletionResponse(ChatCompletionResponse):
@@ -291,6 +307,16 @@ class DeleteInstanceResponse(BaseModel):
instance_id: InstanceId
class AwaitInstanceReadyMessage(BaseModel):
type: Literal["ready"] = "ready"
instance: Instance
class AwaitInstanceTimeoutMessage(BaseModel):
type: Literal["timeout"] = "timeout"
message: str
class CancelCommandResponse(BaseModel):
message: str
command_id: CommandId
+11
View File
@@ -15,6 +15,10 @@ from exo.download.download_utils import (
resolve_existing_model,
)
from exo.download.shard_downloader import ShardDownloader
from exo.routing.event_router import (
EventRouterBrokenResourceError,
EventRouterClosedResourceError,
)
from exo.shared.constants import EXO_DEFAULT_MODELS_DIR, EXO_MODELS_READ_ONLY_DIRS
from exo.shared.models import model_cards
from exo.shared.models.model_cards import ModelId
@@ -139,7 +143,14 @@ class DownloadCoordinator:
async with self._tg as tg:
tg.start_soon(self._command_processor)
tg.start_soon(self._emit_existing_download_progress)
except* (EventRouterBrokenResourceError, EventRouterClosedResourceError):
# Event router has been closed (try-star syntax handles error groups)
pass
finally:
# don't forget to clean up resources
self.download_command_receiver.close()
self.event_sender.close()
self._stopped.set()
async def shutdown(self) -> None:
+101 -35
View File
@@ -8,23 +8,26 @@ from dataclasses import dataclass, field
from typing import Self
import anyio
from anyio.lowlevel import checkpoint as anyio_checkpoint
from daemon import DaemonContext # pyright: ignore[reportMissingTypeStubs]
from exo_rs import Pidfile, PidfileError
from loguru import logger
from pydantic import PositiveInt
import exo.routing.topics as topics
from exo import __version__
from exo.api.main import API
from exo.download.coordinator import DownloadCoordinator
from exo.download.impl_shard_downloader import exo_shard_downloader
from exo.master.main import Master
from exo.routing.event_router import EventRouter
from exo.routing.router import Router, get_node_id_keypair
from exo.shared.constants import EXO_DEFAULT_MODELS_DIR, EXO_LOG
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
from exo.shared.types.common import NodeId, SessionId
from exo.utils import STDIO_FDS
from exo.utils.channels import Receiver, channel
from exo.utils.daemon import detach_stdio_to_devnull
from exo.utils.pidfile import PidfileLockError, acquire_exo_pidfile
from exo.utils.pydantic_ext import FrozenModel
from exo.utils.task_group import TaskGroup
from exo.worker.main import Worker
@@ -48,13 +51,13 @@ class Node:
@classmethod
async def create(cls, args: "Args") -> Self:
keypair = get_node_id_keypair()
node_id = NodeId(keypair.to_node_id())
node_id = get_node_zid()
session_id = SessionId(master_node_id=node_id, election_clock=0)
router = Router.create(
keypair,
bootstrap_peers=args.bootstrap_peers,
listen_port=args.libp2p_port,
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)
@@ -190,7 +193,7 @@ class Node:
# - Shut down and re-create the API
if result.is_new_master:
await anyio.sleep(0)
await anyio_checkpoint()
self.event_router.shutdown()
self.event_router = EventRouter(
result.session_id,
@@ -203,7 +206,10 @@ class Node:
result.session_id.master_node_id == self.node_id
and self.master is not None
):
logger.info("Node elected Master")
assert not result.is_new_master, (
"cannot be new master if we remain master"
)
logger.info("Node elected Master - maintaining self")
elif (
result.session_id.master_node_id == self.node_id
and self.master is None
@@ -270,14 +276,60 @@ class Node:
def main():
# Exit early if no PID file (not compatible with double-for daemonization yet)
try:
pidfile = acquire_exo_pidfile()
except PidfileLockError as exception:
print(exception, file=sys.stderr)
raise SystemExit(1) from exception
# Parse args first => --help or bad args don't require PID-locking
args = Args.parse()
# Exit early if cannot acquire PID file
try:
pidfile = Pidfile(EXO_PID_FILE, 0o0600)
except PidfileError as e:
print(e, file=sys.stderr)
raise SystemExit(1) from e
try:
if args.legacy_daemon:
# keep stdio backed by explicit /dev/null streams. multiprocessing spawn expects
# valid stdio FDs; letting DaemonContext close/reopen them can break runner startup.
for stream in (sys.stdout, sys.stderr, sys.__stdout__, sys.__stderr__):
if stream is not None:
stream.flush()
stdin = open(os.devnull, "r") # noqa: SIM115
stdout = open(os.devnull, "w") # noqa: SIM115
stderr = open(os.devnull, "w") # noqa: SIM115
with DaemonContext(
detach_process=True,
files_preserve=[pidfile.as_raw_fd()],
stdin=stdin,
stdout=stdout,
stderr=stderr,
):
# cleanup loose file descriptors (as long as they aren't stdio)
for f in (
f for f in (stdin, stdout, stderr) if f.fileno() not in STDIO_FDS
):
f.close()
# 1) if daemonizing => fork then write PID
try:
pidfile.write()
except PidfileError as e:
print(e, file=sys.stderr)
raise SystemExit(1) from e
main_inner(args)
else:
# 2) otherwise => just write PID
try:
pidfile.write()
except PidfileError as e:
print(e, file=sys.stderr)
raise SystemExit(1) from e
main_inner(args)
finally:
pidfile.close()
def main_inner(args: "Args"):
soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
target = min(max(soft, 65535), hard)
resource.setrlimit(resource.RLIMIT_NOFILE, (target, hard))
@@ -286,20 +338,19 @@ def main():
# TODO: Refactor the current verbosity system
logger_setup(EXO_LOG, args.verbosity)
if args.no_stdio:
detach_stdio_to_devnull()
logger.info("Detached stdio to /dev/null")
logger.info(f"{'=' * 40}")
logger.info(f"Starting EXO | pid={os.getpid()}")
logger.info(f"{'=' * 40}")
logger.info(f"EXO_LIBP2P_NAMESPACE: {os.getenv('EXO_LIBP2P_NAMESPACE')}")
logger.info(f"pid = {os.getpid()}")
if os.getenv("EXO_LIBP2P_NAMESPACE"):
raise ValueError(
"EXO_LIBP2P_NAMESPACE has been removed - use EXO_ZENOH_NAMESPACE instead"
)
logger.info(f"EXO_ZENOH_NAMESPACE: {os.getenv('EXO_ZENOH_NAMESPACE')}")
if args.offline:
logger.info("Running in OFFLINE mode — no internet checks, local models only")
if args.bootstrap_peers:
logger.info(f"Bootstrap peers: {args.bootstrap_peers}")
raise ValueError("Bootstrap peers has been temporarily removed")
if args.no_batch:
os.environ["EXO_NO_BATCH"] = "1"
@@ -324,7 +375,6 @@ def main():
finally:
logger.info("EXO Shutdown complete")
logger_cleanup()
del pidfile
class Args(FrozenModel):
@@ -338,9 +388,11 @@ class Args(FrozenModel):
offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true"
no_batch: bool = False
fast_synch: bool | None = None # None = auto, True = force on, False = force off
no_stdio: bool = False
legacy_daemon: bool = False
bootstrap_peers: list[str] = []
libp2p_port: int
namespace: str
zenoh_port: int
discovery_port: int
@classmethod
def parse(cls) -> Self:
@@ -399,9 +451,9 @@ class Args(FrozenModel):
help="Disable continuous batching, use sequential generation",
)
parser.add_argument(
"--no-stdio",
"--legacy-daemon",
action="store_true",
help="Detach stdin/stdout/stderr to /dev/null after logging is configured",
help="Run as a legacy SysV-style background daemon using double-fork daemonization",
)
parser.add_argument(
"--bootstrap-peers",
@@ -413,11 +465,25 @@ class Args(FrozenModel):
help="Comma-separated libp2p multiaddrs to dial on startup (env: EXO_BOOTSTRAP_PEERS)",
)
parser.add_argument(
"--libp2p-port",
"--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,
dest="libp2p_port",
help="Fixed TCP port for libp2p to listen on (0 = OS-assigned).",
default=52414,
dest="zenoh_port",
help="Fixed TCP port for zenoh to listen.",
)
parser.add_argument(
"--discovery-port",
type=int,
default=52413,
dest="discovery_port",
help="Fixed UDP port for the discovery service.",
)
fast_synch_group = parser.add_mutually_exclusive_group()
fast_synch_group.add_argument(
+17 -4
View File
@@ -11,6 +11,10 @@ from exo.master.placement import (
place_instance,
)
from exo.master.placement_utils import find_ip_prioritised
from exo.routing.event_router import (
EventRouterBrokenResourceError,
EventRouterClosedResourceError,
)
from exo.shared.apply import apply
from exo.shared.constants import EXO_EVENT_LOG_DIR, EXO_TRACING_ENABLED
from exo.shared.types.commands import (
@@ -151,6 +155,9 @@ class Master:
tg.start_soon(self._event_processor)
tg.start_soon(self._command_processor)
tg.start_soon(self._plan)
except* (EventRouterBrokenResourceError, EventRouterClosedResourceError):
# Event router has been closed (try-star syntax handles error groups)
pass
finally:
self._event_log.close()
self.global_event_sender.close()
@@ -174,6 +181,7 @@ class Master:
case TestCommand():
pass
case TextGeneration():
# set-difference => prefill-only nodes
prefill_only: set[InstanceId] = set()
for link in self.state.instance_links.values():
prefill_only.update(link.prefill_instances)
@@ -181,11 +189,13 @@ class Master:
prefill_only.difference_update(link.decode_instances)
for instance in self.state.instances.values():
# NON-prefill-only instances matching the model ID
if (
instance.shard_assignments.model_id
== command.task_params.model
and instance.instance_id not in prefill_only
):
# count in-flight tasks of that instance
in_flight = {TaskStatus.Pending, TaskStatus.Running}
task_count = sum(
1
@@ -197,6 +207,7 @@ class Master:
task_count
)
# there are no NON-prefill-only instances matching this model ID
if not instance_task_counts:
raise ValueError(
f"No instance found for model {command.task_params.model}"
@@ -448,10 +459,12 @@ class Master:
self._event_log.read_range(command.since_idx, end),
start=command.since_idx,
):
await self._send_event(IndexedEvent(idx=i, event=event))
await self._send_indexed_event(
IndexedEvent(idx=i, event=event)
)
for event in generated_events:
await self.event_sender.send(event)
except ValueError as e:
except Exception 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
@@ -506,10 +519,10 @@ class Master:
self.state = apply(self.state, indexed)
self._event_log.append(event)
await self._send_event(indexed)
await self._send_indexed_event(indexed)
# This function is re-entrant, take care!
async def _send_event(self, event: IndexedEvent):
async def _send_indexed_event(self, event: IndexedEvent):
# Convenience method since this line is ugly
await self.global_event_sender.send(
GlobalForwarderEvent(
+3 -4
View File
@@ -6,7 +6,7 @@ import pytest
from loguru import logger
from exo.master.main import Master
from exo.routing.router import get_node_id_keypair
from exo.routing.router import get_node_zid
from exo.shared.models.model_cards import ModelCard, ModelTask
from exo.shared.types.backends import Backend
from exo.shared.types.commands import (
@@ -16,7 +16,7 @@ from exo.shared.types.commands import (
PlaceInstance,
TextGeneration,
)
from exo.shared.types.common import ModelId, NodeId, SessionId, SystemId
from exo.shared.types.common import ModelId, SessionId, SystemId
from exo.shared.types.events import (
Event,
GlobalForwarderEvent,
@@ -49,8 +49,7 @@ from exo.utils.info_gatherer.info_gatherer import NodeBackends
@pytest.mark.asyncio
async def test_master():
keypair = get_node_id_keypair()
node_id = NodeId(keypair.to_node_id())
node_id = get_node_zid()
session_id = SessionId(master_node_id=node_id, election_clock=0)
ge_sender, global_event_receiver = channel[GlobalForwarderEvent]()
+3 -5
View File
@@ -1,15 +1,13 @@
from exo_pyo3_bindings import PyFromSwarm
from exo_rs import FromSwarm
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)
def from_update(cls, update: FromSwarm.Connection) -> "ConnectionMessage":
return cls(connected=update.connected)
+21 -2
View File
@@ -15,11 +15,30 @@ from exo.shared.types.events import (
IndexedEvent,
LocalForwarderEvent,
)
from exo.utils import channels
from exo.utils.channels import Receiver, Sender, channel
from exo.utils.event_buffer import OrderedBuffer
from exo.utils.task_group import TaskGroup
class EventRouterClosedResourceError(ClosedResourceError):
pass
class EventRouterBrokenResourceError(BrokenResourceError):
pass
# Event Router is created and destroyed before consumers of its channels are,
# hence its nice to have tagged errors for event-router channels being closed
#
# so consumers can catch specifically these errors, rather than the generic ones
_ERROR_CFG = channels.ErrorOverride(
closed_resource_error=EventRouterClosedResourceError,
broken_resource_error=EventRouterBrokenResourceError,
)
@dataclass
class EventRouter:
session_id: SessionId
@@ -64,7 +83,7 @@ class EventRouter:
await self.external_outbound.send(event)
def sender(self) -> Sender[Event]:
send, recv = channel[Event]()
send, recv = channel[Event](error_override_config=_ERROR_CFG)
if self._tg.is_running():
self._tg.start_soon(self._ingest, SystemId(), recv)
else:
@@ -73,7 +92,7 @@ class EventRouter:
def receiver(self) -> Receiver[IndexedEvent]:
assert not self._tg.is_running()
send, recv = channel[IndexedEvent]()
send, recv = channel[IndexedEvent](error_override_config=_ERROR_CFG)
self.internal_outbound.append(send)
return recv
+25 -36
View File
@@ -1,8 +1,7 @@
from collections.abc import Sequence
import os
from copy import copy
from itertools import count
from math import inf
from os import PathLike
from pathlib import Path
from typing import cast
@@ -12,18 +11,14 @@ from anyio import (
move_on_after,
sleep_forever,
)
from exo_pyo3_bindings import (
AllQueuesFullError,
Keypair,
MessageTooLargeError,
from exo_rs import (
FromSwarm,
NetworkingHandle,
NoPeersSubscribedToTopicError,
PyFromSwarm,
)
from filelock import FileLock
from loguru import logger
from exo.shared.constants import EXO_NODE_ID_KEYPAIR
from exo.shared.constants import EXO_NODE_ZID
from exo.shared.types.common import NodeId
from exo.utils.channels import Receiver, Sender, channel
from exo.utils.pydantic_ext import FrozenModel
from exo.utils.task_group import TaskGroup
@@ -105,12 +100,15 @@ class Router:
@classmethod
def create(
cls,
identity: Keypair,
bootstrap_peers: Sequence[str] = (),
listen_port: int = 0,
identity: str,
namespace: str,
listen_port: int,
discovery_service_port: int,
) -> "Router":
return cls(
handle=NetworkingHandle(identity, list(bootstrap_peers), listen_port)
handle=NetworkingHandle.new(
identity, namespace, listen_port, discovery_service_port
)
)
def __init__(self, handle: NetworkingHandle):
@@ -191,10 +189,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 FromSwarm.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}"
@@ -202,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}"
@@ -225,33 +221,25 @@ 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"
)
await self._net.gossipsub_publish(topic, data)
def get_node_id_keypair(
path: str | bytes | PathLike[str] | PathLike[bytes] = EXO_NODE_ID_KEYPAIR,
) -> Keypair:
def get_node_zid(
path: Path = EXO_NODE_ZID,
) -> NodeId:
"""
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()
return NodeId(os.urandom(16).hex().lstrip("0"))
"""
def lock_path(path: str | bytes | PathLike[str] | PathLike[bytes]) -> Path:
return Path(str(path) + ".lock")
@@ -273,3 +261,4 @@ def get_node_id_keypair(
keypair = Keypair.generate()
f.write(keypair.to_bytes())
return keypair
"""
+1 -1
View File
@@ -76,7 +76,7 @@ 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"
EXO_NODE_ZID = EXO_CACHE_HOME / "node_zid"
EXO_CONFIG_FILE = EXO_CONFIG_HOME / "config.toml"
# libp2p topics for event forwarding
+2 -1
View File
@@ -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_rs").setLevel(logging.INFO)
logging.getLogger("networking").setLevel(logging.INFO)
logging.getLogger("httpx").setLevel(logging.WARNING)
logging.getLogger("httpcore").setLevel(logging.WARNING)
+9 -2
View File
@@ -81,8 +81,10 @@ class _CardCache:
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
except (ValidationError, TOMLKitError) as e:
logger.opt(exception=e).warning(
f"failed to validate model card at {toml_file}"
)
async def refresh(self) -> None:
for path in _BUILTIN_CARD_DIRS:
@@ -197,6 +199,11 @@ class ModelCard(FrozenModel):
def _validate_tasks(cls, v: list[str | ModelTask]) -> list[ModelTask]:
return [item if isinstance(item, ModelTask) else ModelTask(item) for item in v]
@field_validator("backends", mode="before")
@classmethod
def _validate_backends(cls, v: list[str | Backend]) -> list[Backend]:
return [item if isinstance(item, Backend) else Backend(item) for item in v]
async def save(self, path: Path) -> None:
async with await open_file(path, "w") as f:
py = self.model_dump(exclude_none=True, exclude={"is_custom"})
+1 -1
View File
@@ -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:
@@ -10,8 +10,8 @@ from multiprocessing.synchronize import Semaphore as SemaphoreT
from loguru import logger
from pytest import LogCaptureFixture, mark
from exo.routing.router import get_node_id_keypair
from exo.shared.constants import EXO_NODE_ID_KEYPAIR
from exo.routing.router import get_node_zid
from exo.shared.constants import EXO_NODE_ZID
NUM_CONCURRENT_PROCS = 10
@@ -23,7 +23,7 @@ def _get_keypair_concurrent_subprocess_task(
sem.release()
# wait to be told to begin simultaneous read
ev.wait()
queue.put(get_node_id_keypair().to_bytes())
queue.put(get_node_zid().encode())
def _get_keypair_concurrent(num_procs: int) -> bytes:
@@ -79,7 +79,7 @@ def test_node_id_fetching(caplog: LogCaptureFixture):
reps = 10
# delete current file and write a new one
_delete_if_exists(EXO_NODE_ID_KEYPAIR)
_delete_if_exists(EXO_NODE_ZID)
kp = _get_keypair_concurrent(NUM_CONCURRENT_PROCS)
with caplog.at_level(101): # supress logs
@@ -88,6 +88,6 @@ def test_node_id_fetching(caplog: LogCaptureFixture):
assert kp == _get_keypair_concurrent(NUM_CONCURRENT_PROCS)
# make sure that after deleting, we are not fetching the same value
_delete_if_exists(EXO_NODE_ID_KEYPAIR)
_delete_if_exists(EXO_NODE_ZID)
for _ in range(reps):
assert kp != _get_keypair_concurrent(NUM_CONCURRENT_PROCS)
-7
View File
@@ -97,13 +97,6 @@ def test_macos_uses_traditional_paths():
assert home / ".exo" == constants.EXO_CACHE_HOME
def test_node_id_in_config_dir():
"""Test that node ID keypair is in the config directory."""
import exo.shared.constants as constants
assert constants.EXO_NODE_ID_KEYPAIR.parent == constants.EXO_CONFIG_HOME
def test_models_in_data_dir():
"""Test that default models directory is in the data directory."""
# Clear EXO_MODELS_DIRS to test default behavior
-1
View File
@@ -13,7 +13,6 @@ from exo.shared.models.model_cards import ModelId
from exo.utils.pydantic_ext import TaggedModel
from exo.worker.runner.diagnostics import KnownRunnerDiagnostic
from ...worker.runner.diagnostics import KnownRunnerDiagnostic
from .common import CommandId
+5
View File
@@ -2,6 +2,11 @@ from typing import Any, Type
from .phantom import PhantomData
STDIN_FD = 0
STDOUT_FD = 1
STDERR_FD = 2
STDIO_FDS = (STDIN_FD, STDOUT_FD, STDERR_FD)
def ensure_type[T](obj: Any, expected_type: Type[T]) -> T: # type: ignore
if not isinstance(obj, expected_type):
+4 -5
View File
@@ -25,10 +25,9 @@ from anyio import (
from anyio.abc import TaskStatus
from loguru import logger
from exo.utils import STDERR_FD, STDIO_FDS, STDOUT_FD
from exo.utils.channels import Receiver, Sender, channel
_STDOUT_FD = 1
_STDERR_FD = 2
_READ_CHUNK_SIZE = 64 * 1024
_JOIN_GRACE_SECONDS = 3.0
_TERMINATE_GRACE_SECONDS = 5.0
@@ -256,11 +255,11 @@ def _run_with_captured_stdio(
stderr_fd = stderr.detach()
try:
os.dup2(stdout_fd, _STDOUT_FD)
os.dup2(stderr_fd, _STDERR_FD)
os.dup2(stdout_fd, STDOUT_FD)
os.dup2(stderr_fd, STDERR_FD)
finally:
for fd in (stdout_fd, stderr_fd):
if fd not in (_STDOUT_FD, _STDERR_FD):
if fd not in STDIO_FDS:
_close_fd(fd)
faulthandler.enable(file=sys.stderr, all_threads=True)
+1 -1
View File
@@ -38,7 +38,7 @@ def print_startup_banner(port: int) -> None:
🌐 Dashboard & API Ready
Dashboard & API Ready
{dashboard_url}{" " * (69 - len(dashboard_url))}
+155 -8
View File
@@ -1,13 +1,16 @@
import contextlib
import multiprocessing as mp
from dataclasses import dataclass, field
from functools import wraps
from inspect import iscoroutinefunction
from math import inf
from multiprocessing.synchronize import Event
from queue import Empty, Full
from types import TracebackType
from typing import Any, Self
from types import CoroutineType, TracebackType
from typing import Any, Callable, NoReturn, Self, cast, overload, override
from anyio import (
BrokenResourceError,
CapacityLimiter,
ClosedResourceError,
EndOfStream,
@@ -20,35 +23,172 @@ from anyio.streams.memory import (
from anyio.streams.memory import (
MemoryObjectSendStream as AnyioSender,
)
from anyio.streams.memory import (
MemoryObjectStreamState,
)
from anyio.streams.memory import (
MemoryObjectStreamState as AnyioState,
)
@dataclass(eq=False)
class ErrorOverride:
closed_resource_error: type[ClosedResourceError] = field(
default=ClosedResourceError,
)
broken_resource_error: type[BrokenResourceError] = field(
default=BrokenResourceError,
)
end_of_stream: type[EndOfStream] = field(
default=EndOfStream,
)
would_block: type[WouldBlock] = field(
default=WouldBlock,
)
@overload
def patch[**P, R](
self,
fn: Callable[P, CoroutineType[Any, Any, R]],
/,
) -> Callable[P, CoroutineType[Any, Any, R]]: ...
@overload
def patch[**P, R](
self,
fn: Callable[P, R],
/,
) -> Callable[P, R]: ...
def patch[**P, R](self, fn: Callable[P, Any], /) -> Callable[P, Any]:
"""
Returns a function with all these exceptions replaced by their overrides
"""
if iscoroutinefunction(fn):
async_fn = cast(Callable[P, CoroutineType[Any, Any, R]], fn)
@wraps(async_fn)
async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
try:
return await async_fn(*args, **kwargs)
except ClosedResourceError as e:
self._raise_replace(self.closed_resource_error, e)
except BrokenResourceError as e:
self._raise_replace(self.broken_resource_error, e)
except EndOfStream as e:
self._raise_replace(self.end_of_stream, e)
except WouldBlock as e:
self._raise_replace(self.would_block, e)
return async_wrapper
else:
sync_fn = cast(Callable[P, R], fn)
@wraps(sync_fn)
def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
try:
return sync_fn(*args, **kwargs)
except ClosedResourceError as e:
self._raise_replace(self.closed_resource_error, e)
except BrokenResourceError as e:
self._raise_replace(self.broken_resource_error, e)
except EndOfStream as e:
self._raise_replace(self.end_of_stream, e)
except WouldBlock as e:
self._raise_replace(self.would_block, e)
return sync_wrapper
@staticmethod
def _raise_replace(replacement: type[BaseException], e: BaseException) -> NoReturn:
if isinstance(e, replacement):
raise
raise replacement() from e
class Sender[T](AnyioSender[T]):
def __init__(
self,
state: MemoryObjectStreamState[T],
error_override_config: ErrorOverride | None,
):
super().__init__(_state=state)
# patch the methods we want to override errors for
#
# NOTE: it is very important that new methods which are added,
# and which can throw, are patched in this block
if (e := error_override_config) is not None:
# new methods of this class
self.clone_receiver = e.patch(self.clone_receiver)
# overridden methods
self.clone = e.patch(self.clone)
# parent methods
self.send_nowait = e.patch(self.send_nowait)
self.send = e.patch(self.send)
self.close = e.patch(self.close)
self.aclose = e.patch(self.aclose)
self.statistics = e.patch(self.statistics)
self.err_config = error_override_config
@override
def clone(self) -> "Sender[T]":
if self._closed:
raise ClosedResourceError
return Sender(_state=self._state)
return Sender(self._state, self.err_config)
def clone_receiver(self) -> "Receiver[T]":
"""Constructs a Receiver using a Senders shared state - similar to calling Receiver.clone() without needing the receiver"""
if self._closed:
raise ClosedResourceError
return Receiver(_state=self._state)
return Receiver(self._state, self.err_config)
class Receiver[T](AnyioReceiver[T]):
def __init__(
self,
state: MemoryObjectStreamState[T],
error_override_config: ErrorOverride | None,
):
super().__init__(_state=state)
# patch the methods we want to override errors for
#
# NOTE: it is very important that new methods which are added,
# and which can throw, are patched in this block
if (e := error_override_config) is not None:
# new methods of this class
self.clone_sender = e.patch(self.clone_sender)
self.collect = e.patch(self.collect)
self.receive_at_least = e.patch(self.receive_at_least)
# overridden methods
self.clone = e.patch(self.clone)
# parent methods
self.receive_nowait = e.patch(self.receive_nowait)
self.receive = e.patch(self.receive)
self.close = e.patch(self.close)
self.aclose = e.patch(self.aclose)
self.statistics = e.patch(self.statistics)
self.err_config = error_override_config
@override
def clone(self) -> "Receiver[T]":
if self._closed:
raise ClosedResourceError
return Receiver(_state=self._state)
return Receiver(self._state, self.err_config)
def clone_sender(self) -> Sender[T]:
"""Constructs a Sender using a Receivers shared state - similar to calling Sender.clone() without needing the sender"""
if self._closed:
raise ClosedResourceError
return Sender(_state=self._state)
return Sender(self._state, self.err_config)
def collect(self) -> list[T]:
"""Collect all currently available items from this receiver"""
@@ -70,6 +210,7 @@ class Receiver[T](AnyioReceiver[T]):
out.extend(self.collect())
return out
@override
def __enter__(self) -> Self:
return self
@@ -285,11 +426,17 @@ class MpReceiver[T]:
class channel[T]: # noqa: N801
"""Create a pair of asynchronous channels for communicating within the same process"""
def __new__(cls, max_buffer_size: float = inf) -> tuple[Sender[T], Receiver[T]]:
def __new__(
cls,
max_buffer_size: float = inf,
error_override_config: ErrorOverride | None = None,
) -> tuple[Sender[T], Receiver[T]]:
if max_buffer_size != inf and not isinstance(max_buffer_size, int):
raise ValueError("max_buffer_size must be either an integer or math.inf")
state = AnyioState[T](max_buffer_size)
return Sender(_state=state), Receiver(_state=state)
return Sender(state, error_override_config), Receiver(
state, error_override_config
)
class mp_channel[T]: # noqa: N801
-28
View File
@@ -1,28 +0,0 @@
import os
import sys
_STDIN_FD = 0
_STDOUT_FD = 1
_STDERR_FD = 2
def detach_stdio_to_devnull() -> None:
"""Redirect process stdio file descriptors to /dev/null."""
for stream in (sys.stdout, sys.stderr, sys.__stdout__, sys.__stderr__):
if stream is not None:
stream.flush()
stdin_fd = os.open(os.devnull, os.O_RDONLY)
stdout_fd = os.open(os.devnull, os.O_WRONLY)
stderr_fd = os.open(os.devnull, os.O_WRONLY)
try:
# dup2 closes the target fd first, but leaves the source fd open.
os.dup2(stdin_fd, _STDIN_FD)
os.dup2(stdout_fd, _STDOUT_FD)
os.dup2(stderr_fd, _STDERR_FD)
finally:
for fd in (stdin_fd, stdout_fd, stderr_fd):
if fd not in (_STDIN_FD, _STDOUT_FD, _STDERR_FD):
os.close(fd)
@@ -630,6 +630,14 @@ class InfoGatherer:
f"MacMon failed with return code {e.returncode}: {stderr_msg}"
)
self._tg.start_soon(self._monitor_memory_usage, 1)
except ProcessLookupError:
# usually throws by the process' context manager on exit
# when we ctrl+c, hence usually should be ignored;
# if anything else throws it, we explicitly don't care:
# process is dead anyways ;)
logger.warning(
"Macmon process not found - shutting down macmon monitor"
)
except Exception as e:
logger.opt(exception=e).warning("Error in macmon monitor")
self._tg.start_soon(self._monitor_memory_usage, 1)
-41
View File
@@ -1,6 +1,3 @@
import os
import shutil
import subprocess
from typing import Self
from pydantic import BaseModel
@@ -71,41 +68,3 @@ class MacmonMetrics(TaggedModel):
@classmethod
def from_raw_json(cls, json: str) -> Self:
return cls.from_raw(RawMacmonMetrics.model_validate_json(json))
def read_macmon_metrics_once(
macmon_path: str | None = None,
*,
timeout: float = 5,
) -> MacmonMetrics | None:
"""
Read a single macmon sample, returning None when macmon is unavailable.
"""
resolved_macmon_path = (
macmon_path or os.getenv("EXO_MACMON_PATH") or shutil.which("macmon")
)
if resolved_macmon_path is None:
return None
try:
result = subprocess.run(
[resolved_macmon_path, "pipe", "--samples", "1", "--interval", "100"],
capture_output=True,
check=False,
text=True,
timeout=timeout,
)
except (OSError, subprocess.SubprocessError):
return None
if result.returncode != 0:
return None
lines = result.stdout.strip().splitlines()
if not lines:
return None
try:
return MacmonMetrics.from_raw_json(lines[0])
except ValueError:
return None
-28
View File
@@ -1,28 +0,0 @@
from __future__ import annotations
import os
from typing import Final
from exo_pyo3_bindings import Pidfile, PidfileError
from exo.shared.constants import EXO_PID_FILE
_PIDFILE_MODE: Final = 0o600
class PidfileLockError(RuntimeError):
pass
def acquire_exo_pidfile() -> Pidfile:
path = EXO_PID_FILE
os.makedirs(os.path.dirname(path), exist_ok=True)
try:
pidfile = Pidfile(path, _PIDFILE_MODE)
pidfile.write()
except (OSError, PidfileError) as exception:
raise PidfileLockError(
f"Failed to acquire EXO pidfile at {path}: {exception}"
) from exception
return pidfile
+122
View File
@@ -24,6 +24,7 @@ class PowerSampler:
] = defaultdict(list)
self._start_time: float | None = None
self._stopped = False
self._prefill_done_at: float | None = None
def _take_sample(self, t_rel: float | None = None) -> None:
assert self._start_time is not None
@@ -38,14 +39,35 @@ class PowerSampler:
await anyio.sleep(self._interval)
self._take_sample()
def mark_prefill_done(self) -> None:
"""Anchor the prefill→generation boundary on a fresh sample.
Idempotent. Safe to call before `run()`; boundary then lands at t=0.
"""
if self._prefill_done_at is not None:
return
if self._start_time is None:
self._prefill_done_at = 0.0
return
t_rel = time.perf_counter() - self._start_time
self._take_sample(t_rel=t_rel)
self._prefill_done_at = t_rel
def result(self) -> PowerUsage:
self._stopped = True
assert self._start_time is not None, "result() called before run()"
elapsed = time.perf_counter() - self._start_time
self._take_sample(t_rel=elapsed)
# Clamp the split point to [0, elapsed] in case timing is weird (e.g.
# mark called after result, or sampler ran for < the prefill window).
split = self._prefill_done_at
if split is not None:
split = max(0.0, min(elapsed, split))
node_stats: list[NodePowerStats] = []
total_energy_j = 0.0
total_prefill_energy_j = 0.0
total_generation_energy_j = 0.0
for node_id, ts_profiles in self._samples.items():
n = len(ts_profiles)
if n == 0:
@@ -53,20 +75,68 @@ class PowerSampler:
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
prefill_e: float | None = None
generation_e: float | None = None
prefill_avg: float | None = None
generation_avg: float | None = None
if split is not None:
prefill_e = trapezoidal_energy_range(ts_profiles, 0.0, split)
generation_e = trapezoidal_energy_range(ts_profiles, split, elapsed)
total_prefill_energy_j += prefill_e
total_generation_energy_j += generation_e
prefill_dt = split
generation_dt = elapsed - split
prefill_avg = prefill_e / prefill_dt if prefill_dt > 0 else 0.0
generation_avg = (
generation_e / generation_dt if generation_dt > 0 else 0.0
)
node_stats.append(
NodePowerStats(
node_id=node_id,
samples=n,
avg_sys_power=avg_power_w,
prefill_avg_sys_power=prefill_avg,
generation_avg_sys_power=generation_avg,
prefill_energy_joules=prefill_e,
generation_energy_joules=generation_e,
)
)
total_avg_sys_w = total_energy_j / elapsed if elapsed > 0 else 0.0
prefill_seconds: float | None = None
generation_seconds: float | None = None
prefill_energy_joules: float | None = None
generation_energy_joules: float | None = None
prefill_avg_w: float | None = None
generation_avg_w: float | None = None
if split is not None:
prefill_seconds = split
generation_seconds = elapsed - split
prefill_energy_joules = total_prefill_energy_j
generation_energy_joules = total_generation_energy_j
prefill_avg_w = (
total_prefill_energy_j / prefill_seconds if prefill_seconds > 0 else 0.0
)
generation_avg_w = (
total_generation_energy_j / generation_seconds
if generation_seconds > 0
else 0.0
)
return PowerUsage(
elapsed_seconds=elapsed,
nodes=node_stats,
total_avg_sys_power_watts=total_avg_sys_w,
total_energy_joules=total_energy_j,
prefill_seconds=prefill_seconds,
generation_seconds=generation_seconds,
prefill_energy_joules=prefill_energy_joules,
generation_energy_joules=generation_energy_joules,
prefill_avg_sys_power_watts=prefill_avg_w,
generation_avg_sys_power_watts=generation_avg_w,
)
@@ -89,3 +159,55 @@ def trapezoidal_energy(
continue
energy_j += (p_prev.sys_power + p_cur.sys_power) / 2.0 * dt
return energy_j
def trapezoidal_energy_range(
ts_profiles: list[tuple[float, SystemPerformanceProfile]],
t_start: float,
t_end: float,
) -> float:
"""Integrate sys_power(t) over [t_start, t_end] using the trapezoidal rule.
Linearly interpolates power at the endpoints when they fall between
existing samples, so callers can integrate over arbitrary sub-windows
(e.g. the prefill segment) without losing accuracy. Returns 0 for an
empty or zero-length window. Falls back to constant-power assumption
when only one sample exists.
"""
if t_end <= t_start:
return 0.0
if len(ts_profiles) == 0:
return 0.0
if len(ts_profiles) == 1:
return ts_profiles[0][1].sys_power * (t_end - t_start)
def power_at(t: float) -> float:
if t <= ts_profiles[0][0]:
return ts_profiles[0][1].sys_power
if t >= ts_profiles[-1][0]:
return ts_profiles[-1][1].sys_power
for i in range(1, len(ts_profiles)):
t_cur, p_cur = ts_profiles[i]
if t_cur >= t:
t_prev, p_prev = ts_profiles[i - 1]
span = t_cur - t_prev
if span <= 0:
return p_cur.sys_power
frac = (t - t_prev) / span
return p_prev.sys_power + frac * (p_cur.sys_power - p_prev.sys_power)
return ts_profiles[-1][1].sys_power
p_start = power_at(t_start)
p_end = power_at(t_end)
in_range: list[tuple[float, float]] = [
(t, profile.sys_power) for t, profile in ts_profiles if t_start < t < t_end
]
seq: list[tuple[float, float]] = [(t_start, p_start)] + in_range + [(t_end, p_end)]
energy_j = 0.0
for i in range(1, len(seq)):
dt = seq[i][0] - seq[i - 1][0]
if dt <= 0:
continue
energy_j += (seq[i - 1][1] + seq[i][1]) / 2.0 * dt
return energy_j
+121
View File
@@ -0,0 +1,121 @@
import multiprocessing as mp
import time
import pytest
from anyio import (
BrokenResourceError,
ClosedResourceError,
EndOfStream,
WouldBlock,
fail_after,
)
from loguru import logger
from exo.utils.channels import ErrorOverride, MpReceiver, MpSender, channel, mp_channel
class CustomClosedResourceError(ClosedResourceError):
pass
class CustomBrokenResourceError(BrokenResourceError):
pass
class CustomEndOfStream(EndOfStream):
pass
class CustomWouldBlock(WouldBlock):
pass
ERROR_OVERRIDE = ErrorOverride(
closed_resource_error=CustomClosedResourceError,
broken_resource_error=CustomBrokenResourceError,
end_of_stream=CustomEndOfStream,
would_block=CustomWouldBlock,
)
def foo(recv: MpReceiver[str]):
expected = ["hi", "hi 2", "bye"]
with recv as r:
for item in r:
assert item == expected.pop(0)
def bar(send: MpSender[str]):
logger.warning("hi")
send.send("hi")
time.sleep(0.1)
logger.warning("hi 2")
send.send("hi 2")
time.sleep(0.1)
logger.warning("bye")
send.send("bye")
time.sleep(0.1)
send.close()
@pytest.mark.anyio
async def test_channel_ipc():
with fail_after(0.5):
s, r = mp_channel[str]()
p1 = mp.Process(target=foo, args=(r,))
p2 = mp.Process(target=bar, args=(s,))
p1.start()
p2.start()
p1.join()
p2.join()
def test_channel_error_override_replaces_sync_errors_with_subclasses():
send, recv = channel[int](0, error_override_config=ERROR_OVERRIDE)
with pytest.raises(CustomWouldBlock) as would_block_info:
send.send_nowait(1)
assert type(would_block_info.value.__cause__) is WouldBlock
recv.close()
with pytest.raises(CustomBrokenResourceError) as broken_resource_info:
send.send_nowait(1)
assert type(broken_resource_info.value.__cause__) is BrokenResourceError
send.close()
with pytest.raises(CustomClosedResourceError) as closed_resource_info:
send.send_nowait(1)
assert type(closed_resource_info.value.__cause__) is ClosedResourceError
@pytest.mark.anyio
async def test_channel_error_override_replaces_async_errors_with_subclasses():
send, recv = channel[int](0, error_override_config=ERROR_OVERRIDE)
recv.close()
with pytest.raises(CustomBrokenResourceError) as broken_resource_info:
await send.send(1)
assert type(broken_resource_info.value.__cause__) is BrokenResourceError
send, recv = channel[int](error_override_config=ERROR_OVERRIDE)
send.close()
with pytest.raises(CustomEndOfStream) as end_of_stream_info:
await recv.receive()
assert type(end_of_stream_info.value.__cause__) is EndOfStream
@pytest.mark.anyio
async def test_channel_error_override_is_preserved_by_clones():
send, recv = channel[int](0, error_override_config=ERROR_OVERRIDE)
send_clone = send.clone()
recv.close()
with pytest.raises(CustomBrokenResourceError):
await send_clone.send(1)
send, recv = channel[int](0, error_override_config=ERROR_OVERRIDE)
cloned_send = recv.clone_sender()
recv.close()
with pytest.raises(CustomBrokenResourceError):
await cloned_send.send(1)
-168
View File
@@ -1,168 +0,0 @@
import contextlib
import os
from collections.abc import AsyncIterator
import anyio
import pytest
from anyio import EndOfStream, create_task_group, fail_after
from exo.utils.async_process import AsyncProcess
from exo.utils.channels import MpReceiver, MpSender, Receiver, mp_channel
from exo.utils.daemon import detach_stdio_to_devnull
def _write_before_and_after_detach() -> None:
os.write(1, b"before stdout\n")
os.write(2, b"before stderr\n")
detach_stdio_to_devnull()
os.write(1, b"after stdout\n")
os.write(2, b"after stderr\n")
def _write_grandchild_stdio(label: str) -> None:
os.write(1, f"{label} stdout\n".encode())
os.write(2, f"{label} stderr\n".encode())
async def _spawn_grandchild_and_report(
result_sender: MpSender[tuple[int, bytes, bytes]],
label: str,
) -> None:
result_sender.send(await _collect_spawned_child(label))
result_sender.close()
async def _collect_spawned_child(label: str) -> tuple[int, bytes, bytes]:
process = AsyncProcess(_write_grandchild_stdio, args=(label,))
async with _started_process(process):
return await _collect_process_output(process)
def _detach_stdio_then_spawn_captured_child(
result_sender: MpSender[tuple[int, bytes, bytes]],
) -> None:
detach_stdio_to_devnull()
anyio.run(_spawn_grandchild_and_report, result_sender, "grandchild")
def _detach_stdio_then_spawn_captured_children_sequentially(
result_sender: MpSender[list[tuple[int, bytes, bytes]]],
) -> None:
async def run_children() -> list[tuple[int, bytes, bytes]]:
results: list[tuple[int, bytes, bytes]] = []
for index in range(5):
results.append(await _collect_spawned_child(f"grandchild-{index}"))
return results
detach_stdio_to_devnull()
result_sender.send(anyio.run(run_children))
result_sender.close()
async def _collect_stream(stream: Receiver[bytes], output: bytearray) -> None:
while True:
try:
output.extend(await stream.receive())
except EndOfStream:
return
async def _collect_process_output(
process: AsyncProcess,
) -> tuple[int, bytes, bytes]:
stdout = bytearray()
stderr = bytearray()
exitcodes: list[int] = []
async with create_task_group() as collect_group:
collect_group.start_soon(_collect_stream, process.stdout, stdout)
collect_group.start_soon(_collect_stream, process.stderr, stderr)
exitcodes.append(await process.wait())
if not exitcodes:
raise RuntimeError("process exited without a return code")
return exitcodes[0], bytes(stdout), bytes(stderr)
@contextlib.asynccontextmanager
async def _started_process(process: AsyncProcess) -> AsyncIterator[None]:
async with create_task_group() as task_group:
await task_group.start(process.run)
try:
yield
finally:
await process.stop()
async def _run_process_and_receive[T](
process: AsyncProcess,
recv: MpReceiver[T],
*,
timeout: float,
) -> tuple[int, T]:
async with _started_process(process):
with fail_after(timeout):
result = await recv.receive_async()
exitcode = await process.wait()
return exitcode, result
@pytest.mark.anyio
async def test_detach_stdio_to_devnull_redirects_stdio_away_from_capture() -> None:
process = AsyncProcess(_write_before_and_after_detach)
async with _started_process(process):
exitcode, stdout, stderr = await _collect_process_output(process)
assert exitcode == 0
assert stdout == b"before stdout\n"
assert stderr == b"before stderr\n"
@pytest.mark.anyio
async def test_detached_stdio_process_can_spawn_and_capture_child_stdio() -> None:
send, recv = mp_channel[tuple[int, bytes, bytes]]()
process = AsyncProcess(_detach_stdio_then_spawn_captured_child, args=(send,))
try:
daemonized_parent_exitcode, result = await _run_process_and_receive(
process, recv, timeout=5
)
finally:
recv.close()
child_exitcode, child_stdout, child_stderr = result
assert daemonized_parent_exitcode == 0
assert child_exitcode == 0
assert child_stdout == b"grandchild stdout\n"
assert child_stderr == b"grandchild stderr\n"
@pytest.mark.anyio
async def test_detached_stdio_process_can_spawn_captured_children_sequentially() -> (
None
):
send, recv = mp_channel[list[tuple[int, bytes, bytes]]]()
process = AsyncProcess(
_detach_stdio_then_spawn_captured_children_sequentially,
args=(send,),
)
try:
daemonized_parent_exitcode, results = await _run_process_and_receive(
process, recv, timeout=10
)
finally:
recv.close()
assert daemonized_parent_exitcode == 0
assert results == [
(
0,
f"grandchild-{index} stdout\n".encode(),
f"grandchild-{index} stderr\n".encode(),
)
for index in range(5)
]
-40
View File
@@ -1,40 +0,0 @@
import multiprocessing as mp
import time
import pytest
from anyio import fail_after
from loguru import logger
from exo.utils.channels import MpReceiver, MpSender, mp_channel
def foo(recv: MpReceiver[str]):
expected = ["hi", "hi 2", "bye"]
with recv as r:
for item in r:
assert item == expected.pop(0)
def bar(send: MpSender[str]):
logger.warning("hi")
send.send("hi")
time.sleep(0.1)
logger.warning("hi 2")
send.send("hi 2")
time.sleep(0.1)
logger.warning("bye")
send.send("bye")
time.sleep(0.1)
send.close()
@pytest.mark.anyio
async def test_channel_ipc():
with fail_after(0.5):
s, r = mp_channel[str]()
p1 = mp.Process(target=foo, args=(r,))
p2 = mp.Process(target=bar, args=(s,))
p1.start()
p2.start()
p1.join()
p2.join()
+14 -24
View File
@@ -8,36 +8,28 @@ import textwrap
from pathlib import Path
from typing import Final
import pytest
import exo.utils.pidfile as pidfile
from exo.utils.pidfile import acquire_exo_pidfile
from exo_rs import 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
from exo_rs import Pidfile, PidfileError
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
path = Path(sys.argv[1])
try:
handle = Pidfile(path, 0o0600)
handle.write()
except (OSError, PidfileError) as exception:
print(f"Failed to acquire EXO pidfile at {path}: {exception}")
raise SystemExit(73) from exception
del handle
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)],
@@ -49,12 +41,11 @@ def _run_child_acquire_pidfile(path: Path) -> subprocess.CompletedProcess[str]:
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()
handle = Pidfile(path, 0o0600)
handle.write()
assert path.read_text() == str(os.getpid())
del handle
@@ -65,12 +56,11 @@ def test_acquire_exo_pidfile_writes_current_pid_and_removes_on_drop(
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()
handle = Pidfile(path, 0o0600)
handle.write()
try:
blocked_child = _run_child_acquire_pidfile(path)
assert blocked_child.returncode == 73
+136
View File
@@ -141,6 +141,142 @@ def test_trapezoidal_unit_single_sample() -> None:
assert trapezoidal_energy(samples, elapsed=3.0) == 42.0 * 3.0
def test_trapezoidal_range_interpolation() -> None:
"""Sub-window integration should linearly interpolate at the boundaries."""
from exo.utils.power_sampler import trapezoidal_energy_range
# Two samples: t=0 W=10, t=10 W=20 -> power(t) = 10 + t
samples = [
(0.0, _make_profile(10.0)),
(10.0, _make_profile(20.0)),
]
# Integral from t=4 to t=6: power goes 14 -> 16, mean 15, dt=2 -> 30 J
assert abs(trapezoidal_energy_range(samples, 4.0, 6.0) - 30.0) < 1e-9
# Integral over the full window matches the full trapezoidal integral.
full = trapezoidal_energy_range(samples, 0.0, 10.0)
assert abs(full - 150.0) < 1e-9
def test_trapezoidal_range_zero_window() -> None:
"""Zero-length or reversed windows integrate to zero."""
from exo.utils.power_sampler import trapezoidal_energy_range
samples = [(0.0, _make_profile(10.0)), (5.0, _make_profile(20.0))]
assert trapezoidal_energy_range(samples, 3.0, 3.0) == 0.0
assert trapezoidal_energy_range(samples, 5.0, 3.0) == 0.0
def test_trapezoidal_range_splits_sum_to_full() -> None:
"""Energy split at an arbitrary boundary should sum back to the full integral."""
from exo.utils.power_sampler import (
trapezoidal_energy,
trapezoidal_energy_range,
)
samples = [
(0.0, _make_profile(10.0)),
(1.0, _make_profile(20.0)),
(3.0, _make_profile(15.0)),
(5.0, _make_profile(25.0)),
]
full = trapezoidal_energy(samples, elapsed=5.0)
# Split at t=2.5 (between samples) — interpolation should be exact.
left = trapezoidal_energy_range(samples, 0.0, 2.5)
right = trapezoidal_energy_range(samples, 2.5, 5.0)
assert abs((left + right) - full) < 1e-9
async def test_prefill_generation_split() -> None:
"""When mark_prefill_done() is called, the result should split energy."""
state: dict[NodeId, SystemPerformanceProfile] = {
NODE_A: _make_profile(10.0),
}
sampler = PowerSampler(get_node_system=lambda: state, interval=0.02)
async with anyio.create_task_group() as tg:
tg.start_soon(sampler.run)
# "Prefill" phase: power = 10 W
await anyio.sleep(0.1)
# Mark the boundary BEFORE changing state — this matches what
# _collect_text_generation_with_stats does in production: the mark
# fires on the first non-prefill chunk, so the boundary sample is
# the genuine end-of-prefill reading rather than the new phase's.
sampler.mark_prefill_done()
state[NODE_A] = _make_profile(30.0)
# "Generation" phase: power = 30 W
await anyio.sleep(0.1)
tg.cancel_scope.cancel()
result = sampler.result()
assert result.prefill_seconds is not None
assert result.generation_seconds is not None
assert result.prefill_energy_joules is not None
assert result.generation_energy_joules is not None
assert result.prefill_avg_sys_power_watts is not None
assert result.generation_avg_sys_power_watts is not None
# Phase durations should sum to the elapsed seconds.
assert (
abs(
(result.prefill_seconds + result.generation_seconds)
- result.elapsed_seconds
)
< 1e-6
)
# Phase energies should sum to (approximately) the total.
assert (
abs(
(result.prefill_energy_joules + result.generation_energy_joules)
- result.total_energy_joules
)
< 1e-6
)
# With the boundary sample anchored at the genuine end-of-prefill (10 W),
# prefill avg should converge tightly on 10 W and generation on 30 W.
# 15 W cleanly separates the two and would catch any cross-contamination.
assert result.prefill_avg_sys_power_watts < 15.0
assert result.generation_avg_sys_power_watts > 15.0
assert result.nodes[0].prefill_avg_sys_power is not None
assert result.nodes[0].generation_avg_sys_power is not None
async def test_no_split_when_unmarked() -> None:
"""If mark_prefill_done() is never called, phase fields stay None."""
state: dict[NodeId, SystemPerformanceProfile] = {
NODE_A: _make_profile(10.0),
}
sampler = PowerSampler(get_node_system=lambda: state, interval=0.02)
async with anyio.create_task_group() as tg:
tg.start_soon(sampler.run)
await anyio.sleep(0.05)
tg.cancel_scope.cancel()
result = sampler.result()
assert result.prefill_seconds is None
assert result.generation_seconds is None
assert result.prefill_energy_joules is None
assert result.generation_energy_joules is None
assert result.nodes[0].prefill_energy_joules is None
assert result.nodes[0].generation_energy_joules is None
async def test_mark_prefill_done_is_idempotent() -> None:
"""Only the first call to mark_prefill_done() should take effect."""
state: dict[NodeId, SystemPerformanceProfile] = {
NODE_A: _make_profile(10.0),
}
sampler = PowerSampler(get_node_system=lambda: state, interval=0.02)
async with anyio.create_task_group() as tg:
tg.start_soon(sampler.run)
await anyio.sleep(0.05)
sampler.mark_prefill_done()
first_prefill_at = sampler._prefill_done_at # pyright: ignore[reportPrivateUsage]
await anyio.sleep(0.05)
sampler.mark_prefill_done()
assert sampler._prefill_done_at == first_prefill_at # pyright: ignore[reportPrivateUsage]
tg.cancel_scope.cancel()
async def test_result_stops_sampling() -> None:
"""Calling result() should stop the sampler's run loop."""
state: dict[NodeId, SystemPerformanceProfile] = {
+4 -68
View File
@@ -1,4 +1,3 @@
from itertools import pairwise
import json
import os
import re
@@ -6,7 +5,6 @@ import sys
import tempfile
import time
from collections.abc import Generator
from math import isqrt
from pathlib import Path
from typing import TYPE_CHECKING, Any, cast
@@ -788,75 +786,13 @@ def mlx_force_oom(size: int = 200000) -> None:
a = mx.random.uniform(shape=(size, size), dtype=mx.float32)
b = mx.random.uniform(shape=(size, size), dtype=mx.float32)
mx.eval(a, b)
c = mx.matmul(a, b) # (size,size)
d = mx.matmul(a, c) # (size,size)
e = mx.matmul(b, c) # (size,size)
f = mx.sigmoid(d + e) # (size,size)
c = mx.matmul(a, b)
d = mx.matmul(a, c)
e = mx.matmul(b, c)
f = mx.sigmoid(d + e)
mx.eval(f)
def mlx_force_oom2(bytes_alloc: int = 1024**5): # the default is 1 petabyte lol
"""
Force an Out-Of-Memory (OOM) error in MLX by performing large tensor operations.
NOTE: probably only works correctly on Apple unified memory
"""
# TODO: I give up. this either uses swap (inconsistently) and only sometimes OOMs
# or if I tune it to be more aggressive then it kerenel panics entirely;
# there MIGHT be a way to make it not use swap memory but I'm not determined enough
# to figure out how :)
def get_size(memory: int):
mat_elem = -(-memory // 4) # per-matrix elements (4 bytes per elem)
# square root to get size (round up if not integer)
root = isqrt(mat_elem)
return root if root**2 == mat_elem else root + 1
def oom(size: int):
mx.set_default_device(mx.gpu)
mx.clear_cache()
# allocate a lot
z = mx.zeros(shape=(size, size), dtype=mx.float32)
t1 = [mx.random.uniform(shape=(size, size), dtype=mx.float32) for _ in range(2)]
# mat mul cycle
t2: list[mx.array] = []
for m1, m2 in pairwise(t1):
t2.append(mx.matmul(m1, m2))
print("t2-run")
mx.eval(*t2)
print("t2-eval")
# sigmoid sum
f = mx.sigmoid(sum(t1, start=z) + sum(t2, start=z))
print("f-run")
mx.eval(f)
print("f-eval")
# use supplied size, or computer appropriate size otherwise
fail_num = 0
while True:
try:
print(f"size {bytes_alloc / 1024**3} GB")
oom(get_size(bytes_alloc))
break
except RuntimeError as e:
max_bytes = re.compile(
r"\[metal::malloc\] Attempting to allocate (?:\d+) bytes which is greater than the maximum allowed buffer size of (?P<max_bytes>\d+) bytes."
).match(str(e))
if max_bytes is None:
raise RuntimeError(
"Tried to get max buffer, but wrong error format"
) from e
bytes_alloc = round(int(max_bytes.group("max_bytes")) * 0.95**fail_num)
fail_num += 1
mlx_force_oom2()
def set_wired_limit_for_model(model_size: Memory):
"""
A context manager to temporarily change the wired limit.
+7 -1
View File
@@ -8,6 +8,10 @@ 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 (
EventRouterBrokenResourceError,
EventRouterClosedResourceError,
)
from exo.shared.apply import apply
from exo.shared.constants import EXO_MAX_INSTANCE_RETRIES
from exo.shared.models.model_cards import ModelId, card_cache
@@ -109,7 +113,9 @@ class Worker:
tg.start_soon(self._event_applier)
tg.start_soon(self._poll_connection_updates)
tg.start_soon(self._reconcile_custom_cards)
except* (EventRouterBrokenResourceError, EventRouterClosedResourceError):
# Event router has been closed (try-star syntax handles error groups)
pass
finally:
# Actual shutdown code - waits for all tasks to complete before executing.
logger.info("Stopping Worker")
-264
View File
@@ -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)
-85
View File
@@ -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)))
@@ -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
+27 -3
View File
@@ -7,6 +7,9 @@ set -uo pipefail
HOST="${1:-localhost:52415}"
MODEL_ID="KevTheHermit/security-testing"
ENCODED_MODEL_ID=$(
python3 -c 'import sys, urllib.parse; print(urllib.parse.quote(sys.argv[1], safe=""))' "$MODEL_ID"
)
CUSTOM_CARDS_DIR="$HOME/.exo/custom_model_cards"
CARD_FILE="$CUSTOM_CARDS_DIR/KevTheHermit--security-testing.toml"
@@ -71,9 +74,30 @@ PLACE_BODY=$(echo "$PLACE_RESPONSE" | sed '$d')
echo " HTTP $PLACE_CODE"
echo " Response: $PLACE_BODY"
# Step 3b: Send a chat completion to actually trigger tokenizer loading
if [ "$PLACE_CODE" -ge 400 ]; then
echo " Placement failed; cannot trigger tokenizer loading."
exit 1
fi
# Step 3b: Wait for placement to materialize before inference.
echo ""
echo "[3b] Sending chat completion to trigger tokenizer load ..."
echo "[3b] Waiting for placed instance ..."
if ! AWAIT_RESPONSE=$(curl -fsS --max-time 65 \
"http://$HOST/instance/await?model_id=$ENCODED_MODEL_ID&timeout_seconds=60" |
awk '/^data: / { sub(/^data: /, ""); print; exit }'); then
echo " Timed out waiting for an instance for $MODEL_ID"
exit 1
fi
if ! printf '%s' "$AWAIT_RESPONSE" | grep -q '"type":"ready"'; then
echo " Timed out waiting for an instance for $MODEL_ID"
exit 1
fi
echo " Instance ready"
# Step 3c: Send a chat completion to actually trigger tokenizer loading
echo ""
echo "[3c] Sending chat completion to trigger tokenizer load ..."
CHAT_RESPONSE=$(curl -s -w "\n%{http_code}" --max-time 30 -X POST "http://$HOST/v1/chat/completions" \
-H "Content-Type: application/json" \
-d "{\"model\":\"$MODEL_ID\",\"messages\":[{\"role\":\"user\",\"content\":\"hello\"}],\"max_tokens\":1}")
@@ -82,7 +106,7 @@ CHAT_BODY=$(echo "$CHAT_RESPONSE" | sed '$d')
echo " HTTP $CHAT_CODE"
echo " Response: $CHAT_BODY"
echo ""
echo "[3c] Checking for RCE proof ..."
echo "[3d] Checking for RCE proof ..."
sleep 5
if [ -f /tmp/exo-rce-proof.txt ]; then
echo " VULNERABLE: Remote code executed!"
Generated
+39 -16
View File
@@ -24,7 +24,7 @@ prerelease-mode = "allow"
members = [
"exo",
"exo-bench",
"exo-pyo3-bindings",
"exo-rs",
"exo-tools",
]
overrides = [{ name = "opencv-python", marker = "python_full_version < '0'" }]
@@ -423,7 +423,7 @@ dependencies = [
{ name = "aiofiles", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "aiohttp", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "anyio", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "exo-pyo3-bindings", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "exo-rs", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "fastapi", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "filelock", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "httpx", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
@@ -434,6 +434,7 @@ dependencies = [
{ name = "openai-harmony", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "psutil", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "pydantic", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "python-daemon", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "python-multipart", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "rustworkx", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "tiktoken", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
@@ -449,7 +450,7 @@ build = [
]
mlx = [
{ name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0.dev20260512+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0.dev20260506+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl" }, marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
@@ -460,7 +461,7 @@ mlx = [
]
mlx-cpu = [
{ name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mlx", version = "0.32.0.dev20260512+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-mlx-cpu') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0.dev20260506+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-mlx-cpu') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cpu') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra != 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra != 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra != 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl" }, marker = "(platform_machine != 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cpu') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra != 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra != 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra != 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx-cpu", marker = "sys_platform == 'linux'" },
@@ -472,7 +473,7 @@ mlx-cpu = [
]
mlx-cuda12 = [
{ name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mlx", version = "0.32.0.dev20260512+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0.dev20260506+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra != 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl" }, marker = "(platform_machine != 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra != 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx-cuda-12", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx_cuda_12-0.32.0-py3-none-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra != 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
@@ -486,7 +487,7 @@ mlx-cuda12 = [
]
mlx-cuda13 = [
{ name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mlx", version = "0.32.0.dev20260512+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0.dev20260506+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl" }, marker = "(platform_machine != 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx-cuda-13", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx_cuda_13-0.32.0-py3-none-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none')" },
@@ -522,7 +523,7 @@ requires-dist = [
{ name = "exo", extras = ["mlx"], marker = "extra == 'mlx-cpu'" },
{ name = "exo", extras = ["mlx"], marker = "extra == 'mlx-cuda12'" },
{ name = "exo", extras = ["mlx"], marker = "extra == 'mlx-cuda13'" },
{ name = "exo-pyo3-bindings", editable = "rust/exo_pyo3_bindings" },
{ name = "exo-rs", editable = "rust/exo_rs" },
{ name = "fastapi", specifier = ">=0.116.1" },
{ name = "filelock", specifier = ">=3.18.0" },
{ name = "httpx", specifier = ">=0.28.1" },
@@ -548,6 +549,7 @@ requires-dist = [
{ name = "openai-harmony", specifier = ">=0.0.8" },
{ name = "psutil", specifier = ">=7.0.0" },
{ name = "pydantic", specifier = ">=2.11.7" },
{ name = "python-daemon", specifier = ">=3.1.2" },
{ name = "python-multipart", specifier = ">=0.0.21" },
{ name = "rustworkx", specifier = ">=0.17.1" },
{ name = "tiktoken", specifier = ">=0.12.0" },
@@ -620,13 +622,13 @@ requires-dist = [
]
[[package]]
name = "exo-pyo3-bindings"
version = "0.2.2"
source = { editable = "rust/exo_pyo3_bindings" }
name = "exo-rs"
version = "0.3.0"
source = { editable = "rust/exo_rs" }
[package.dev-dependencies]
dev = [
{ name = "exo-pyo3-bindings", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "exo-rs", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "pytest", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "pytest-asyncio", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
]
@@ -635,7 +637,7 @@ dev = [
[package.metadata.requires-dev]
dev = [
{ name = "exo-pyo3-bindings", editable = "rust/exo_pyo3_bindings" },
{ name = "exo-rs", editable = "rust/exo_rs" },
{ name = "pytest", specifier = ">=8.4.0" },
{ name = "pytest-asyncio", specifier = ">=1.0.0" },
]
@@ -1162,6 +1164,15 @@ math = [
{ name = "sympy", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
]
[[package]]
name = "lockfile"
version = "0.12.2"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/17/47/72cb04a58a35ec495f96984dddb48232b551aafb95bde614605b754fe6f7/lockfile-0.12.2.tar.gz", hash = "sha256:6aed02de03cba24efabcd600b30540140634fc06cfa603822d508d5361e9f799", size = 20874, upload-time = "2015-11-25T18:29:58.279Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/c8/22/9460e311f340cb62d26a38c419b1381b8593b0bb6b5d1f056938b086d362/lockfile-0.12.2-py2.py3-none-any.whl", hash = "sha256:6c3cb24f344923d30b2785d5ad75182c8ea7ac1b6171b08657258ec7429d50fa", size = 13564, upload-time = "2015-11-25T18:29:51.462Z" },
]
[[package]]
name = "loguru"
version = "0.7.3"
@@ -1332,7 +1343,7 @@ dependencies = [
{ name = "hf-transfer", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "matplotlib", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0.dev20260512+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0.dev20260506+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl" }, marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
@@ -1363,7 +1374,7 @@ sdist = { url = "https://files.pythonhosted.org/packages/55/fa/96d4cc7ada2833571
[[package]]
name = "mlx"
version = "0.32.0.dev20260512+cc3f3e60"
version = "0.32.0.dev20260506+cc3f3e60"
source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }
resolution-markers = [
"sys_platform == 'darwin'",
@@ -1538,7 +1549,7 @@ version = "0.31.3"
source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Fdeepseek-v4#6a3df6cd6b00a347ee40f12d97a182aaf86ea599" }
dependencies = [
{ name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0.dev20260512+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0.dev20260506+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
@@ -1554,7 +1565,7 @@ dependencies = [
{ name = "datasets", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "fastapi", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "miniaudio", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0.dev20260512+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0.dev20260506+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl" }, marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
{ name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
@@ -2300,6 +2311,18 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/27/98/822b924a4a3eb58aacba84444c7439fce32680592f394de26af9c76e2569/pytest_env-1.2.0-py3-none-any.whl", hash = "sha256:d7e5b7198f9b83c795377c09feefa45d56083834e60d04767efd64819fc9da00", size = 6251, upload-time = "2025-10-09T19:15:46.077Z" },
]
[[package]]
name = "python-daemon"
version = "3.1.2"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "lockfile", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
]
sdist = { url = "https://files.pythonhosted.org/packages/3d/37/4f10e37bdabc058a32989da2daf29e57dc59dbc5395497f3d36d5f5e2694/python_daemon-3.1.2.tar.gz", hash = "sha256:f7b04335adc473de877f5117e26d5f1142f4c9f7cd765408f0877757be5afbf4", size = 71576, upload-time = "2024-12-03T08:41:07.843Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/45/3c/b88167e2d6785c0e781ee5d498b07472aeb9b6765da3b19e7cc9e0813841/python_daemon-3.1.2-py3-none-any.whl", hash = "sha256:b906833cef63502994ad48e2eab213259ed9bb18d54fa8774dcba2ff7864cec6", size = 30872, upload-time = "2024-12-03T08:41:03.322Z" },
]
[[package]]
name = "python-dateutil"
version = "2.9.0.post0"