mirror of
https://github.com/exo-explore/exo.git
synced 2026-09-08 11:35:40 -04:00
Compare commits
94
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0abd91c6ca | ||
|
|
6440046af1 | ||
|
|
7f0ba1628a | ||
|
|
c2efbfd35d | ||
|
|
9fb66e313b | ||
|
|
22adc34311 | ||
|
|
04a688083f | ||
|
|
9e1a41b687 | ||
|
|
87657cf061 | ||
|
|
7258bdbfd1 | ||
|
|
41afeb7c5e | ||
|
|
dc8e3a97fa | ||
|
|
f3f126b69e | ||
|
|
8dd3715d13 | ||
|
|
3b49c1b493 | ||
|
|
b6b9761d5f | ||
|
|
de7454e873 | ||
|
|
5a2a96b8c5 | ||
|
|
4f2427bd22 | ||
|
|
fa6faced53 | ||
|
|
98901d3234 | ||
|
|
21236dd34e | ||
|
|
b7e374699c | ||
|
|
8f80dc23d1 | ||
|
|
fa4ec5a979 | ||
|
|
a434944921 | ||
|
|
8688ff1855 | ||
|
|
c405d89a26 | ||
|
|
dbfc7f8ba6 | ||
|
|
c5d8f8c331 | ||
|
|
91c4d2b5fc | ||
|
|
adec95a0a3 | ||
|
|
9d0d4e878c | ||
|
|
ced27f050d | ||
|
|
6bc00555ac | ||
|
|
8f97457a0d | ||
|
|
d395c50fba | ||
|
|
2a507509fa | ||
|
|
2a230be446 | ||
|
|
d5ab01d9a1 | ||
|
|
dc30ca7fa6 | ||
|
|
04fec92987 | ||
|
|
6035f66b61 | ||
|
|
733cd1b9c2 | ||
|
|
462a72e22d | ||
|
|
b2c8dc9c49 | ||
|
|
d5f1d0792f | ||
|
|
683e1b1a29 | ||
|
|
d4c86b68b1 | ||
|
|
8c3e360bd1 | ||
|
|
b00b1882da | ||
|
|
22228c5d02 | ||
|
|
37440ce6a0 | ||
|
|
8c9b000ec3 | ||
|
|
f1f393bfdf | ||
|
|
131e3af4ff | ||
|
|
f4a2ffa577 | ||
|
|
817c556851 | ||
|
|
1e8d4abe94 | ||
|
|
c2ecc8b59e | ||
|
|
a5ec6f783f | ||
|
|
e04208605e | ||
|
|
06fa9c3fee | ||
|
|
f54a701979 | ||
|
|
24cab4799c | ||
|
|
fdf5f0c00b | ||
|
|
9604a1a18c | ||
|
|
13e5bf8c16 | ||
|
|
ae3b195868 | ||
|
|
c9c6b59562 | ||
|
|
72d3bfc088 | ||
|
|
2db2abbb1e | ||
|
|
2cacfb5a9b | ||
|
|
051563a303 | ||
|
|
d3d680f569 | ||
|
|
c7c449f550 | ||
|
|
df2925ce15 | ||
|
|
13b4ac4162 | ||
|
|
4883bcd3a9 | ||
|
|
d4a61620d2 | ||
|
|
6649ce7f0c | ||
|
|
bc06e029be | ||
|
|
92e9c9f8c2 | ||
|
|
5abd06735b | ||
|
|
7186ec2423 | ||
|
|
f9fda49ae8 | ||
|
|
e892f7fb8a | ||
|
|
b7730f743d | ||
|
|
09f9ea313f | ||
|
|
e12744edd6 | ||
|
|
8506e7a4dc | ||
|
|
81d7cb0fcd | ||
|
|
439f59924a | ||
|
|
629c55d6ba |
No files matched your search
@@ -1,8 +1 @@
|
||||
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
|
||||
@@ -34,7 +34,7 @@ jobs:
|
||||
SPARKLE_S3_PREFIX: ${{ secrets.SPARKLE_S3_PREFIX }}
|
||||
AWS_REGION: ${{ secrets.AWS_REGION }}
|
||||
EXO_BUILD_NUMBER: ${{ github.run_number }}
|
||||
EXO_LIBP2P_NAMESPACE: ${{ github.ref_name }}
|
||||
EXO_NAMESPACE: ${{ github.ref_name }}
|
||||
|
||||
steps:
|
||||
# ============================================================
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -15,8 +15,8 @@ cd dashboard && npm install && npm run build && cd ..
|
||||
# Run exo (starts both master and worker with API at http://localhost:52415)
|
||||
uv run exo
|
||||
|
||||
# Run with verbose logging
|
||||
uv run exo -v # or -vv for more verbose
|
||||
# Run with debug logging
|
||||
uv run exo --verbosity=debug # or: uv run exo -v debug
|
||||
|
||||
# Run tests (excludes slow tests by default)
|
||||
uv run pytest
|
||||
@@ -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
|
||||
|
||||
Generated
+3019
-2337
File diff suppressed because it is too large.
Load diff
+57
-10
@@ -1,6 +1,6 @@
|
||||
[workspace]
|
||||
resolver = "3"
|
||||
members = ["rust/networking", "rust/exo_pyo3_bindings", "rust/util"]
|
||||
members = ["rust/exo_rs", "rust/networking", "rust/util"]
|
||||
|
||||
[workspace.package]
|
||||
version = "0.0.1"
|
||||
@@ -22,29 +22,76 @@ opt-level = 3
|
||||
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"
|
||||
derive_more = "2.1"
|
||||
|
||||
# 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 = "=0.3.1"
|
||||
|
||||
[patch.crates-io]
|
||||
pidfile-rs = { git = "https://github.com/AndreiCravtov/pidfile-rs", rev = "faf6b79" }
|
||||
|
||||
# zenoh
|
||||
zenoh = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-buffers = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-codec = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-collections = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-config = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-core = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-crypto = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-keyexpr = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-link = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-link-commons = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-link-quic = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-link-quic_datagram = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-link-tcp = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-link-tls = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-link-udp = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-link-unixsock_stream = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-link-ws = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-macros = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-plugin-trait = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-protocol = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-result = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-runtime = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-sync = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-task = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-transport = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh-util = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
zenoh_backend_traits = { git = "https://github.com/evanev7/zenoh.git", rev = "2fed58e" }
|
||||
|
||||
[workspace.lints.rust]
|
||||
static_mut_refs = "warn" # Or use "warn" instead of deny
|
||||
|
||||
@@ -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).
|
||||
|
||||
|
||||
@@ -229,18 +229,24 @@ 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:**
|
||||
|
||||
The macOS app includes a custom namespace feature that allows you to isolate your exo cluster from others on the same network. This is configured through the `EXO_LIBP2P_NAMESPACE` setting:
|
||||
The macOS app includes a custom namespace feature that allows you to isolate your exo cluster from others on the same network. The app passes this to exo with `--namespace`; source runs can use `--namespace` or `EXO_NAMESPACE`.
|
||||
|
||||
- **Use cases**:
|
||||
- Running multiple separate exo clusters on the same network
|
||||
- Isolating development/testing clusters from production clusters
|
||||
- Preventing accidental cluster joining
|
||||
|
||||
- **Configuration**: Access this setting in the app's Advanced settings (or set the `EXO_LIBP2P_NAMESPACE` environment variable when running from source)
|
||||
- **Configuration**: Access this setting in the app's Advanced settings (or use `--namespace` / `EXO_NAMESPACE` when running from source)
|
||||
|
||||
The namespace is logged on startup for debugging purposes.
|
||||
|
||||
@@ -295,20 +301,51 @@ After that, RDMA will be enabled in macOS and exo will take care of the rest.
|
||||
|
||||
---
|
||||
|
||||
## Environment Variables
|
||||
## Configuration
|
||||
|
||||
exo supports several environment variables for configuration:
|
||||
exo reads configuration in this order:
|
||||
|
||||
1. Built-in defaults
|
||||
2. `config.toml`
|
||||
3. Environment variables
|
||||
4. CLI arguments
|
||||
|
||||
If a default `config.toml` exists in the resolved Exo config directory, exo reads
|
||||
it automatically. Use `--config-file` or `EXO_CONFIG_FILE` to point at an
|
||||
existing custom file.
|
||||
|
||||
Example `config.toml`:
|
||||
|
||||
```toml
|
||||
verbosity = "info"
|
||||
continuous_batching_enabled = true
|
||||
max_concurrent_requests = 8
|
||||
offline = false
|
||||
image_models_enabled = false
|
||||
tracing_enabled = false
|
||||
disaggregation_enabled = false
|
||||
# fast_synch = true # or false; omit for automatic behavior
|
||||
```
|
||||
|
||||
Some paths are resolved before `config.toml` can be loaded, so they are configured
|
||||
only through CLI arguments or environment variables.
|
||||
|
||||
| Variable | Description | Default |
|
||||
|----------|-------------|---------|
|
||||
| `EXO_DEFAULT_MODELS_DIR` | Default directory for model downloads and caches. Always first in the writable dirs list. | `~/.local/share/exo/models` (Linux) or `~/.exo/models` (macOS) |
|
||||
| `EXO_HOME` | Base directory for Exo config, data, and cache paths. CLI form: `--exo-home`. | Platform config/data/cache dirs |
|
||||
| `EXO_CONFIG_FILE` | Path to Exo's `config.toml`. CLI form: `--config-file`. | Resolved Exo config dir + `config.toml` |
|
||||
| `EXO_DEFAULT_MODELS_DIR` | Default directory for model downloads and caches. Always first in the writable dirs list. | Platform data dir + `models` |
|
||||
| `EXO_MODELS_DIRS` | Colon-separated additional writable directories for model downloads. Checked in order after the default; first with enough free space is used. | None |
|
||||
| `EXO_MODELS_READ_ONLY_DIRS` | Colon-separated read-only directories to search for pre-downloaded models (e.g., NFS mounts, shared storage). Models here cannot be deleted. | None |
|
||||
| `EXO_OFFLINE` | Run without internet connection (uses only local models) | `false` |
|
||||
| `EXO_ENABLE_IMAGE_MODELS` | Enable image model support | `false` |
|
||||
| `EXO_LIBP2P_NAMESPACE` | Custom namespace for cluster isolation | None |
|
||||
| `EXO_FAST_SYNCH` | Control MLX_METAL_FAST_SYNCH behavior (for JACCL backend) | Auto |
|
||||
| `EXO_TRACING_ENABLED` | Enable distributed tracing for performance analysis | `false` |
|
||||
| `EXO_NAMESPACE` | Custom namespace for cluster isolation. CLI form: `--namespace`. | Package version |
|
||||
| `EXO_VERBOSITY` | Verbosity filter. CLI forms: `--verbosity=<LEVEL>` or `-v <LEVEL>`. Valid levels: `off`, `error`, `warn`, `info`, `debug`, `trace`. | `info` |
|
||||
| `EXO_NO_BATCH` | Boolean inverse of `continuous_batching_enabled`. CLI form: `--no-batch[=true\|false]`. | `false` |
|
||||
| `EXO_MAX_CONCURRENT_REQUESTS` | Maximum number of concurrent generation requests per runner. CLI form: `--max-concurrent-requests`. | `8` |
|
||||
| `EXO_OFFLINE` | Boolean equivalent for offline mode. CLI form: `--offline[=true\|false]`. Uses only local models and skips internet checks. | `false` |
|
||||
| `EXO_IMAGE_MODELS_ENABLED` | Boolean equivalent for image model support. CLI form: `--enable-image-models[=true\|false]`. | `false` |
|
||||
| `EXO_TRACING_ENABLED` | Boolean equivalent for distributed tracing. CLI form: `--enable-tracing[=true\|false]`. | `false` |
|
||||
| `EXO_DISAGGREGATION_ENABLED` | Boolean equivalent for prefill/decode disaggregation. CLI form: `--enable-disaggregation[=true\|false]`. | `false` |
|
||||
| `EXO_FAST_SYNCH` | Boolean FAST_SYNCH override for JACCL backend. CLI form: `--fast-synch=true\|false`. Leave unset or omit `fast_synch` from `config.toml` for automatic behavior. | Auto |
|
||||
|
||||
**Example usage:**
|
||||
|
||||
@@ -319,14 +356,51 @@ EXO_MODELS_READ_ONLY_DIRS=/mnt/nfs/models:/opt/ai-models uv run exo
|
||||
# Download models to an external SSD (falls back to default dir if full)
|
||||
EXO_MODELS_DIRS=/Volumes/ExternalSSD/exo-models uv run exo
|
||||
|
||||
# Run in offline mode
|
||||
# Run it in offline mode (CLI or ENV arg)
|
||||
uv run exo --offline
|
||||
EXO_OFFLINE=true uv run exo
|
||||
|
||||
# Enable image models
|
||||
EXO_ENABLE_IMAGE_MODELS=true uv run exo
|
||||
uv run exo --enable-image-models
|
||||
uv run exo --enable-image-models=false
|
||||
EXO_IMAGE_MODELS_ENABLED=true uv run exo
|
||||
|
||||
# Enable distributed tracing
|
||||
uv run exo --enable-tracing
|
||||
uv run exo --enable-tracing=false
|
||||
EXO_TRACING_ENABLED=true uv run exo
|
||||
|
||||
# Enable prefill/decode disaggregation
|
||||
uv run exo --enable-disaggregation
|
||||
uv run exo --enable-disaggregation=false
|
||||
EXO_DISAGGREGATION_ENABLED=true uv run exo
|
||||
|
||||
# Disable continuous batching
|
||||
uv run exo --no-batch
|
||||
EXO_NO_BATCH=true uv run exo
|
||||
|
||||
# Set request concurrency
|
||||
uv run exo --max-concurrent-requests 16
|
||||
EXO_MAX_CONCURRENT_REQUESTS=16 uv run exo
|
||||
|
||||
# Use custom namespace for cluster isolation
|
||||
EXO_LIBP2P_NAMESPACE=my-dev-cluster uv run exo
|
||||
uv run exo --namespace=my-dev-cluster
|
||||
EXO_NAMESPACE=my-dev-cluster uv run exo
|
||||
|
||||
# Set logging verbosity (CLI or ENV arg)
|
||||
uv run exo --verbosity=debug
|
||||
uv run exo -v debug
|
||||
EXO_VERBOSITY=debug uv run exo
|
||||
|
||||
# Only show error logs
|
||||
uv run exo --quiet
|
||||
|
||||
# Disable logging output
|
||||
uv run exo --verbosity=off
|
||||
|
||||
# Force MLX FAST_SYNCH off (CLI or ENV arg)
|
||||
uv run exo --fast-synch=false
|
||||
EXO_FAST_SYNCH=false uv run exo
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
@@ -1,14 +1,13 @@
|
||||
3. Task cancellation. When API http request gets cancelled, it should cancel corresponding task.
|
||||
1. Restore bootstrap peers support.
|
||||
|
||||
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
|
||||
@@ -183,6 +183,7 @@ final class ExoProcessController: ObservableObject {
|
||||
at: exoHomeURL, withIntermediateDirectories: true
|
||||
)
|
||||
child.currentDirectoryURL = exoHomeURL
|
||||
child.arguments = makeArgs()
|
||||
child.environment = makeEnvironment(for: runtimeURL)
|
||||
|
||||
child.standardOutput = FileHandle.nullDevice
|
||||
@@ -352,21 +353,12 @@ 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()
|
||||
if !hfToken.isEmpty {
|
||||
environment["HF_TOKEN"] = hfToken
|
||||
}
|
||||
if !hfEndpoint.isEmpty {
|
||||
environment["HF_ENDPOINT"] = hfEndpoint
|
||||
}
|
||||
if enableImageModels {
|
||||
environment["EXO_ENABLE_IMAGE_MODELS"] = "true"
|
||||
}
|
||||
if offlineMode {
|
||||
environment["EXO_OFFLINE"] = "true"
|
||||
}
|
||||
environment["EXO_FAST_SYNCH"] = fastSynchEnabled ? "true" : "false"
|
||||
|
||||
var paths: [String] = []
|
||||
if let existing = environment["PATH"], !existing.isEmpty {
|
||||
paths = existing.split(separator: ":").map(String.init)
|
||||
@@ -416,6 +408,20 @@ final class ExoProcessController: ObservableObject {
|
||||
return environment
|
||||
}
|
||||
|
||||
private func makeArgs() -> [String] {
|
||||
var args = [
|
||||
"--namespace=\(computeNamespace())",
|
||||
"--fast-synch=\(fastSynchEnabled ? "true" : "false")",
|
||||
]
|
||||
if offlineMode {
|
||||
args.append("--offline")
|
||||
}
|
||||
if enableImageModels {
|
||||
args.append("--enable-image-models")
|
||||
}
|
||||
return args
|
||||
}
|
||||
|
||||
private func buildTag() -> String {
|
||||
if let tag = Bundle.main.infoDictionary?["EXOBuildTag"] as? String, !tag.isEmpty {
|
||||
return tag
|
||||
|
||||
@@ -40,7 +40,7 @@
|
||||
{:else if !enabled}
|
||||
<div class="text-exo-light-gray/60 text-sm">
|
||||
No advanced features enabled. Set <code
|
||||
class="text-exo-yellow font-mono">ENABLE_DISAGGREGATION=true</code
|
||||
class="text-exo-yellow font-mono">--enable-disaggregation</code
|
||||
> on the cluster to access prefill/decode disaggregation.
|
||||
</div>
|
||||
{:else}
|
||||
|
||||
@@ -194,7 +194,8 @@
|
||||
>
|
||||
<div class="text-sm">No traces found.</div>
|
||||
<div class="text-xs text-exo-light-gray/70">
|
||||
Run exo with EXO_TRACING_ENABLED=1 to collect traces.
|
||||
Run exo with --enable-tracing or EXO_TRACING_ENABLED=true to collect
|
||||
traces.
|
||||
</div>
|
||||
</div>
|
||||
{:else}
|
||||
|
||||
@@ -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
@@ -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": {
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -23,7 +23,7 @@ sync-clean:
|
||||
|
||||
rust-rebuild:
|
||||
PYO3_PYTHON="$(uv run python -c 'import sys; print(sys.executable)')" cargo run --bin stub_gen
|
||||
uv sync --reinstall-package exo_pyo3_bindings
|
||||
uv sync --reinstall-package exo_rs
|
||||
|
||||
build-dashboard:
|
||||
#!/usr/bin/env bash
|
||||
|
||||
+11
-12
@@ -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",
|
||||
@@ -76,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'" },
|
||||
@@ -101,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'" },
|
||||
@@ -241,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"]
|
||||
|
||||
+12
-8
@@ -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
|
||||
'';
|
||||
};
|
||||
@@ -210,6 +211,9 @@ let
|
||||
text: name: pkgs.writeShellApplication {
|
||||
inherit name;
|
||||
text = ''
|
||||
unset PYTHONPATH
|
||||
unset PYTHONHOME
|
||||
|
||||
LD_LIBRARY_PATH="${libPath}''${LD_LIBRARY_PATH:+:}''${LD_LIBRARY_PATH:-}" exec \
|
||||
${lib.optionalString cudaSupport "nixglhost "} ${text}
|
||||
'';
|
||||
@@ -223,7 +227,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 {
|
||||
|
||||
@@ -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 = { git = "https://github.com/AndreiCravtov/pidfile-rs" }
|
||||
|
||||
# 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"
|
||||
@@ -1,152 +0,0 @@
|
||||
# This file is automatically generated by pyo3_stub_gen
|
||||
# ruff: noqa: E501, F401
|
||||
|
||||
import builtins
|
||||
import os
|
||||
import pathlib
|
||||
import typing
|
||||
|
||||
@typing.final
|
||||
class AllQueuesFullError(builtins.Exception):
|
||||
def __new__(cls, *args: typing.Any) -> AllQueuesFullError: ...
|
||||
def __repr__(self) -> builtins.str: ...
|
||||
def __str__(self) -> builtins.str: ...
|
||||
|
||||
@typing.final
|
||||
class Keypair:
|
||||
r"""
|
||||
Identity keypair of a node.
|
||||
"""
|
||||
@staticmethod
|
||||
def generate() -> Keypair:
|
||||
r"""
|
||||
Generate a new Ed25519 keypair.
|
||||
"""
|
||||
@staticmethod
|
||||
def from_bytes(bytes: bytes) -> Keypair:
|
||||
r"""
|
||||
Construct an Ed25519 keypair from secret key bytes
|
||||
"""
|
||||
def to_bytes(self) -> bytes:
|
||||
r"""
|
||||
Get the secret key bytes underlying the keypair
|
||||
"""
|
||||
def to_node_id(self) -> builtins.str:
|
||||
r"""
|
||||
Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
|
||||
"""
|
||||
|
||||
@typing.final
|
||||
class MessageTooLargeError(builtins.Exception):
|
||||
def __new__(cls, *args: typing.Any) -> MessageTooLargeError: ...
|
||||
def __repr__(self) -> builtins.str: ...
|
||||
def __str__(self) -> builtins.str: ...
|
||||
|
||||
@typing.final
|
||||
class NetworkingHandle:
|
||||
def __new__(cls, identity: Keypair, bootstrap_peers: typing.Sequence[builtins.str], listen_port: builtins.int) -> NetworkingHandle: ...
|
||||
async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
|
||||
r"""
|
||||
Subscribe to a `GossipSub` topic.
|
||||
|
||||
Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
|
||||
"""
|
||||
async def gossipsub_unsubscribe(self, topic: builtins.str) -> builtins.bool:
|
||||
r"""
|
||||
Unsubscribes from a `GossipSub` topic.
|
||||
|
||||
Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
|
||||
"""
|
||||
async def gossipsub_publish(self, topic: builtins.str, data: bytes) -> None:
|
||||
r"""
|
||||
Publishes a message with multiple topics to the `GossipSub` network.
|
||||
|
||||
If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
|
||||
"""
|
||||
async def recv(self) -> PyFromSwarm: ...
|
||||
|
||||
@typing.final
|
||||
class NoPeersSubscribedToTopicError(builtins.Exception):
|
||||
def __new__(cls, *args: typing.Any) -> NoPeersSubscribedToTopicError: ...
|
||||
def __repr__(self) -> builtins.str: ...
|
||||
def __str__(self) -> builtins.str: ...
|
||||
|
||||
@typing.final
|
||||
class Pidfile:
|
||||
r"""
|
||||
A PID file protected with a lock.
|
||||
|
||||
An instance of `Pidfile` can be used to manage a PID file: create it,
|
||||
lock it, detect already running daemons. It is backed by [`pidfile`][]
|
||||
functions of `libbsd`/`libutil` which use `flopen` to lock the PID
|
||||
file.
|
||||
|
||||
When a PID file is created, the process ID of the current process is
|
||||
*not* written there, making it possible to lock the PID file before
|
||||
forking and only write the ID of the forked process when it is ready.
|
||||
|
||||
The PID file is deleted automatically when the `Pidfile` comes out of
|
||||
the scope. To close the PID file without deleting it, for example, in
|
||||
the parent process of a forked daemon, call `close()`.
|
||||
|
||||
[`exit`]: https://doc.rust-lang.org/std/process/fn.exit.html
|
||||
[`pidfile`]: https://linux.die.net/man/3/pidfile
|
||||
[`daemon`(3)]: https://linux.die.net/man/3/daemon
|
||||
"""
|
||||
def __new__(cls, path: builtins.str | os.PathLike | pathlib.Path, mode: builtins.int) -> Pidfile:
|
||||
r"""
|
||||
Creates a new PID file and locks it.
|
||||
|
||||
If the PID file cannot be locked, returns `PidfileError::AlreadyRunning` with
|
||||
a PID of the already running process, or `None` if no PID has been written to
|
||||
the PID file yet.
|
||||
"""
|
||||
def write(self) -> None:
|
||||
r"""
|
||||
Writes the current process ID to the PID file.
|
||||
|
||||
The file is truncated before writing.
|
||||
"""
|
||||
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: ...
|
||||
|
||||
...
|
||||
|
||||
@@ -1,8 +0,0 @@
|
||||
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()?;
|
||||
stub.generate()?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -1,49 +0,0 @@
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
from _pytest.capture import CaptureFixture
|
||||
from exo_pyo3_bindings import (
|
||||
Keypair,
|
||||
NetworkingHandle,
|
||||
NoPeersSubscribedToTopicError,
|
||||
Pidfile,
|
||||
PyFromSwarm,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sleep_on_multiple_items() -> None:
|
||||
print("PYTHON: starting handle")
|
||||
h = NetworkingHandle(Keypair.generate(), [], 0)
|
||||
|
||||
rt = asyncio.create_task(_await_recv(h))
|
||||
|
||||
# sleep for 4 ticks
|
||||
for i in range(4):
|
||||
await asyncio.sleep(1)
|
||||
|
||||
try:
|
||||
await h.gossipsub_publish("topic", b"somehting or other")
|
||||
except NoPeersSubscribedToTopicError as e:
|
||||
print("caught it", e)
|
||||
|
||||
|
||||
def test_pidfile(capsys: CaptureFixture[str]):
|
||||
with capsys.disabled():
|
||||
print("\nbefore python")
|
||||
scoped_lock_file()
|
||||
print("after python")
|
||||
|
||||
|
||||
async def _await_recv(h: NetworkingHandle):
|
||||
while True:
|
||||
event = await h.recv()
|
||||
match event:
|
||||
case PyFromSwarm.Connection() as c:
|
||||
print(f"PYTHON: connection update: {c}")
|
||||
case PyFromSwarm.Message() as m:
|
||||
print(f"PYTHON: message: {m}")
|
||||
|
||||
|
||||
def scoped_lock_file():
|
||||
a = Pidfile("/tmp/lock.pid", 0o0600)
|
||||
@@ -0,0 +1,71 @@
|
||||
[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
|
||||
util.workspace = true
|
||||
extend.workspace = true
|
||||
itertools = "0.14.0"
|
||||
|
||||
# interop
|
||||
pyo3 = { workspace = true, features = ["serde", "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 }
|
||||
|
||||
# config
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
serde_with = { version = "3.21", features = ["macros"] }
|
||||
postcard = { version = "1.1.3", features = ["alloc"] }
|
||||
clap = { version = "4.6.1", features = [
|
||||
"derive",
|
||||
"env",
|
||||
"unicode",
|
||||
"string",
|
||||
"wrap_help",
|
||||
] }
|
||||
figment = { version = "0.10", features = ["env", "toml"] }
|
||||
dirs = "6.0"
|
||||
toml_edit = "0.25"
|
||||
|
||||
|
||||
# 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.
@@ -0,0 +1,490 @@
|
||||
# This file is automatically generated by pyo3_stub_gen
|
||||
# ruff: noqa: E501, F401, F403, F405
|
||||
|
||||
import builtins
|
||||
import enum
|
||||
import os
|
||||
import pathlib
|
||||
import typing
|
||||
__all__ = [
|
||||
"AppArgs",
|
||||
"AppSettings",
|
||||
"BootstrapArgs",
|
||||
"BootstrapSettings",
|
||||
"CliArgs",
|
||||
"ExoHome",
|
||||
"FromSwarm",
|
||||
"LogFiles",
|
||||
"ModelsDirs",
|
||||
"NetworkingHandle",
|
||||
"Pidfile",
|
||||
"PidfileError",
|
||||
"RejectedArgs",
|
||||
"VerbosityFilter",
|
||||
]
|
||||
|
||||
@typing.final
|
||||
class AppArgs:
|
||||
r"""
|
||||
Arguments that participate in application settings resolution.
|
||||
|
||||
These values may come from defaults, `config.toml`, environment variables, or
|
||||
CLI arguments. Unlike [`BootstrapArgs`](crate::config::bootstrap::BootstrapArgs),
|
||||
they do not participate in finding or loading `config.toml`.
|
||||
|
||||
# Important
|
||||
- Make sure all fields are [`Option<T>`] so they can be layered with other
|
||||
settings sources.
|
||||
"""
|
||||
@property
|
||||
def verbosity(self) -> typing.Optional[VerbosityFilter]: ...
|
||||
@verbosity.setter
|
||||
def verbosity(self, value: typing.Optional[VerbosityFilter]) -> None: ...
|
||||
@property
|
||||
def continuous_batching_enabled(self) -> typing.Optional[builtins.bool]: ...
|
||||
@continuous_batching_enabled.setter
|
||||
def continuous_batching_enabled(self, value: typing.Optional[builtins.bool]) -> None: ...
|
||||
@property
|
||||
def max_concurrent_requests(self) -> typing.Optional[builtins.int]: ...
|
||||
@max_concurrent_requests.setter
|
||||
def max_concurrent_requests(self, value: typing.Optional[builtins.int]) -> None: ...
|
||||
@property
|
||||
def offline(self) -> typing.Optional[builtins.bool]: ...
|
||||
@offline.setter
|
||||
def offline(self, value: typing.Optional[builtins.bool]) -> None: ...
|
||||
@property
|
||||
def image_models_enabled(self) -> typing.Optional[builtins.bool]: ...
|
||||
@image_models_enabled.setter
|
||||
def image_models_enabled(self, value: typing.Optional[builtins.bool]) -> None: ...
|
||||
@property
|
||||
def tracing_enabled(self) -> typing.Optional[builtins.bool]: ...
|
||||
@tracing_enabled.setter
|
||||
def tracing_enabled(self, value: typing.Optional[builtins.bool]) -> None: ...
|
||||
@property
|
||||
def disaggregation_enabled(self) -> typing.Optional[builtins.bool]: ...
|
||||
@disaggregation_enabled.setter
|
||||
def disaggregation_enabled(self, value: typing.Optional[builtins.bool]) -> None: ...
|
||||
@property
|
||||
def fast_synch(self) -> typing.Optional[builtins.bool]: ...
|
||||
@fast_synch.setter
|
||||
def fast_synch(self, value: typing.Optional[builtins.bool]) -> None: ...
|
||||
|
||||
@typing.final
|
||||
class AppSettings:
|
||||
@property
|
||||
def verbosity(self) -> VerbosityFilter: ...
|
||||
@verbosity.setter
|
||||
def verbosity(self, value: VerbosityFilter) -> None: ...
|
||||
@property
|
||||
def continuous_batching_enabled(self) -> builtins.bool: ...
|
||||
@continuous_batching_enabled.setter
|
||||
def continuous_batching_enabled(self, value: builtins.bool) -> None: ...
|
||||
@property
|
||||
def max_concurrent_requests(self) -> builtins.int: ...
|
||||
@max_concurrent_requests.setter
|
||||
def max_concurrent_requests(self, value: builtins.int) -> None: ...
|
||||
@property
|
||||
def offline(self) -> builtins.bool: ...
|
||||
@offline.setter
|
||||
def offline(self, value: builtins.bool) -> None: ...
|
||||
@property
|
||||
def image_models_enabled(self) -> builtins.bool: ...
|
||||
@image_models_enabled.setter
|
||||
def image_models_enabled(self, value: builtins.bool) -> None: ...
|
||||
@property
|
||||
def tracing_enabled(self) -> builtins.bool: ...
|
||||
@tracing_enabled.setter
|
||||
def tracing_enabled(self, value: builtins.bool) -> None: ...
|
||||
@property
|
||||
def disaggregation_enabled(self) -> builtins.bool: ...
|
||||
@disaggregation_enabled.setter
|
||||
def disaggregation_enabled(self, value: builtins.bool) -> None: ...
|
||||
@property
|
||||
def fast_synch(self) -> typing.Optional[builtins.bool]: ...
|
||||
@fast_synch.setter
|
||||
def fast_synch(self, value: typing.Optional[builtins.bool]) -> None: ...
|
||||
@staticmethod
|
||||
def default() -> AppSettings:
|
||||
r"""
|
||||
Create default instance.
|
||||
"""
|
||||
@staticmethod
|
||||
def from_env_only() -> AppSettings:
|
||||
r"""
|
||||
Create only from environment variables.
|
||||
"""
|
||||
@staticmethod
|
||||
def resolve(args: AppArgs, bootstrap: BootstrapSettings) -> AppSettings: ...
|
||||
def to_bytes(self) -> builtins.list[builtins.int]: ...
|
||||
@staticmethod
|
||||
def from_bytes(bytes: typing.Sequence[builtins.int]) -> AppSettings: ...
|
||||
def __reduce__(self) -> tuple[typing.Any, tuple]: ...
|
||||
|
||||
@typing.final
|
||||
class BootstrapArgs:
|
||||
r"""
|
||||
Arguments that are needed to resolve bootstrap settings.
|
||||
|
||||
These values are resolved before `config.toml` can be loaded. For example, the
|
||||
`config.toml` path itself depends on these values, so these arguments cannot be
|
||||
specified by `config.toml`.
|
||||
|
||||
By default, any path-like argument goes here, but it can be moved to
|
||||
[`AppArgs`](crate::config::app::AppArgs) if it no longer participates in bootstrap
|
||||
resolution.
|
||||
"""
|
||||
@property
|
||||
def exo_home(self) -> typing.Optional[pathlib.Path]: ...
|
||||
@exo_home.setter
|
||||
def exo_home(self, value: typing.Optional[builtins.str | os.PathLike | pathlib.Path]) -> None: ...
|
||||
@property
|
||||
def default_models_dir(self) -> typing.Optional[pathlib.Path]: ...
|
||||
@default_models_dir.setter
|
||||
def default_models_dir(self, value: typing.Optional[builtins.str | os.PathLike | pathlib.Path]) -> None: ...
|
||||
@property
|
||||
def models_read_only_dirs(self) -> typing.Optional[builtins.list[pathlib.Path]]: ...
|
||||
@models_read_only_dirs.setter
|
||||
def models_read_only_dirs(self, value: typing.Optional[typing.Sequence[builtins.str | os.PathLike | pathlib.Path]]) -> None: ...
|
||||
@property
|
||||
def models_dirs(self) -> typing.Optional[builtins.list[pathlib.Path]]: ...
|
||||
@models_dirs.setter
|
||||
def models_dirs(self, value: typing.Optional[typing.Sequence[builtins.str | os.PathLike | pathlib.Path]]) -> None: ...
|
||||
@property
|
||||
def config_file(self) -> typing.Optional[pathlib.Path]: ...
|
||||
@config_file.setter
|
||||
def config_file(self, value: typing.Optional[builtins.str | os.PathLike | pathlib.Path]) -> None: ...
|
||||
|
||||
@typing.final
|
||||
class BootstrapSettings:
|
||||
@property
|
||||
def exo_home(self) -> ExoHome: ...
|
||||
@exo_home.setter
|
||||
def exo_home(self, value: ExoHome) -> None: ...
|
||||
@property
|
||||
def models_dirs(self) -> ModelsDirs: ...
|
||||
@models_dirs.setter
|
||||
def models_dirs(self, value: ModelsDirs) -> None: ...
|
||||
@property
|
||||
def log_files(self) -> LogFiles: ...
|
||||
@log_files.setter
|
||||
def log_files(self, value: LogFiles) -> None: ...
|
||||
@property
|
||||
def pid_file(self) -> pathlib.Path: ...
|
||||
@pid_file.setter
|
||||
def pid_file(self, value: builtins.str | os.PathLike | pathlib.Path) -> None: ...
|
||||
@property
|
||||
def node_zid(self) -> pathlib.Path: ...
|
||||
@node_zid.setter
|
||||
def node_zid(self, value: builtins.str | os.PathLike | pathlib.Path) -> None: ...
|
||||
@property
|
||||
def config_file(self) -> pathlib.Path: ...
|
||||
@config_file.setter
|
||||
def config_file(self, value: builtins.str | os.PathLike | pathlib.Path) -> None: ...
|
||||
@property
|
||||
def custom_model_cards_dir(self) -> pathlib.Path: ...
|
||||
@custom_model_cards_dir.setter
|
||||
def custom_model_cards_dir(self, value: builtins.str | os.PathLike | pathlib.Path) -> None: ...
|
||||
@property
|
||||
def event_log_dir(self) -> pathlib.Path: ...
|
||||
@event_log_dir.setter
|
||||
def event_log_dir(self, value: builtins.str | os.PathLike | pathlib.Path) -> None: ...
|
||||
@property
|
||||
def image_cache_dir(self) -> pathlib.Path: ...
|
||||
@image_cache_dir.setter
|
||||
def image_cache_dir(self, value: builtins.str | os.PathLike | pathlib.Path) -> None: ...
|
||||
@property
|
||||
def tracing_cache_dir(self) -> pathlib.Path: ...
|
||||
@tracing_cache_dir.setter
|
||||
def tracing_cache_dir(self, value: builtins.str | os.PathLike | pathlib.Path) -> None: ...
|
||||
@staticmethod
|
||||
def default() -> BootstrapSettings:
|
||||
r"""
|
||||
Create default instance
|
||||
"""
|
||||
@staticmethod
|
||||
def from_env_only() -> BootstrapSettings:
|
||||
r"""
|
||||
Create only from env-variables
|
||||
"""
|
||||
@staticmethod
|
||||
def resolve(args: BootstrapArgs) -> BootstrapSettings: ...
|
||||
def to_bytes(self) -> builtins.list[builtins.int]: ...
|
||||
@staticmethod
|
||||
def from_bytes(bytes: typing.Sequence[builtins.int]) -> BootstrapSettings: ...
|
||||
def __reduce__(self) -> tuple[typing.Any, tuple]: ...
|
||||
|
||||
@typing.final
|
||||
class CliArgs:
|
||||
@property
|
||||
def force_master(self) -> builtins.bool: ...
|
||||
@force_master.setter
|
||||
def force_master(self, value: builtins.bool) -> None: ...
|
||||
@property
|
||||
def api_enabled(self) -> builtins.bool: ...
|
||||
@api_enabled.setter
|
||||
def api_enabled(self, value: builtins.bool) -> None: ...
|
||||
@property
|
||||
def api_port(self) -> builtins.int: ...
|
||||
@api_port.setter
|
||||
def api_port(self, value: builtins.int) -> None: ...
|
||||
@property
|
||||
def worker_enabled(self) -> builtins.bool: ...
|
||||
@worker_enabled.setter
|
||||
def worker_enabled(self, value: builtins.bool) -> None: ...
|
||||
@property
|
||||
def downloads_enabled(self) -> builtins.bool: ...
|
||||
@downloads_enabled.setter
|
||||
def downloads_enabled(self, value: builtins.bool) -> None: ...
|
||||
@property
|
||||
def legacy_daemon(self) -> builtins.bool: ...
|
||||
@legacy_daemon.setter
|
||||
def legacy_daemon(self, value: builtins.bool) -> None: ...
|
||||
@property
|
||||
def namespace(self) -> builtins.str: ...
|
||||
@namespace.setter
|
||||
def namespace(self, value: builtins.str) -> None: ...
|
||||
@property
|
||||
def zenoh_port(self) -> builtins.int: ...
|
||||
@zenoh_port.setter
|
||||
def zenoh_port(self, value: builtins.int) -> None: ...
|
||||
@property
|
||||
def discovery_port(self) -> builtins.int: ...
|
||||
@discovery_port.setter
|
||||
def discovery_port(self, value: builtins.int) -> None: ...
|
||||
@property
|
||||
def bootstrap(self) -> BootstrapArgs: ...
|
||||
@bootstrap.setter
|
||||
def bootstrap(self, value: BootstrapArgs) -> None: ...
|
||||
@property
|
||||
def app(self) -> AppArgs: ...
|
||||
@app.setter
|
||||
def app(self, value: AppArgs) -> None: ...
|
||||
@property
|
||||
def rejected(self) -> RejectedArgs: ...
|
||||
@rejected.setter
|
||||
def rejected(self, value: RejectedArgs) -> None: ...
|
||||
@staticmethod
|
||||
def from_env_only() -> CliArgs:
|
||||
r"""
|
||||
Create only from env-variables
|
||||
"""
|
||||
@staticmethod
|
||||
def parse_from(argv: typing.Sequence[builtins.str]) -> CliArgs: ...
|
||||
@staticmethod
|
||||
def parse() -> CliArgs: ...
|
||||
def to_bytes(self) -> builtins.list[builtins.int]: ...
|
||||
@staticmethod
|
||||
def from_bytes(bytes: typing.Sequence[builtins.int]) -> CliArgs: ...
|
||||
def __reduce__(self) -> tuple[typing.Any, tuple]: ...
|
||||
|
||||
@typing.final
|
||||
class ExoHome:
|
||||
@property
|
||||
def config(self) -> pathlib.Path: ...
|
||||
@config.setter
|
||||
def config(self, value: builtins.str | os.PathLike | pathlib.Path) -> None: ...
|
||||
@property
|
||||
def data(self) -> pathlib.Path: ...
|
||||
@data.setter
|
||||
def data(self, value: builtins.str | os.PathLike | pathlib.Path) -> None: ...
|
||||
@property
|
||||
def cache(self) -> pathlib.Path: ...
|
||||
@cache.setter
|
||||
def cache(self, value: builtins.str | os.PathLike | pathlib.Path) -> None: ...
|
||||
|
||||
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 LogFiles:
|
||||
@property
|
||||
def exo_log_dir(self) -> pathlib.Path: ...
|
||||
@exo_log_dir.setter
|
||||
def exo_log_dir(self, value: builtins.str | os.PathLike | pathlib.Path) -> None: ...
|
||||
@property
|
||||
def exo_log(self) -> pathlib.Path: ...
|
||||
@exo_log.setter
|
||||
def exo_log(self, value: builtins.str | os.PathLike | pathlib.Path) -> None: ...
|
||||
@property
|
||||
def exo_runner_log_dir(self) -> pathlib.Path: ...
|
||||
@exo_runner_log_dir.setter
|
||||
def exo_runner_log_dir(self, value: builtins.str | os.PathLike | pathlib.Path) -> None: ...
|
||||
@property
|
||||
def exo_runner_stdout_log(self) -> pathlib.Path: ...
|
||||
@exo_runner_stdout_log.setter
|
||||
def exo_runner_stdout_log(self, value: builtins.str | os.PathLike | pathlib.Path) -> None: ...
|
||||
@property
|
||||
def exo_runner_stderr_log(self) -> pathlib.Path: ...
|
||||
@exo_runner_stderr_log.setter
|
||||
def exo_runner_stderr_log(self, value: builtins.str | os.PathLike | pathlib.Path) -> None: ...
|
||||
|
||||
@typing.final
|
||||
class ModelsDirs:
|
||||
@property
|
||||
def default_models_dir(self) -> pathlib.Path: ...
|
||||
@default_models_dir.setter
|
||||
def default_models_dir(self, value: builtins.str | os.PathLike | pathlib.Path) -> None: ...
|
||||
@property
|
||||
def models_read_only_dirs(self) -> builtins.list[pathlib.Path]: ...
|
||||
@models_read_only_dirs.setter
|
||||
def models_read_only_dirs(self, value: typing.Sequence[builtins.str | os.PathLike | pathlib.Path]) -> None: ...
|
||||
@property
|
||||
def models_dirs(self) -> builtins.list[pathlib.Path]: ...
|
||||
@models_dirs.setter
|
||||
def models_dirs(self, value: typing.Sequence[builtins.str | os.PathLike | pathlib.Path]) -> None: ...
|
||||
|
||||
@typing.final
|
||||
class NetworkingHandle:
|
||||
@staticmethod
|
||||
def new(identity: builtins.str, namespace: builtins.str, listen_port: builtins.int, discovery_service_port: builtins.int) -> NetworkingHandle: ...
|
||||
async def recv(self) -> FromSwarm: ...
|
||||
async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
|
||||
r"""
|
||||
Subscribe to a `GossipSub` topic.
|
||||
|
||||
Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
|
||||
"""
|
||||
async def gossipsub_unsubscribe(self, topic: builtins.str) -> builtins.bool:
|
||||
r"""
|
||||
Unsubscribes from a `GossipSub` topic.
|
||||
|
||||
Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
|
||||
"""
|
||||
async def gossipsub_publish(self, topic: builtins.str, data: bytes) -> None:
|
||||
r"""
|
||||
Publishes a message with multiple topics to the `GossipSub` network.
|
||||
|
||||
If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
|
||||
"""
|
||||
|
||||
@typing.final
|
||||
class Pidfile:
|
||||
r"""
|
||||
A PID file protected with a lock.
|
||||
|
||||
An instance of `Pidfile` can be used to manage a PID file: create it,
|
||||
lock it, detect already running daemons. It is backed by [`pidfile`]
|
||||
functions of `libbsd`/`libutil` which use `flopen` to lock the PID
|
||||
file.
|
||||
|
||||
When a PID file is created, the process ID of the current process is
|
||||
*not* written there, making it possible to lock the PID file before
|
||||
forking and only write the ID of the forked process when it is ready.
|
||||
|
||||
The PID file is deleted automatically when the `Pidfile` comes out of
|
||||
the scope. To close the PID file without deleting it, for example, in
|
||||
the parent process of a forked daemon, call `close()`.
|
||||
|
||||
[`exit`]: https://doc.rust-lang.org/std/process/fn.exit.html
|
||||
[`pidfile`]: https://linux.die.net/man/3/pidfile
|
||||
[`daemon`(3)]: https://linux.die.net/man/3/daemon
|
||||
"""
|
||||
def __new__(cls, path: builtins.str | os.PathLike | pathlib.Path, mode: builtins.int) -> Pidfile:
|
||||
r"""
|
||||
Creates a new PID file and locks it.
|
||||
|
||||
If the PID file cannot be locked, returns `PidfileError::AlreadyRunning` with
|
||||
a PID of the already running process, or `None` if no PID has been written to
|
||||
the PID file yet.
|
||||
"""
|
||||
def write(self) -> None:
|
||||
r"""
|
||||
Writes the current process ID to the PID file.
|
||||
|
||||
The file is truncated before writing.
|
||||
"""
|
||||
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: ...
|
||||
|
||||
@typing.final
|
||||
class RejectedArgs:
|
||||
r"""
|
||||
Rejected arguments go here.
|
||||
|
||||
# Important
|
||||
- Make sure all are `hide = true` so it won't appear in `--help`
|
||||
- Make sure all are [`Option<T>`] so them being missing doesn't cause issues
|
||||
"""
|
||||
@property
|
||||
def bootstrap_peers(self) -> typing.Optional[builtins.list[builtins.str]]: ...
|
||||
@bootstrap_peers.setter
|
||||
def bootstrap_peers(self, value: typing.Optional[typing.Sequence[builtins.str]]) -> None: ...
|
||||
@property
|
||||
def libp2p_port(self) -> typing.Optional[builtins.int]: ...
|
||||
@libp2p_port.setter
|
||||
def libp2p_port(self, value: typing.Optional[builtins.int]) -> None: ...
|
||||
@property
|
||||
def libp2p_namespace(self) -> typing.Optional[builtins.str]: ...
|
||||
@libp2p_namespace.setter
|
||||
def libp2p_namespace(self, value: typing.Optional[builtins.str]) -> None: ...
|
||||
@property
|
||||
def zenoh_namespace(self) -> typing.Optional[builtins.str]: ...
|
||||
@zenoh_namespace.setter
|
||||
def zenoh_namespace(self, value: typing.Optional[builtins.str]) -> None: ...
|
||||
@property
|
||||
def enable_image_models(self) -> typing.Optional[builtins.bool]: ...
|
||||
@enable_image_models.setter
|
||||
def enable_image_models(self, value: typing.Optional[builtins.bool]) -> None: ...
|
||||
@property
|
||||
def enable_disaggregation(self) -> typing.Optional[builtins.bool]: ...
|
||||
@enable_disaggregation.setter
|
||||
def enable_disaggregation(self, value: typing.Optional[builtins.bool]) -> None: ...
|
||||
@property
|
||||
def no_fast_synch(self) -> typing.Optional[builtins.bool]: ...
|
||||
@no_fast_synch.setter
|
||||
def no_fast_synch(self, value: typing.Optional[builtins.bool]) -> None: ...
|
||||
@property
|
||||
def verbose(self) -> typing.Optional[builtins.bool]: ...
|
||||
@verbose.setter
|
||||
def verbose(self, value: typing.Optional[builtins.bool]) -> None: ...
|
||||
|
||||
@typing.final
|
||||
class VerbosityFilter(enum.Enum):
|
||||
r"""
|
||||
Verbosity level used by EXO's logger.
|
||||
"""
|
||||
Off = ...
|
||||
Error = ...
|
||||
Warn = ...
|
||||
Info = ...
|
||||
Debug = ...
|
||||
Trace = ...
|
||||
|
||||
def __lt__(self, other: object, /) -> bool: ...
|
||||
def __le__(self, other: object, /) -> bool: ...
|
||||
def __gt__(self, other: object, /) -> bool: ...
|
||||
def __ge__(self, other: object, /) -> bool: ...
|
||||
|
||||
@@ -3,27 +3,31 @@ requires = ["maturin>=1.0,<2.0"]
|
||||
build-backend = "maturin"
|
||||
|
||||
[project]
|
||||
name = "exo_pyo3_bindings"
|
||||
version = "0.2.10"
|
||||
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"
|
||||
features = ["pyo3/extension-module", "pyo3/experimental-async"]
|
||||
module-name = "exo_rs"
|
||||
features = ["pyo3/extension-module", "pyo3/serde", "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" }]
|
||||
@@ -1,4 +1,4 @@
|
||||
//! SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
|
||||
//! SEE: https://pyo3.rs/v0.28.3/async-await#detaching-from-the-interpreter-across-await
|
||||
//!
|
||||
|
||||
use pin_project::pin_project;
|
||||
@@ -9,7 +9,7 @@ use std::{
|
||||
task::{Context, Poll},
|
||||
};
|
||||
|
||||
/// SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
|
||||
/// SEE: https://pyo3.rs/v0.28.3/async-await#detaching-from-the-interpreter-across-await
|
||||
#[pin_project]
|
||||
#[repr(transparent)]
|
||||
pub(crate) struct AllowThreads<F>(#[pin] F);
|
||||
@@ -0,0 +1,27 @@
|
||||
use clap::Parser;
|
||||
use exo_rs::config::app::AppSettings;
|
||||
use exo_rs::config::bootstrap::BootstrapSettings;
|
||||
use exo_rs::config::cli::CliArgs;
|
||||
use pyo3::Python;
|
||||
use pyo3_stub_gen::Result;
|
||||
|
||||
fn main() -> Result<()> {
|
||||
Python::initialize();
|
||||
|
||||
Python::attach(|py| -> Result<()> {
|
||||
let a = CliArgs::parse();
|
||||
println!("{a:?}\n");
|
||||
let b = BootstrapSettings::resolve(py, &a.bootstrap.borrow(py))?;
|
||||
println!("{b:?}\n");
|
||||
let app = AppSettings::resolve(&a.app.borrow(py), &b)?;
|
||||
println!("{app:?}\n");
|
||||
Ok(())
|
||||
})?;
|
||||
|
||||
// return Ok(());
|
||||
|
||||
env_logger::Builder::from_env(env_logger::Env::default().filter_or("RUST_LOG", "info")).init();
|
||||
let stub = exo_rs::stub_info()?;
|
||||
stub.generate()?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,217 @@
|
||||
use crate::config::bootstrap::BootstrapSettings;
|
||||
use crate::config::cli::CliArgs;
|
||||
use crate::config::{VerbosityFilter, default};
|
||||
use crate::ext::ResultExt;
|
||||
use crate::pickle_reduce;
|
||||
use clap::{
|
||||
ArgAction,
|
||||
builder::{BoolishValueParser, TypedValueParser},
|
||||
};
|
||||
use figment::Figment;
|
||||
use figment::providers::{Format, Serialized, Toml};
|
||||
use pyo3::prelude::{PyModule, PyModuleMethods};
|
||||
use pyo3::types::PyTuple;
|
||||
use pyo3::{Bound, PyAny, PyResult, Python, pyclass, pymethods};
|
||||
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_with::skip_serializing_none;
|
||||
|
||||
/// Arguments that participate in application settings resolution.
|
||||
///
|
||||
/// These values may come from defaults, `config.toml`, environment variables, or
|
||||
/// CLI arguments. Unlike [`BootstrapArgs`](crate::config::bootstrap::BootstrapArgs),
|
||||
/// they do not participate in finding or loading `config.toml`.
|
||||
///
|
||||
/// # Important
|
||||
/// - Make sure all fields are [`Option<T>`] so they can be layered with other
|
||||
/// settings sources.
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass(from_py_object)]
|
||||
#[skip_serializing_none]
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize, clap::Args)]
|
||||
#[command(about = None, long_about = None)]
|
||||
pub struct AppArgs {
|
||||
#[arg(
|
||||
short = 'q',
|
||||
long = "quiet",
|
||||
action = ArgAction::SetTrue,
|
||||
conflicts_with = "verbosity",
|
||||
help = "Only show error logs (alias for --verbosity=error)"
|
||||
)]
|
||||
#[serde(skip)]
|
||||
pub verbosity_off: bool,
|
||||
#[arg(
|
||||
short = 'v',
|
||||
long,
|
||||
env = "EXO_VERBOSITY",
|
||||
value_enum,
|
||||
default_value_if("verbosity_off", "true", Some("error")),
|
||||
value_name = "LEVEL",
|
||||
conflicts_with = "verbosity_off",
|
||||
help = "Verbosity filter of the application"
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub verbosity: Option<VerbosityFilter>,
|
||||
|
||||
// this parser cannot use the default boolean parser + ArgAction::SetFalse
|
||||
// since it needs to logically invert --no-batch and EXO_NO_BATCH
|
||||
#[arg(
|
||||
long = "no-batch",
|
||||
env = "EXO_NO_BATCH",
|
||||
num_args = 0..=1,
|
||||
require_equals = true,
|
||||
default_missing_value = "true",
|
||||
value_parser = BoolishValueParser::new().map(|no_batch| !no_batch),
|
||||
value_name = "BOOL",
|
||||
help = "Disable continuous batching, use sequential generation"
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub continuous_batching_enabled: Option<bool>,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "EXO_MAX_CONCURRENT_REQUESTS",
|
||||
value_parser = clap::value_parser!(u16).range(1..),
|
||||
value_name = "NUM",
|
||||
help = "Maximum number of concurrent generation requests per runner"
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub max_concurrent_requests: Option<u16>,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "EXO_OFFLINE",
|
||||
num_args = 0..=1,
|
||||
require_equals = true,
|
||||
default_missing_value = "true",
|
||||
value_name = "BOOL",
|
||||
help = "Run in offline/air-gapped mode: skip internet checks, use only pre-staged local models"
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub offline: Option<bool>,
|
||||
|
||||
#[arg(
|
||||
long = "enable-image-models",
|
||||
env = "EXO_IMAGE_MODELS_ENABLED",
|
||||
num_args = 0..=1,
|
||||
require_equals = true,
|
||||
default_missing_value = "true",
|
||||
value_name = "BOOL",
|
||||
help = "Enable image model support"
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub image_models_enabled: Option<bool>,
|
||||
|
||||
#[arg(
|
||||
long = "enable-tracing",
|
||||
env = "EXO_TRACING_ENABLED",
|
||||
num_args = 0..=1,
|
||||
require_equals = true,
|
||||
default_missing_value = "true",
|
||||
value_name = "BOOL",
|
||||
help = "Enable distributed tracing for performance analysis"
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub tracing_enabled: Option<bool>,
|
||||
|
||||
#[arg(
|
||||
long = "enable-disaggregation",
|
||||
env = "EXO_DISAGGREGATION_ENABLED",
|
||||
num_args = 0..=1,
|
||||
require_equals = true,
|
||||
default_missing_value = "true",
|
||||
value_name = "BOOL",
|
||||
help = "Enable prefill/decode disaggregation"
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub disaggregation_enabled: Option<bool>,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "EXO_FAST_SYNCH",
|
||||
value_name = "BOOL",
|
||||
help = "Force MLX FAST_SYNCH on/off (for JACCL backend)"
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub fast_synch: Option<bool>,
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass(module = "exo_rs", from_py_object)]
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct AppSettings {
|
||||
#[pyo3(get, set)]
|
||||
pub verbosity: VerbosityFilter,
|
||||
#[pyo3(get, set)]
|
||||
pub continuous_batching_enabled: bool,
|
||||
#[pyo3(get, set)]
|
||||
pub max_concurrent_requests: u16,
|
||||
#[pyo3(get, set)]
|
||||
pub offline: bool,
|
||||
#[pyo3(get, set)]
|
||||
pub image_models_enabled: bool,
|
||||
#[pyo3(get, set)]
|
||||
pub tracing_enabled: bool,
|
||||
#[pyo3(get, set)]
|
||||
pub disaggregation_enabled: bool,
|
||||
#[pyo3(get, set)]
|
||||
pub fast_synch: Option<bool>,
|
||||
}
|
||||
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl AppSettings {
|
||||
/// Create default instance.
|
||||
#[staticmethod]
|
||||
#[pyo3(name = "default")]
|
||||
pub fn py_default(py: Python<'_>) -> PyResult<Self> {
|
||||
let bootstrap = BootstrapSettings::py_default(py)?;
|
||||
let args = AppArgs::default();
|
||||
Self::resolve(&args, &bootstrap)
|
||||
}
|
||||
|
||||
/// Create only from environment variables.
|
||||
#[staticmethod]
|
||||
pub fn from_env_only(py: Python<'_>) -> PyResult<Self> {
|
||||
let args = CliArgs::from_env_only();
|
||||
let bootstrap = BootstrapSettings::resolve(py, &*args.bootstrap.borrow(py))?;
|
||||
Self::resolve(&*args.app.borrow(py), &bootstrap)
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
pub fn resolve(args: &AppArgs, bootstrap: &BootstrapSettings) -> PyResult<Self> {
|
||||
let mut f = Figment::new()
|
||||
// merge default CLI values
|
||||
.merge(Serialized::defaults(default::APP_ARGS));
|
||||
// merge configuration file IF EXISTS!!
|
||||
if bootstrap.config_file.exists() {
|
||||
f = f.merge(Toml::file(&bootstrap.config_file))
|
||||
}
|
||||
// merge CLI args (with ENV already merged)
|
||||
f.merge(Serialized::defaults(args.clone()))
|
||||
.extract::<Self>()
|
||||
.pyerr()
|
||||
}
|
||||
|
||||
// -------- SERDE/PICKLING support --------
|
||||
|
||||
pub fn to_bytes(&self) -> PyResult<Vec<u8>> {
|
||||
postcard::to_allocvec(self).pyerr()
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
pub fn from_bytes(bytes: Vec<u8>) -> PyResult<Self> {
|
||||
postcard::from_bytes(&bytes).pyerr()
|
||||
}
|
||||
|
||||
pub fn __reduce__(slf: Bound<'_, Self>) -> PyResult<(Bound<'_, PyAny>, Bound<'_, PyTuple>)> {
|
||||
pickle_reduce(slf, "from_bytes", Self::to_bytes)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn app_submodule(m: &Bound<PyModule>) -> PyResult<()> {
|
||||
m.add_class::<AppArgs>()?;
|
||||
m.add_class::<AppSettings>()?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,326 @@
|
||||
use crate::config::cli::CliArgs;
|
||||
use crate::config::cli::{PathBufValueParserExt, parse_path};
|
||||
use crate::ext::ResultExt;
|
||||
use crate::newtype::NewPy;
|
||||
use crate::pickle_reduce;
|
||||
use pyo3::prelude::{PyModule, PyModuleMethods};
|
||||
use pyo3::types::PyTuple;
|
||||
use pyo3::{Bound, PyAny, PyResult, Python, pyclass, pymethods};
|
||||
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::io;
|
||||
use std::path::PathBuf;
|
||||
use util::VecExt;
|
||||
|
||||
/// Arguments that are needed to resolve bootstrap settings.
|
||||
///
|
||||
/// These values are resolved before `config.toml` can be loaded. For example, the
|
||||
/// `config.toml` path itself depends on these values, so these arguments cannot be
|
||||
/// specified by `config.toml`.
|
||||
///
|
||||
/// By default, any path-like argument goes here, but it can be moved to
|
||||
/// [`AppArgs`](crate::config::app::AppArgs) if it no longer participates in bootstrap
|
||||
/// resolution.
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass(from_py_object)]
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize, clap::Args)]
|
||||
#[command(about = None, long_about = None)]
|
||||
pub struct BootstrapArgs {
|
||||
#[arg(
|
||||
long,
|
||||
env = "EXO_HOME",
|
||||
value_parser = parse_path().is_dir_if_exists(),
|
||||
value_name = "PATH",
|
||||
help = "Path to Exo's home directory"
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub exo_home: Option<PathBuf>,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "EXO_DEFAULT_MODELS_DIR",
|
||||
value_parser = parse_path().is_dir_if_exists(),
|
||||
value_name = "PATH",
|
||||
help = "Default models directory; always included as first entry in writable models directories"
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub default_models_dir: Option<PathBuf>,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
value_delimiter = ':',
|
||||
env = "EXO_MODELS_READ_ONLY_DIRS",
|
||||
value_parser = parse_path().dir_exists(),
|
||||
value_name = "PATHS",
|
||||
help = "Read-only model directories (colon-separated); never written to or deleted from"
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub models_read_only_dirs: Option<Vec<PathBuf>>,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
value_delimiter = ':',
|
||||
env = "EXO_MODELS_DIRS",
|
||||
value_parser = parse_path().is_dir_if_exists(),
|
||||
value_name = "PATHS",
|
||||
help = "Writable model directories (colon-separated); default directory is always prepended"
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub models_dirs: Option<Vec<PathBuf>>,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "EXO_CONFIG_FILE",
|
||||
value_parser = parse_path().toml_file_exists(),
|
||||
value_name = "PATH",
|
||||
help = "Path to Exo's .toml config file"
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub config_file: Option<PathBuf>,
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass(module = "exo_rs", skip_from_py_object)]
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct BootstrapSettings {
|
||||
#[pyo3(get, set)]
|
||||
pub exo_home: NewPy<ExoHome>,
|
||||
|
||||
#[pyo3(get, set)]
|
||||
pub models_dirs: NewPy<ModelsDirs>,
|
||||
|
||||
#[pyo3(get, set)]
|
||||
pub log_files: NewPy<LogFiles>,
|
||||
|
||||
// other
|
||||
#[pyo3(get, set)]
|
||||
pub pid_file: PathBuf,
|
||||
|
||||
#[pyo3(get, set)]
|
||||
pub node_zid: PathBuf,
|
||||
|
||||
#[pyo3(get, set)]
|
||||
pub config_file: PathBuf,
|
||||
|
||||
#[pyo3(get, set)]
|
||||
pub custom_model_cards_dir: PathBuf,
|
||||
|
||||
#[pyo3(get, set)]
|
||||
pub event_log_dir: PathBuf,
|
||||
|
||||
#[pyo3(get, set)]
|
||||
pub image_cache_dir: PathBuf,
|
||||
|
||||
#[pyo3(get, set)]
|
||||
pub tracing_cache_dir: PathBuf,
|
||||
}
|
||||
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl BootstrapSettings {
|
||||
/// Create default instance
|
||||
#[staticmethod]
|
||||
#[pyo3(name = "default")]
|
||||
pub fn py_default(py: Python<'_>) -> PyResult<Self> {
|
||||
// resolve from env only
|
||||
Self::resolve(py, &BootstrapArgs::default())
|
||||
}
|
||||
|
||||
/// Create only from env-variables
|
||||
#[staticmethod]
|
||||
pub fn from_env_only(py: Python<'_>) -> PyResult<Self> {
|
||||
// resolve from env only
|
||||
Self::resolve(py, &*CliArgs::from_env_only().bootstrap.borrow(py))
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
pub fn resolve(py: Python<'_>, args: &BootstrapArgs) -> PyResult<Self> {
|
||||
let exo_home = ExoHome::resolve(args)?;
|
||||
let models_dirs = ModelsDirs::resolve(args, &exo_home)?;
|
||||
let log_files = LogFiles::resolve(&exo_home)?;
|
||||
|
||||
// PID file
|
||||
let pid_file = exo_home.cache.join("exo.pid");
|
||||
|
||||
// Identity (config)
|
||||
let node_zid = exo_home.cache.join("node_zid");
|
||||
let config_file = args
|
||||
.config_file
|
||||
.clone()
|
||||
.unwrap_or_else(|| exo_home.config.join("config.toml"));
|
||||
|
||||
// custom model card dirs TODO: see model_cards.py "todo"
|
||||
let custom_model_cards_dir = exo_home.data.join("custom_model_cards");
|
||||
|
||||
let event_log_dir = exo_home.data.join("event_log");
|
||||
let image_cache_dir = exo_home.cache.join("images");
|
||||
let tracing_cache_dir = exo_home.cache.join("traces");
|
||||
|
||||
Ok(Self {
|
||||
exo_home: NewPy::py_try_new_with(py, exo_home)?,
|
||||
models_dirs: NewPy::py_try_new_with(py, models_dirs)?,
|
||||
log_files: NewPy::py_try_new_with(py, log_files)?,
|
||||
pid_file,
|
||||
node_zid,
|
||||
config_file,
|
||||
custom_model_cards_dir,
|
||||
event_log_dir,
|
||||
image_cache_dir,
|
||||
tracing_cache_dir,
|
||||
})
|
||||
}
|
||||
|
||||
// -------- SERDE/PICKLING support --------
|
||||
|
||||
pub fn to_bytes(&self) -> PyResult<Vec<u8>> {
|
||||
postcard::to_allocvec(self).pyerr()
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
pub fn from_bytes(bytes: Vec<u8>) -> PyResult<Self> {
|
||||
postcard::from_bytes(&bytes).pyerr()
|
||||
}
|
||||
|
||||
pub fn __reduce__(slf: Bound<'_, Self>) -> PyResult<(Bound<'_, PyAny>, Bound<'_, PyTuple>)> {
|
||||
pickle_reduce(slf, "from_bytes", Self::to_bytes)
|
||||
}
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass(skip_from_py_object)]
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ExoHome {
|
||||
#[pyo3(get, set)]
|
||||
pub config: PathBuf,
|
||||
#[pyo3(get, set)]
|
||||
pub data: PathBuf,
|
||||
#[pyo3(get, set)]
|
||||
pub cache: PathBuf,
|
||||
}
|
||||
|
||||
impl ExoHome {
|
||||
/// Get the home directory for a specific purpose, with this precedence:
|
||||
/// 1. Prioritize `exo_home` if set
|
||||
/// 2. Fall back to "`<dir>`/exo" if specified; should be [XDG Directories] on Linux,
|
||||
/// and [Standard Directories] on macOS
|
||||
/// 3. Fall back to "$HOME/.exo" if all else fails
|
||||
///
|
||||
/// [XDG Directories]: https://specifications.freedesktop.org/basedir/latest/
|
||||
/// [Standard Directories]: https://developer.apple.com/library/archive/documentation/FileManagement/Conceptual/FileSystemProgrammingGuide/FileSystemOverview/FileSystemOverview.html#//apple_ref/doc/uid/TP40010672-CH2-SW6
|
||||
fn get_home_dir(
|
||||
exo_home: &Option<PathBuf>,
|
||||
get_dir: impl FnOnce() -> Option<PathBuf>,
|
||||
) -> io::Result<PathBuf> {
|
||||
// TODO: This splits the folders into separate ones in macOS **too**
|
||||
// so we can have persistent IDs in cache folder and avoid the copy bug Evan mentioned
|
||||
// BUT the user encountered the bug when he used "macOS time machine" or something
|
||||
// so test that the "macOS time machine" doesn't copy the cache folder
|
||||
|
||||
exo_home
|
||||
.clone()
|
||||
.or_else(|| get_dir().map(|p| p.join("exo")))
|
||||
.or_else(|| dirs::home_dir().map(|p| p.join(".exo")))
|
||||
.ok_or_else(|| {
|
||||
io::Error::new(
|
||||
io::ErrorKind::NotFound,
|
||||
"no home EXO home directory found: none specified, and $HOME directory doesn't exist",
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn resolve(args: &BootstrapArgs) -> io::Result<Self> {
|
||||
// resolve config/data/cache folders which the rest of the paths are derived from
|
||||
Ok(Self {
|
||||
config: Self::get_home_dir(&args.exo_home, dirs::config_dir)?,
|
||||
data: Self::get_home_dir(&args.exo_home, dirs::data_dir)?,
|
||||
cache: Self::get_home_dir(&args.exo_home, dirs::cache_dir)?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass(skip_from_py_object)]
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ModelsDirs {
|
||||
#[pyo3(get, set)]
|
||||
pub default_models_dir: PathBuf,
|
||||
#[pyo3(get, set)]
|
||||
pub models_read_only_dirs: Vec<PathBuf>,
|
||||
#[pyo3(get, set)]
|
||||
pub models_dirs: Vec<PathBuf>,
|
||||
}
|
||||
|
||||
impl ModelsDirs {
|
||||
pub fn resolve(args: &BootstrapArgs, exo_home: &ExoHome) -> io::Result<Self> {
|
||||
// resolve default models dir
|
||||
let default_models_dir = args
|
||||
.default_models_dir
|
||||
.clone()
|
||||
.unwrap_or_else(|| exo_home.data.join("models"));
|
||||
|
||||
// set of read-only directories
|
||||
let mut models_read_only_dirs = args.models_read_only_dirs.clone().unwrap_or_else(Vec::new);
|
||||
models_read_only_dirs.dedup_preserve_order();
|
||||
|
||||
// set of mutable directories includes default directory and excludes read-only ones
|
||||
let mut models_dirs = vec![default_models_dir.clone()];
|
||||
if let Some(ref dirs) = args.models_dirs {
|
||||
models_dirs.extend(dirs.clone())
|
||||
};
|
||||
models_dirs.dedup_preserve_order();
|
||||
models_dirs.retain(|d| !models_read_only_dirs.contains(d));
|
||||
|
||||
Ok(Self {
|
||||
default_models_dir,
|
||||
models_read_only_dirs,
|
||||
models_dirs,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass(skip_from_py_object)]
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct LogFiles {
|
||||
#[pyo3(get, set)]
|
||||
pub exo_log_dir: PathBuf,
|
||||
#[pyo3(get, set)]
|
||||
pub exo_log: PathBuf,
|
||||
#[pyo3(get, set)]
|
||||
pub exo_runner_log_dir: PathBuf,
|
||||
#[pyo3(get, set)]
|
||||
pub exo_runner_stdout_log: PathBuf,
|
||||
#[pyo3(get, set)]
|
||||
pub exo_runner_stderr_log: PathBuf,
|
||||
}
|
||||
|
||||
impl LogFiles {
|
||||
pub fn resolve(exo_home: &ExoHome) -> io::Result<Self> {
|
||||
// Exo log
|
||||
let exo_log_dir = exo_home.cache.join("exo_log");
|
||||
let exo_log = exo_log_dir.join("exo.log");
|
||||
|
||||
// Exo runner log
|
||||
let exo_runner_log_dir = exo_log_dir.join("runner_log");
|
||||
let exo_runner_stdout_log = exo_runner_log_dir.join("stdout.log");
|
||||
let exo_runner_stderr_log = exo_runner_log_dir.join("stderr.log");
|
||||
|
||||
Ok(Self {
|
||||
exo_log_dir,
|
||||
exo_log,
|
||||
exo_runner_log_dir,
|
||||
exo_runner_stdout_log,
|
||||
exo_runner_stderr_log,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub fn bootstrap_submodule(m: &Bound<PyModule>) -> PyResult<()> {
|
||||
m.add_class::<BootstrapArgs>()?;
|
||||
m.add_class::<BootstrapSettings>()?;
|
||||
m.add_class::<ExoHome>()?;
|
||||
m.add_class::<ModelsDirs>()?;
|
||||
m.add_class::<LogFiles>()?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,428 @@
|
||||
use crate::config::app::AppArgs;
|
||||
use crate::config::bootstrap::BootstrapArgs;
|
||||
use crate::config::default;
|
||||
use crate::ext::ResultExt;
|
||||
use crate::{pickle_reduce, version};
|
||||
use clap::{ArgAction, Parser};
|
||||
use pyo3::prelude::{PyAnyMethods, PyModuleMethods};
|
||||
use pyo3::types::{PyModule, PyTuple};
|
||||
use pyo3::{Bound, PyAny, PyResult, Python, pyclass, pymethods};
|
||||
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::ffi::OsString;
|
||||
|
||||
// re-export
|
||||
use crate::newtype::NewPy;
|
||||
pub use parser_impl::*;
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass(module = "exo_rs", from_py_object)]
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Parser)]
|
||||
#[command(name = "EXO", version = version::version(), about, long_about = None)]
|
||||
pub struct CliArgs {
|
||||
#[arg(
|
||||
short = 'm',
|
||||
long,
|
||||
action = ArgAction::SetTrue,
|
||||
help = "Force node to be master"
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub force_master: bool,
|
||||
|
||||
#[arg(
|
||||
long = "no-api",
|
||||
action = ArgAction::SetFalse,
|
||||
default_value_t = true,
|
||||
help = "Disable the API"
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub api_enabled: bool,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
default_value_t = default::API_PORT,
|
||||
value_name = "PORT",
|
||||
help = "Port on which the API runs"
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub api_port: u16,
|
||||
|
||||
#[arg(
|
||||
long = "no-worker",
|
||||
action = ArgAction::SetFalse,
|
||||
default_value_t = true,
|
||||
help = "Disable the worker"
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub worker_enabled: bool,
|
||||
|
||||
#[arg(
|
||||
long = "no-downloads",
|
||||
action = ArgAction::SetFalse,
|
||||
default_value_t = true,
|
||||
help = "Disable the download coordinator (node won't download models)"
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub downloads_enabled: bool,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
action = ArgAction::SetTrue,
|
||||
help = "Run as a legacy SysV-style background daemon using double-fork daemonization"
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub legacy_daemon: bool,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "EXO_NAMESPACE",
|
||||
default_value_t = default::NAMESPACE(),
|
||||
value_name = "STRING",
|
||||
help = "Discovery namespace, nodes with different namespaces will not connect"
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub namespace: String,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
default_value_t = default::ZENOH_PORT,
|
||||
value_name = "PORT",
|
||||
help = "Fixed TCP port for zenoh to listen"
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub zenoh_port: u16,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
default_value_t = default::DISCOVERY_PORT,
|
||||
value_name = "PORT",
|
||||
help = "Fixed UDP port for the discovery service"
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub discovery_port: u16,
|
||||
|
||||
// -------- FLATTENED SUBCOMMANDS --------
|
||||
#[command(flatten)]
|
||||
#[pyo3(get, set)]
|
||||
pub bootstrap: NewPy<BootstrapArgs>,
|
||||
|
||||
#[command(flatten)]
|
||||
#[pyo3(get, set)]
|
||||
pub app: NewPy<AppArgs>,
|
||||
|
||||
#[command(flatten)]
|
||||
#[pyo3(get, set)]
|
||||
pub rejected: NewPy<RejectedArgs>,
|
||||
}
|
||||
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl CliArgs {
|
||||
/// Create only from env-variables
|
||||
#[staticmethod]
|
||||
pub fn from_env_only() -> Self {
|
||||
// parse only from env - no arguments
|
||||
CliArgs::parse_from(&["exo"])
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
#[pyo3(name = "parse_from")]
|
||||
pub fn py_parse_from(argv: Vec<OsString>) -> Self {
|
||||
CliArgs::parse_from(argv)
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
#[pyo3(name = "parse")]
|
||||
pub fn py_parse(py: Python<'_>) -> PyResult<Self> {
|
||||
// the correct CLI args to parse is `sys.argv`, because the original ones
|
||||
// (i.e. `sys.orig_argv`) may contain extra arguments which would mess up parsing
|
||||
let argv: Vec<OsString> = PyModule::import(py, "sys")?.getattr("argv")?.extract()?;
|
||||
Ok(CliArgs::parse_from(argv))
|
||||
}
|
||||
|
||||
// -------- SERDE/PICKLING support --------
|
||||
|
||||
pub fn to_bytes(&self) -> PyResult<Vec<u8>> {
|
||||
postcard::to_allocvec(self).pyerr()
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
pub fn from_bytes(bytes: Vec<u8>) -> PyResult<Self> {
|
||||
postcard::from_bytes(&bytes).pyerr()
|
||||
}
|
||||
|
||||
pub fn __reduce__(slf: Bound<'_, Self>) -> PyResult<(Bound<'_, PyAny>, Bound<'_, PyTuple>)> {
|
||||
pickle_reduce(slf, "from_bytes", Self::to_bytes)
|
||||
}
|
||||
}
|
||||
|
||||
/// Rejected arguments go here.
|
||||
///
|
||||
/// # Important
|
||||
/// - Make sure all are `hide = true` so it won't appear in `--help`
|
||||
/// - Make sure all are [`Option<T>`] so them being missing doesn't cause issues
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass(from_py_object)]
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, clap::Args)]
|
||||
#[command(about = None, long_about = None)]
|
||||
pub struct RejectedArgs {
|
||||
// -------- temporarily unavailable --------
|
||||
#[arg(
|
||||
long,
|
||||
env = "EXO_BOOTSTRAP_PEERS",
|
||||
value_delimiter = ',',
|
||||
value_name = "MULTIADDRS",
|
||||
help = "Comma-separated libp2p multiaddrs to dial on startup",
|
||||
hide = true,
|
||||
value_parser = Rejected::<String>::unavailable(
|
||||
Some("--bootstrap-peers"), None, Some("EXO_BOOTSTRAP_PEERS"),
|
||||
"bootstrap peers are temporarily removed",
|
||||
)
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub bootstrap_peers: Option<Vec<String>>,
|
||||
|
||||
// -------- deprecated --------
|
||||
#[arg(
|
||||
long, value_name = "PORT", hide = true,
|
||||
value_parser = Rejected::<u16>::deprecated(
|
||||
Some("--libp2p-port"), None, None,
|
||||
Some("--zenoh-port"), None, None,
|
||||
)
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub libp2p_port: Option<u16>,
|
||||
|
||||
#[arg(
|
||||
long = Rejected::fake_name(),
|
||||
env = "EXO_LIBP2P_NAMESPACE", value_name = "STRING", hide = true,
|
||||
value_parser = Rejected::<String>::deprecated(
|
||||
None, None, Some("EXO_LIBP2P_NAMESPACE"),
|
||||
Some("--namespace"), None, Some("EXO_NAMESPACE"),
|
||||
)
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub libp2p_namespace: Option<String>,
|
||||
|
||||
#[arg(
|
||||
long = Rejected::fake_name(),
|
||||
env = "EXO_ZENOH_NAMESPACE", value_name = "STRING", hide = true,
|
||||
value_parser = Rejected::<String>::deprecated(
|
||||
None, None, Some("EXO_ZENOH_NAMESPACE"),
|
||||
Some("--namespace"), None, Some("EXO_NAMESPACE"),
|
||||
)
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub zenoh_namespace: Option<String>,
|
||||
|
||||
#[arg(
|
||||
long = Rejected::fake_name(),
|
||||
env = "EXO_ENABLE_IMAGE_MODELS", value_name = "BOOL", hide = true,
|
||||
value_parser = Rejected::<bool>::deprecated(
|
||||
None, None, Some("EXO_ENABLE_IMAGE_MODELS"),
|
||||
Some("--enable-image-models"), None, Some("EXO_IMAGE_MODELS_ENABLED"),
|
||||
)
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub enable_image_models: Option<bool>,
|
||||
|
||||
#[arg(
|
||||
long = Rejected::fake_name(),
|
||||
env = "ENABLE_DISAGGREGATION", value_name = "BOOL", hide = true,
|
||||
value_parser = Rejected::<bool>::deprecated(
|
||||
None, None, Some("ENABLE_DISAGGREGATION"),
|
||||
Some("--enable-disaggregation"), None, Some("EXO_DISAGGREGATION_ENABLED"),
|
||||
)
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub enable_disaggregation: Option<bool>,
|
||||
|
||||
#[arg(
|
||||
long = "no-fast-synch", hide = true,
|
||||
num_args = 0..=1, default_missing_value = "true",
|
||||
value_parser = Rejected::<bool>::deprecated(
|
||||
Some("--no-fast-synch"), None, None,
|
||||
Some("--fast-synch=false"), None, None,
|
||||
)
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub no_fast_synch: Option<bool>,
|
||||
|
||||
#[arg(
|
||||
long = "verbose", hide = true,
|
||||
num_args = 0..=1, default_missing_value = "true",
|
||||
value_parser = Rejected::<bool>::deprecated(
|
||||
Some("--verbose"), None, None,
|
||||
Some("--verbosity=debug"), None, Some("EXO_VERBOSITY=debug"),
|
||||
)
|
||||
)]
|
||||
#[pyo3(get, set)]
|
||||
pub verbose: Option<bool>,
|
||||
}
|
||||
|
||||
mod parser_impl {
|
||||
use clap::builder::PathBufValueParser;
|
||||
use clap::builder::TypedValueParser;
|
||||
use itertools::Itertools;
|
||||
use rand::distr::SampleString;
|
||||
use std::error::Error;
|
||||
use std::ffi::OsStr;
|
||||
use std::fs;
|
||||
use std::marker::PhantomData;
|
||||
use std::path::PathBuf;
|
||||
use std::str::FromStr;
|
||||
use util::path::{PathExt, resolve_path};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Rejected<T = ()> {
|
||||
message: String,
|
||||
_ty: PhantomData<T>,
|
||||
}
|
||||
|
||||
impl Rejected {
|
||||
pub fn fake_name() -> String {
|
||||
rand::distr::Alphanumeric.sample_string(&mut rand::rng(), 32)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Rejected<T> {
|
||||
#[inline(always)]
|
||||
pub fn new(message: impl Into<String>) -> Self {
|
||||
let mut message = message.into();
|
||||
if !message.ends_with('\n') {
|
||||
message.push('\n');
|
||||
}
|
||||
Self {
|
||||
message,
|
||||
_ty: PhantomData,
|
||||
}
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
pub fn deprecated(
|
||||
old_long: Option<&str>,
|
||||
old_short: Option<&str>,
|
||||
old_env: Option<&str>,
|
||||
new_long: Option<&str>,
|
||||
new_short: Option<&str>,
|
||||
new_env: Option<&str>,
|
||||
) -> Self {
|
||||
let old_names = vec![old_short, old_long, old_env]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.join("/");
|
||||
let new_names = vec![new_short, new_long, new_env]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.join("/");
|
||||
Self::new(format!(
|
||||
"the argument {old_names} is deprecated{}",
|
||||
if new_names.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
format!("; use {new_names} instead")
|
||||
}
|
||||
))
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
pub fn unavailable(
|
||||
long: Option<&str>,
|
||||
short: Option<&str>,
|
||||
env: Option<&str>,
|
||||
reason: impl AsRef<str>,
|
||||
) -> Self {
|
||||
let names = vec![short, long, env].into_iter().flatten().join("/");
|
||||
Self::new(format!(
|
||||
"the argument {names} is unavailable: {}",
|
||||
reason.as_ref()
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> TypedValueParser for Rejected<T>
|
||||
where
|
||||
T: Clone + Send + Sync + 'static,
|
||||
{
|
||||
type Value = T;
|
||||
fn parse_ref(
|
||||
&self,
|
||||
cmd: &clap::Command,
|
||||
_arg: Option<&clap::Arg>,
|
||||
_value: &OsStr,
|
||||
) -> Result<Self::Value, clap::Error> {
|
||||
Err(clap::Error::raw(
|
||||
clap::error::ErrorKind::ValueValidation,
|
||||
self.message.clone(),
|
||||
)
|
||||
.with_cmd(cmd))
|
||||
}
|
||||
}
|
||||
|
||||
/// Default path parser that should be used to ensure paths are
|
||||
/// resolved to absolute paths before being further processed.
|
||||
pub fn parse_path() -> impl TypedValueParser<Value = PathBuf> {
|
||||
PathBufValueParser::new().try_map(resolve_path)
|
||||
}
|
||||
|
||||
// extension trait to tack on extra validation on path parsing
|
||||
pub trait PathBufValueParserExt: TypedValueParser<Value = PathBuf> {
|
||||
#[inline]
|
||||
fn canonicalize(self) -> impl TypedValueParser<Value = PathBuf> {
|
||||
self.try_map(|p| p.canonicalize())
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn dir_exists(self) -> impl TypedValueParser<Value = PathBuf> {
|
||||
self.canonicalize()
|
||||
.try_map(|p| p.try_dir_exists().map(|_| p))
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn file_exists(self) -> impl TypedValueParser<Value = PathBuf> {
|
||||
self.canonicalize()
|
||||
.try_map(|p| p.try_file_exists().map(|_| p))
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn is_dir_if_exists(self) -> impl TypedValueParser<Value = PathBuf> {
|
||||
self.try_map(|p| {
|
||||
if p.exists() {
|
||||
p.canonicalize()?.try_dir_exists().map(|_| p)
|
||||
} else {
|
||||
Ok(p)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn is_file_if_exists(self) -> impl TypedValueParser<Value = PathBuf> {
|
||||
self.try_map(|p| {
|
||||
if p.exists() {
|
||||
p.canonicalize()?.try_file_exists().map(|_| p)
|
||||
} else {
|
||||
Ok(p)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn toml_file_exists(self) -> impl TypedValueParser<Value = PathBuf> {
|
||||
self.file_exists()
|
||||
.try_map(|p| -> Result<_, Box<dyn Error + Send + Sync + 'static>> {
|
||||
let toml = fs::read_to_string(&p)?;
|
||||
let _ = toml_edit::Document::from_str(&toml)?;
|
||||
Ok(p)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: TypedValueParser<Value = PathBuf>> PathBufValueParserExt for T {}
|
||||
}
|
||||
|
||||
pub fn cli_submodule(m: &Bound<PyModule>) -> PyResult<()> {
|
||||
m.add_class::<CliArgs>()?;
|
||||
m.add_class::<RejectedArgs>()?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
use crate::config::app::app_submodule;
|
||||
use crate::config::bootstrap::bootstrap_submodule;
|
||||
use crate::config::cli::cli_submodule;
|
||||
use clap::ValueEnum;
|
||||
use pyo3::prelude::{PyModule, PyModuleMethods};
|
||||
use pyo3::{Bound, PyResult, pyclass};
|
||||
use pyo3_stub_gen::derive::gen_stub_pyclass_enum;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
pub mod app;
|
||||
pub mod bootstrap;
|
||||
pub mod cli;
|
||||
|
||||
/// Verbosity level used by EXO's logger.
|
||||
#[gen_stub_pyclass_enum]
|
||||
#[pyclass(eq, eq_int, ord, from_py_object)]
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, ValueEnum)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
#[repr(u8)]
|
||||
pub enum VerbosityFilter {
|
||||
Off = 0,
|
||||
Error,
|
||||
Warn,
|
||||
Info,
|
||||
Debug,
|
||||
Trace,
|
||||
}
|
||||
|
||||
pyo3_stub_gen::inventory::submit! {
|
||||
pyo3_stub_gen::derive::gen_methods_from_python! {
|
||||
r#"
|
||||
class VerbosityFilter:
|
||||
def __lt__(self, other: object, /) -> bool: ...
|
||||
def __le__(self, other: object, /) -> bool: ...
|
||||
def __gt__(self, other: object, /) -> bool: ...
|
||||
def __ge__(self, other: object, /) -> bool: ...
|
||||
"#
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(nonstandard_style)]
|
||||
pub mod default {
|
||||
use crate::config::VerbosityFilter;
|
||||
use crate::config::app::AppArgs;
|
||||
use crate::version;
|
||||
// ---- unclassified defaults (belonging to top-level CLI) ----
|
||||
|
||||
/// Port on which the API runs
|
||||
pub const API_PORT: u16 = 52415;
|
||||
|
||||
/// Discovery namespace, nodes with different namespaces will not connect.
|
||||
pub fn NAMESPACE() -> String {
|
||||
version::version().to_string()
|
||||
}
|
||||
|
||||
/// Fixed TCP port for zenoh to listen
|
||||
pub const ZENOH_PORT: u16 = 52414;
|
||||
|
||||
/// Fixed UDP port for the discovery service
|
||||
pub const DISCOVERY_PORT: u16 = 52413;
|
||||
|
||||
/// Default [`AppArgs`] values
|
||||
pub const APP_ARGS: AppArgs = AppArgs {
|
||||
// verbosity
|
||||
verbosity_off: false,
|
||||
verbosity: Some(VerbosityFilter::Info),
|
||||
|
||||
// rest
|
||||
continuous_batching_enabled: Some(true),
|
||||
max_concurrent_requests: Some(8),
|
||||
offline: Some(false),
|
||||
image_models_enabled: Some(false),
|
||||
tracing_enabled: Some(false),
|
||||
disaggregation_enabled: Some(false),
|
||||
fast_synch: None,
|
||||
};
|
||||
}
|
||||
|
||||
pub fn config_submodule(m: &Bound<PyModule>) -> PyResult<()> {
|
||||
bootstrap_submodule(m)?;
|
||||
app_submodule(m)?;
|
||||
cli_submodule(m)?;
|
||||
|
||||
m.add_class::<VerbosityFilter>()?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -5,23 +5,19 @@
|
||||
//!
|
||||
|
||||
mod allow_threading;
|
||||
mod ident;
|
||||
pub mod config;
|
||||
mod networking;
|
||||
pub mod newtype;
|
||||
mod pidfile;
|
||||
|
||||
use crate::ident::PyKeypair;
|
||||
use crate::config::config_submodule;
|
||||
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::prelude::{PyAnyMethods, PyModule};
|
||||
use pyo3::types::{DerefToPyAny, PyTuple};
|
||||
use pyo3::{Bound, PyAny, PyClass, 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;
|
||||
@@ -52,7 +48,7 @@ pub(crate) mod ext {
|
||||
}
|
||||
|
||||
pub trait FutureExt: Future + Sized {
|
||||
/// SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
|
||||
/// SEE: https://pyo3.rs/v0.28.3/async-await#detaching-from-the-interpreter-across-await
|
||||
fn allow_threads_py(self) -> AllowThreads<Self>
|
||||
where
|
||||
AllowThreads<Self>: Future,
|
||||
@@ -90,13 +86,27 @@ pub(crate) mod ext {
|
||||
|
||||
#[ext(pub, name = TokioRuntimeExt)]
|
||||
impl Runtime {
|
||||
#[inline(always)]
|
||||
fn spawn_with_scope<F>(&self, py: Python<'_>, future: F) -> PyResult<JoinHandle<F::Output>>
|
||||
where
|
||||
F: Future + Send + 'static,
|
||||
F::Output: Send + 'static,
|
||||
{
|
||||
let locals = pyo3_async_runtimes::tokio::get_current_locals(py)?;
|
||||
Ok(self.spawn(pyo3_async_runtimes::tokio::scope(locals, future)))
|
||||
use pyo3_async_runtimes::tokio::{get_current_locals, scope};
|
||||
let locals = get_current_locals(py)?;
|
||||
Ok(self.spawn(scope(locals, future)))
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
async fn run_with_scope<F>(&self, future: F) -> PyResult<F::Output>
|
||||
where
|
||||
F: Future + Send + 'static,
|
||||
F::Output: Send + 'static,
|
||||
{
|
||||
Python::attach(|py| self.spawn_with_scope(py, future))?
|
||||
.allow_threads_py()
|
||||
.await
|
||||
.pyerr()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -150,13 +160,77 @@ pub(crate) mod ext {
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolving the version of the python project
|
||||
pub(crate) mod version {
|
||||
use pyo3::exceptions::PyRuntimeError;
|
||||
use pyo3::prelude::PyAnyMethods;
|
||||
use pyo3::types::PyModule;
|
||||
use pyo3::{PyResult, Python};
|
||||
use std::env;
|
||||
use std::sync::OnceLock;
|
||||
|
||||
const DEFAULT_VERSION: &str = env!("CARGO_PKG_VERSION");
|
||||
static VERSION: OnceLock<String> = OnceLock::new();
|
||||
|
||||
/// Returns either the configured version of Exo (once set by [`set_version_once`])
|
||||
/// or falls back to `CARGO_PKG_VERSION` if that hasn't been configured.
|
||||
pub fn version() -> &'static str {
|
||||
VERSION.get().map_or(DEFAULT_VERSION, String::as_str)
|
||||
}
|
||||
|
||||
/// First tries to find `EXO_PKG_VERSION` env-var, falls back to calling Python
|
||||
/// `importlib.metadata.version("exo")` to resolve the version of Exo
|
||||
pub fn set_version_once(py: Python<'_>) -> PyResult<()> {
|
||||
let v = if let Ok(v) = env::var("EXO_PKG_VERSION") {
|
||||
v
|
||||
} else {
|
||||
// essentially runs:
|
||||
// ```python
|
||||
// from importlib.metadata import version
|
||||
// version("exo")
|
||||
// ```
|
||||
PyModule::import(py, "importlib.metadata")?
|
||||
.getattr("version")?
|
||||
.call1(("exo",))?
|
||||
.extract()?
|
||||
};
|
||||
|
||||
// sets version only once
|
||||
VERSION
|
||||
.set(v)
|
||||
.map_err(|_| PyRuntimeError::new_err("Cannot set exo_rs version twice".to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
/// Default `__reduce__` implementation for Pyo3 classes to support Python pickling.
|
||||
///
|
||||
/// Must provide a "from bytes" class method, a "to bytes" closure and the pyclass must have
|
||||
/// e.g. `module = "exo_rs"` set to prevent Pyo3 from putting it in `builtins.<CLASS_NAME>`
|
||||
/// (which will cause errors when pickling.)
|
||||
#[inline(always)]
|
||||
pub(crate) fn pickle_reduce<'py, T: PyClass + DerefToPyAny>(
|
||||
obj: Bound<'py, T>,
|
||||
from_bytes_method: &str,
|
||||
to_bytes: impl for<'a> FnOnce(&'a T) -> PyResult<Vec<u8>>,
|
||||
) -> PyResult<(Bound<'py, PyAny>, Bound<'py, PyTuple>)> {
|
||||
let callable = obj.getattr("__class__")?.getattr(from_bytes_method)?;
|
||||
let args = PyTuple::new(obj.py(), [to_bytes(&obj.borrow())?])?;
|
||||
|
||||
Ok((callable, args))
|
||||
}
|
||||
|
||||
/// 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();
|
||||
|
||||
// resolve version
|
||||
version::set_version_once(m.py())?;
|
||||
|
||||
// configure runtime
|
||||
let mut builder = tokio::runtime::Builder::new_multi_thread();
|
||||
builder.enable_all();
|
||||
pyo3_async_runtimes::tokio::init(builder);
|
||||
@@ -164,9 +238,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)?;
|
||||
networking_submodule(m)?;
|
||||
config_submodule(m)?;
|
||||
|
||||
// top-level constructs
|
||||
// TODO: ...
|
||||
@@ -0,0 +1,212 @@
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _, TokioRuntimeExt};
|
||||
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, 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)
|
||||
// })
|
||||
// }
|
||||
|
||||
async fn recv(&self) -> PyResult<PyFromSwarm> {
|
||||
let swarm = Arc::clone(&self.swarm);
|
||||
pyo3_async_runtimes::tokio::get_runtime()
|
||||
.run_with_scope(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)
|
||||
})
|
||||
.await?
|
||||
}
|
||||
|
||||
// ---- 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(())
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
mod new_py;
|
||||
|
||||
// re-exports
|
||||
pub use new_py::*;
|
||||
@@ -0,0 +1,277 @@
|
||||
use crate::ext::ResultExt;
|
||||
use clap::{ArgMatches, Args, CommandFactory, FromArgMatches, Parser, Subcommand};
|
||||
use pyo3::pyclass::boolean_struct::False;
|
||||
use pyo3::{
|
||||
Borrowed, FromPyObject, IntoPyObject, Py, PyAny, PyClass, PyClassInitializer, PyErr, PyResult,
|
||||
Python,
|
||||
};
|
||||
use pyo3_stub_gen::{PyStubType, TypeInfo};
|
||||
use serde::{Deserialize, Deserializer, Serialize, Serializer};
|
||||
use std::ffi::OsString;
|
||||
use std::fmt::{Debug, Display, Formatter};
|
||||
use std::ops::{Deref, DerefMut};
|
||||
use std::str::FromStr;
|
||||
|
||||
/// Wrapper around [`Py`] to provide integration with other libraries.
|
||||
#[repr(transparent)]
|
||||
pub struct NewPy<T>(Py<T>);
|
||||
|
||||
impl<T: Debug> Debug for NewPy<T>
|
||||
where
|
||||
T: PyClass,
|
||||
{
|
||||
#[inline(always)]
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
Python::attach(|py| (&*self.borrow(py)).fmt(f))
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Clone> Clone for NewPy<T> {
|
||||
#[inline(always)]
|
||||
fn clone(&self) -> Self {
|
||||
Self::new(Python::attach(|py| self.clone_ref(py)))
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: PartialEq> PartialEq for NewPy<T>
|
||||
where
|
||||
T: PyClass,
|
||||
{
|
||||
#[inline(always)]
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
Python::attach(|py| &*self.borrow(py) == &*other.borrow(py))
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Eq> Eq for NewPy<T> where T: PyClass {}
|
||||
|
||||
impl<T: Display> Display for NewPy<T>
|
||||
where
|
||||
T: PyClass,
|
||||
{
|
||||
#[inline(always)]
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
Python::attach(|py| (&*self.borrow(py)).fmt(f))
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: FromStr> FromStr for NewPy<T>
|
||||
where
|
||||
T: PyClass + Into<PyClassInitializer<T>>,
|
||||
T::Err: ToString,
|
||||
{
|
||||
type Err = PyErr;
|
||||
|
||||
#[inline(always)]
|
||||
fn from_str(s: &str) -> Result<Self, Self::Err> {
|
||||
Self::py_try_new(<T as FromStr>::from_str(s).pyerr()?)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> From<Py<T>> for NewPy<T> {
|
||||
#[inline(always)]
|
||||
fn from(inner: Py<T>) -> Self {
|
||||
Self::new(inner)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Deref for NewPy<T> {
|
||||
type Target = Py<T>;
|
||||
#[inline(always)]
|
||||
fn deref(&self) -> &Self::Target {
|
||||
self.inner()
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> DerefMut for NewPy<T> {
|
||||
#[inline(always)]
|
||||
fn deref_mut(&mut self) -> &mut Self::Target {
|
||||
self.inner_mut()
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, 'py, T> FromPyObject<'a, 'py> for NewPy<T>
|
||||
where
|
||||
Py<T>: FromPyObject<'a, 'py>,
|
||||
{
|
||||
type Error = <Py<T> as FromPyObject<'a, 'py>>::Error;
|
||||
|
||||
#[inline(always)]
|
||||
fn extract(ob: Borrowed<'a, 'py, PyAny>) -> Result<Self, Self::Error> {
|
||||
<Py<T> as FromPyObject<'a, 'py>>::extract(ob).map(Self::new)
|
||||
}
|
||||
}
|
||||
|
||||
impl<'py, T> IntoPyObject<'py> for NewPy<T>
|
||||
where
|
||||
Py<T>: IntoPyObject<'py>,
|
||||
{
|
||||
type Target = <Py<T> as IntoPyObject<'py>>::Target;
|
||||
type Output = <Py<T> as IntoPyObject<'py>>::Output;
|
||||
type Error = <Py<T> as IntoPyObject<'py>>::Error;
|
||||
|
||||
#[inline(always)]
|
||||
fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
|
||||
<Py<T> as IntoPyObject<'py>>::into_pyobject(self.into_inner(), py)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Serialize for NewPy<T>
|
||||
where
|
||||
Py<T>: Serialize,
|
||||
{
|
||||
#[inline(always)]
|
||||
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
|
||||
where
|
||||
S: Serializer,
|
||||
{
|
||||
<Py<T> as Serialize>::serialize(&*self, serializer)
|
||||
}
|
||||
}
|
||||
|
||||
impl<'de, T> Deserialize<'de> for NewPy<T>
|
||||
where
|
||||
Py<T>: Deserialize<'de>,
|
||||
{
|
||||
#[inline(always)]
|
||||
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
|
||||
where
|
||||
D: Deserializer<'de>,
|
||||
{
|
||||
<Py<T> as Deserialize<'de>>::deserialize(deserializer).map(Self::new)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> PyStubType for NewPy<T>
|
||||
where
|
||||
Py<T>: PyStubType,
|
||||
{
|
||||
#[inline(always)]
|
||||
fn type_output() -> TypeInfo {
|
||||
<Py<T> as PyStubType>::type_output()
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> NewPy<T> {
|
||||
#[inline(always)]
|
||||
fn new(inner: impl Into<Py<T>>) -> Self {
|
||||
Self(inner.into())
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
pub fn py_try_new_with(
|
||||
py: Python<'_>,
|
||||
value: impl Into<PyClassInitializer<T>>,
|
||||
) -> PyResult<Self>
|
||||
where
|
||||
T: PyClass,
|
||||
{
|
||||
Py::new(py, value).map(Self)
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
pub fn py_try_new(value: impl Into<PyClassInitializer<T>>) -> PyResult<Self>
|
||||
where
|
||||
T: PyClass,
|
||||
{
|
||||
Python::attach(|py| Self::py_try_new_with(py, value))
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
pub fn clap_try_new(value: impl Into<PyClassInitializer<T>>) -> Result<Self, clap::Error>
|
||||
where
|
||||
T: PyClass,
|
||||
{
|
||||
Self::py_try_new(value).map_err(|e| clap::Error::raw(clap::error::ErrorKind::Io, e))
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
pub fn inner(&self) -> &Py<T> {
|
||||
&self.0
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
pub fn inner_mut(&mut self) -> &mut Py<T> {
|
||||
&mut self.0
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
pub fn into_inner(self) -> Py<T> {
|
||||
self.0
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Parser> Parser for NewPy<T>
|
||||
where
|
||||
T: PyClass<Frozen = False> + Into<PyClassInitializer<T>>,
|
||||
{
|
||||
fn try_parse() -> Result<Self, clap::Error> {
|
||||
<T as Parser>::try_parse().and_then(Self::clap_try_new)
|
||||
}
|
||||
|
||||
fn try_parse_from<I, It>(itr: I) -> Result<Self, clap::Error>
|
||||
where
|
||||
I: IntoIterator<Item = It>,
|
||||
It: Into<OsString> + Clone,
|
||||
{
|
||||
<T as Parser>::try_parse_from(itr).and_then(Self::clap_try_new)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: CommandFactory> CommandFactory for NewPy<T> {
|
||||
fn command() -> clap::Command {
|
||||
<T as CommandFactory>::command()
|
||||
}
|
||||
fn command_for_update() -> clap::Command {
|
||||
<T as CommandFactory>::command_for_update()
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: FromArgMatches> FromArgMatches for NewPy<T>
|
||||
where
|
||||
T: PyClass<Frozen = False> + Into<PyClassInitializer<T>>,
|
||||
{
|
||||
fn from_arg_matches(matches: &ArgMatches) -> Result<Self, clap::Error> {
|
||||
<T as FromArgMatches>::from_arg_matches(matches).and_then(Self::clap_try_new)
|
||||
}
|
||||
fn from_arg_matches_mut(matches: &mut ArgMatches) -> Result<Self, clap::Error> {
|
||||
<T as FromArgMatches>::from_arg_matches_mut(matches).and_then(Self::clap_try_new)
|
||||
}
|
||||
fn update_from_arg_matches(&mut self, matches: &ArgMatches) -> Result<(), clap::Error> {
|
||||
Python::attach(|py| {
|
||||
<T as FromArgMatches>::update_from_arg_matches(&mut *self.borrow_mut(py), matches)
|
||||
})
|
||||
}
|
||||
fn update_from_arg_matches_mut(&mut self, matches: &mut ArgMatches) -> Result<(), clap::Error> {
|
||||
Python::attach(|py| {
|
||||
<T as FromArgMatches>::update_from_arg_matches_mut(&mut *self.borrow_mut(py), matches)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Args> Args for NewPy<T>
|
||||
where
|
||||
T: PyClass<Frozen = False> + Into<PyClassInitializer<T>>,
|
||||
{
|
||||
fn augment_args(cmd: clap::Command) -> clap::Command {
|
||||
<T as Args>::augment_args(cmd)
|
||||
}
|
||||
fn augment_args_for_update(cmd: clap::Command) -> clap::Command {
|
||||
<T as Args>::augment_args_for_update(cmd)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Subcommand> Subcommand for NewPy<T>
|
||||
where
|
||||
T: PyClass<Frozen = False> + Into<PyClassInitializer<T>>,
|
||||
{
|
||||
fn augment_subcommands(cmd: clap::Command) -> clap::Command {
|
||||
<T as Subcommand>::augment_subcommands(cmd)
|
||||
}
|
||||
fn augment_subcommands_for_update(cmd: clap::Command) -> clap::Command {
|
||||
<T as Subcommand>::augment_subcommands_for_update(cmd)
|
||||
}
|
||||
fn has_subcommand(name: &str) -> bool {
|
||||
<T as Subcommand>::has_subcommand(name)
|
||||
}
|
||||
}
|
||||
@@ -38,7 +38,7 @@ impl PyPidfileError {
|
||||
/// A PID file protected with a lock.
|
||||
///
|
||||
/// An instance of `Pidfile` can be used to manage a PID file: create it,
|
||||
/// lock it, detect already running daemons. It is backed by [`pidfile`][]
|
||||
/// lock it, detect already running daemons. It is backed by [`pidfile`]
|
||||
/// functions of `libbsd`/`libutil` which use `flopen` to lock the PID
|
||||
/// file.
|
||||
///
|
||||
File renamed without changes.
@@ -0,0 +1,57 @@
|
||||
import asyncio
|
||||
from multiprocessing.context import SpawnProcess
|
||||
import os
|
||||
import multiprocessing as mp
|
||||
|
||||
import pytest
|
||||
from _pytest.capture import CaptureFixture
|
||||
from exo_rs import (
|
||||
CliArgs,
|
||||
NetworkingHandle,
|
||||
Pidfile,
|
||||
FromSwarm,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sleep_on_multiple_items() -> None:
|
||||
print("PYTHON: starting handle")
|
||||
h = NetworkingHandle.new(os.urandom(16).hex().lstrip("0"), "default", 52414, 52413)
|
||||
print("PYTHON: handle started")
|
||||
|
||||
rt = asyncio.create_task(_await_recv(h))
|
||||
|
||||
# sleep for 4 ticks
|
||||
for i in range(10):
|
||||
await asyncio.sleep(1)
|
||||
|
||||
await h.gossipsub_publish("topic", b"somehting or other")
|
||||
|
||||
|
||||
async def _await_recv(h: NetworkingHandle):
|
||||
while True:
|
||||
event = await h.recv()
|
||||
match event:
|
||||
case FromSwarm.Connection() as c:
|
||||
print(f"PYTHON: connection update: {c}")
|
||||
case FromSwarm.Message() as m:
|
||||
print(f"PYTHON: message: {m}")
|
||||
case _:
|
||||
raise Exception("logical error")
|
||||
|
||||
|
||||
def test_pickling(capsys: CaptureFixture[str]):
|
||||
with capsys.disabled():
|
||||
p = mp.get_context("spawn").Process(
|
||||
target=run_mp, args=(CliArgs.parse_from(["exo"]),)
|
||||
)
|
||||
p.start()
|
||||
p.join()
|
||||
|
||||
|
||||
def run_mp(args: CliArgs):
|
||||
print("it got here")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(test_sleep_on_multiple_items())
|
||||
+20
-35
@@ -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
|
||||
@@ -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 => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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
@@ -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
@@ -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
@@ -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,
|
||||
})
|
||||
}
|
||||
@@ -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");
|
||||
}
|
||||
@@ -1,7 +0,0 @@
|
||||
// maybe this will hold test in the future...??
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
#[test]
|
||||
fn does_nothing() {}
|
||||
}
|
||||
+5
-4
@@ -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,8 +96,8 @@
|
||||
maturin build \
|
||||
--release \
|
||||
--manylinux off \
|
||||
--manifest-path rust/exo_pyo3_bindings/Cargo.toml \
|
||||
--features "pyo3/extension-module,pyo3/experimental-async" \
|
||||
--manifest-path rust/exo_rs/Cargo.toml \
|
||||
--features "pyo3/extension-module,pyo3/serde,pyo3/experimental-async" \
|
||||
--interpreter ${pkgs.python313}/bin/python \
|
||||
--out dist
|
||||
'';
|
||||
|
||||
@@ -13,3 +13,10 @@ path = "src/lib.rs"
|
||||
workspace = true
|
||||
|
||||
[dependencies]
|
||||
extend.workspace = true
|
||||
|
||||
expanduser = "1.2"
|
||||
path-clean = "1.0"
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = "3.23"
|
||||
+18
-1
@@ -1 +1,18 @@
|
||||
pub mod wakerdeque;
|
||||
use extend::ext;
|
||||
use std::collections::HashSet;
|
||||
use std::hash::Hash;
|
||||
|
||||
pub mod path;
|
||||
|
||||
#[ext(pub, name = VecExt)]
|
||||
impl<T> Vec<T> {
|
||||
/// Deduplicates vector while preserving the order.
|
||||
#[inline(always)]
|
||||
fn dedup_preserve_order(&mut self)
|
||||
where
|
||||
T: Eq + Hash + Clone,
|
||||
{
|
||||
let mut set = HashSet::new();
|
||||
self.retain(|x| set.insert(x.clone()));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,289 @@
|
||||
use extend::ext;
|
||||
use path_clean::PathClean;
|
||||
use std::fs::File;
|
||||
use std::path::{Component, Path, PathBuf};
|
||||
use std::{fs, io, path};
|
||||
|
||||
#[ext(pub, name = PathExt)]
|
||||
impl Path {
|
||||
/// Converts path to UTF-8 string, or returns `Err` if not UTF-8.
|
||||
#[inline(always)]
|
||||
fn to_str_utf8(&self) -> io::Result<&str> {
|
||||
self.to_str().ok_or_else(|| {
|
||||
io::Error::new(
|
||||
io::ErrorKind::InvalidFilename,
|
||||
"path contains non-utf8 characters",
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn create_file_if_not_found(&self) -> io::Result<()> {
|
||||
match File::create_new(self) {
|
||||
Ok(_) => Ok(()),
|
||||
Err(e) if e.kind() == io::ErrorKind::AlreadyExists => {
|
||||
if self.is_dir() {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::IsADirectory,
|
||||
format!("{self:?} is a directory, not a file"),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[inline(always)]
|
||||
fn try_dir_exists(&self) -> io::Result<()> {
|
||||
let m = fs::metadata(self)?;
|
||||
if m.is_dir() {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(io::Error::new(
|
||||
io::ErrorKind::NotADirectory,
|
||||
format!("{self:?} is not a directory"),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[inline(always)]
|
||||
fn try_file_exists(&self) -> io::Result<()> {
|
||||
let m = fs::metadata(self)?;
|
||||
if !m.is_dir() {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(io::Error::new(
|
||||
io::ErrorKind::IsADirectory,
|
||||
format!("{self:?} is a directory"),
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolves any path to its true absolute form as much as possible.
|
||||
///
|
||||
/// The path does not need to exist, but some component of it may exist.
|
||||
/// Expands any `~` or `~user`; eliminates any `..` or `.` and resolves
|
||||
/// symlinks by traversing the file system for the real part of the path,
|
||||
/// and performing lexical cleaning for the nonexistent part.
|
||||
#[cfg(unix)]
|
||||
pub fn resolve_path(path: PathBuf) -> io::Result<PathBuf> {
|
||||
// expand user if possible
|
||||
let mut path = expanduser::expanduser(path.to_str_utf8()?)?;
|
||||
|
||||
// convert to absolute, this will slightly clean path even if NOT relative
|
||||
path = path::absolute(&path)?;
|
||||
|
||||
let mut components = path.components().collect::<Vec<_>>();
|
||||
components
|
||||
.first()
|
||||
.filter(|&&c| c == Component::RootDir)
|
||||
.expect("the first component must exist, and be the root directory");
|
||||
|
||||
// resolve real prefix of path with `Path::canonicalize` (which will follow symlinks)
|
||||
// and nonexistent suffix with `PathClean::clean` (which will do lexical cleaning).
|
||||
//
|
||||
// 1) canonicalization is attempted iteratively to determine prefix/suffix split
|
||||
// 2) lexical cleaning ran on suffix which result in leading ".." components
|
||||
fn split_canonicalize(components: &[Component]) -> io::Result<(PathBuf, PathBuf)> {
|
||||
let mut prefix = PathBuf::new();
|
||||
let mut suffix = PathBuf::new();
|
||||
for i in (1..=components.len()).rev() {
|
||||
prefix = PathBuf::from_iter(&components[..i]);
|
||||
suffix = PathBuf::from_iter(&components[i..]);
|
||||
|
||||
match prefix.canonicalize() {
|
||||
Ok(p) => {
|
||||
// ensure non-leaf components are directories
|
||||
if i != components.len() && !fs::metadata(&p)?.is_dir() {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::AddrInUse,
|
||||
format!(
|
||||
"cannot resolve {:?}: {:?} is not a directory",
|
||||
PathBuf::from_iter(components),
|
||||
prefix
|
||||
),
|
||||
));
|
||||
}
|
||||
prefix = p;
|
||||
|
||||
// clean + substitute "." with empty buffer
|
||||
suffix = suffix.clean();
|
||||
if suffix == Path::new(".") {
|
||||
suffix = PathBuf::new()
|
||||
}
|
||||
break;
|
||||
}
|
||||
Err(e)
|
||||
if i > 1
|
||||
&& matches!(
|
||||
e.kind(),
|
||||
io::ErrorKind::NotFound | io::ErrorKind::NotADirectory
|
||||
) =>
|
||||
{
|
||||
continue;
|
||||
}
|
||||
Err(e) => return Err(e),
|
||||
}
|
||||
}
|
||||
Ok((prefix, suffix))
|
||||
}
|
||||
let (mut prefix, suffix) = split_canonicalize(&components)?;
|
||||
// 3) the ".." components are joined with real prefix and canonicalized again
|
||||
let mut suffix_components = suffix.components().collect::<Vec<_>>();
|
||||
suffix_components.reverse();
|
||||
while let Some(&c) = suffix_components.last() {
|
||||
if c != Component::ParentDir {
|
||||
break;
|
||||
}
|
||||
prefix.push(
|
||||
suffix_components
|
||||
.pop()
|
||||
.expect("already checked that its non-empty"),
|
||||
);
|
||||
}
|
||||
suffix_components.reverse();
|
||||
prefix = prefix.canonicalize()?;
|
||||
|
||||
// 4) prefix/suffix joined and 1) & 2) ran again to resolve any new exposed symlinks
|
||||
// NOTE: this time there shouldn't be any ".." in suffix
|
||||
prefix.extend(suffix_components);
|
||||
components = prefix.components().collect::<Vec<_>>();
|
||||
let (mut prefix, suffix) = split_canonicalize(&components)?;
|
||||
assert!(
|
||||
suffix
|
||||
.components()
|
||||
.all(|c| !matches!(c, Component::ParentDir | Component::CurDir)),
|
||||
"all `.` or `..` in suffix should have been eliminated"
|
||||
);
|
||||
|
||||
// 5) prefix/suffix joined in final absolute clean path with symlinks resolved and no ".."
|
||||
if !suffix.as_os_str().is_empty() {
|
||||
prefix.push(&suffix);
|
||||
}
|
||||
let path = prefix;
|
||||
assert!(
|
||||
path.components()
|
||||
.all(|c| !matches!(c, Component::ParentDir | Component::CurDir)),
|
||||
"all `.` or `..` in path should have been eliminated"
|
||||
);
|
||||
|
||||
Ok(path)
|
||||
}
|
||||
|
||||
#[cfg(all(test, unix))]
|
||||
mod tests {
|
||||
use super::resolve_path;
|
||||
use std::fs;
|
||||
use std::os::unix::fs::symlink;
|
||||
use std::path::{Component, Path};
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn assert_is_root_followed_by_normal_components(path: &Path) {
|
||||
let mut components = path.components();
|
||||
|
||||
assert_eq!(
|
||||
components.next(),
|
||||
Some(Component::RootDir),
|
||||
"resolved path should start with root: {}",
|
||||
path.display()
|
||||
);
|
||||
assert!(
|
||||
components.all(|component| matches!(component, Component::Normal(_))),
|
||||
"resolved path should contain only normal components after root: {}",
|
||||
path.display()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_path_lexically_cleans_nonexistent_suffix() {
|
||||
let test_dir = TempDir::new().unwrap();
|
||||
let base = test_dir.path().join("base");
|
||||
fs::create_dir_all(&base).unwrap();
|
||||
|
||||
let resolved = resolve_path(base.join("missing").join("..").join("leaf")).unwrap();
|
||||
|
||||
assert_eq!(resolved, base.canonicalize().unwrap().join("leaf"));
|
||||
assert_is_root_followed_by_normal_components(&resolved);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_path_applies_leading_suffix_parents_to_canonical_prefix() {
|
||||
let test_dir = TempDir::new().unwrap();
|
||||
let base = test_dir.path().join("base");
|
||||
fs::create_dir_all(&base).unwrap();
|
||||
|
||||
let resolved = resolve_path(
|
||||
base.join("missing")
|
||||
.join("..")
|
||||
.join("..")
|
||||
.join("outside")
|
||||
.join("leaf"),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
resolved,
|
||||
test_dir.path().canonicalize().unwrap().join("outside/leaf")
|
||||
);
|
||||
assert_is_root_followed_by_normal_components(&resolved);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_path_resolves_symlinks_exposed_by_cleaned_suffix() {
|
||||
let test_dir = TempDir::new().unwrap();
|
||||
let base = test_dir.path().join("base");
|
||||
let real_target = test_dir.path().join("real-target");
|
||||
let link = base.join("link");
|
||||
fs::create_dir_all(&base).unwrap();
|
||||
fs::create_dir_all(&real_target).unwrap();
|
||||
symlink(&real_target, &link).unwrap();
|
||||
|
||||
let resolved = resolve_path(
|
||||
base.join("missing")
|
||||
.join("..")
|
||||
.join("link")
|
||||
.join("future")
|
||||
.join("..")
|
||||
.join("leaf"),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(resolved, real_target.canonicalize().unwrap().join("leaf"));
|
||||
assert_is_root_followed_by_normal_components(&resolved);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_path_preserves_existing_symlink_parent_semantics() {
|
||||
let test_dir = TempDir::new().unwrap();
|
||||
let real_parent = test_dir.path().join("real-parent");
|
||||
let real_target = real_parent.join("target");
|
||||
let link = test_dir.path().join("link");
|
||||
fs::create_dir_all(&real_target).unwrap();
|
||||
symlink(&real_target, &link).unwrap();
|
||||
|
||||
let resolved = resolve_path(link.join("..")).unwrap();
|
||||
|
||||
assert_eq!(resolved, real_parent.canonicalize().unwrap());
|
||||
assert_is_root_followed_by_normal_components(&resolved);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_path_supports_existing_file() {
|
||||
let test_dir = TempDir::new().unwrap();
|
||||
let file = test_dir.path().join("config.toml");
|
||||
fs::write(&file, "verbosity = \"warn\"\n").unwrap();
|
||||
|
||||
let resolved = resolve_path(file.clone()).unwrap();
|
||||
|
||||
assert_eq!(resolved, file.canonicalize().unwrap());
|
||||
assert_eq!(
|
||||
resolved.canonicalize().unwrap(),
|
||||
file.canonicalize().unwrap()
|
||||
);
|
||||
assert_is_root_followed_by_normal_components(&resolved);
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
import os
|
||||
from importlib.metadata import version
|
||||
|
||||
# set __version__ and env-var
|
||||
__version__ = version("exo")
|
||||
os.environ["EXO_PKG_VERSION"] = __version__
|
||||
+21
-24
@@ -23,6 +23,7 @@ from hypercorn.typing import ASGIFramework
|
||||
from hypercorn.utils import LifespanTimeoutError, ShutdownError
|
||||
from loguru import logger
|
||||
|
||||
import exo.shared.config as config
|
||||
from exo.api.adapters.chat_completions import (
|
||||
chat_request_to_text_generation,
|
||||
collect_chat_response,
|
||||
@@ -128,12 +129,7 @@ from exo.master.placement import place_instance as get_instance_placements
|
||||
from exo.shared.apply import apply
|
||||
from exo.shared.constants import (
|
||||
DASHBOARD_DIR,
|
||||
ENABLE_DISAGGREGATION,
|
||||
EXO_CACHE_HOME,
|
||||
EXO_EVENT_LOG_DIR,
|
||||
EXO_IMAGE_CACHE_DIR,
|
||||
EXO_MAX_CHUNK_SIZE,
|
||||
EXO_TRACING_CACHE_DIR,
|
||||
)
|
||||
from exo.shared.election import ElectionMessage
|
||||
from exo.shared.logging import InterceptLogger
|
||||
@@ -207,9 +203,6 @@ from exo.utils.disk_event_log import DiskEventLog
|
||||
from exo.utils.power_sampler import PowerSampler
|
||||
from exo.utils.task_group import TaskGroup
|
||||
|
||||
_API_EVENT_LOG_DIR = EXO_EVENT_LOG_DIR / "api"
|
||||
ONBOARDING_COMPLETE_FILE = EXO_CACHE_HOME / "onboarding_complete"
|
||||
|
||||
|
||||
def _format_to_content_type(image_format: Literal["png", "jpeg", "webp"] | None) -> str:
|
||||
return f"image/{image_format or 'png'}"
|
||||
@@ -225,12 +218,12 @@ def _ensure_seed(params: AdvancedImageParams | None) -> AdvancedImageParams:
|
||||
|
||||
|
||||
def _require_disaggregation_enabled() -> None:
|
||||
if not ENABLE_DISAGGREGATION:
|
||||
if not config.app().disaggregation_enabled:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.NOT_FOUND,
|
||||
detail=(
|
||||
"Prefill/decode disaggregation is disabled. "
|
||||
"Set ENABLE_DISAGGREGATION=true to enable."
|
||||
"Run with --enable-disaggregation or EXO_DISAGGREGATION_ENABLED=true to enable."
|
||||
),
|
||||
)
|
||||
|
||||
@@ -248,7 +241,8 @@ class API:
|
||||
election_receiver: Receiver[ElectionMessage],
|
||||
) -> None:
|
||||
self.state = State()
|
||||
self._event_log = DiskEventLog(_API_EVENT_LOG_DIR)
|
||||
self._api_event_log_dir = config.bootstrap().event_log_dir / "api"
|
||||
self._event_log = DiskEventLog(self._api_event_log_dir)
|
||||
self._system_id = SystemId()
|
||||
self.command_sender = command_sender
|
||||
self.download_command_sender = download_command_sender
|
||||
@@ -258,6 +252,10 @@ class API:
|
||||
self.last_completed_election: int = 0
|
||||
self.port = port
|
||||
self._sent_image_hashes: set[str] = set()
|
||||
self._tracing_cache_dir = config.bootstrap().tracing_cache_dir
|
||||
self._onboarding_complete_file = (
|
||||
config.bootstrap().exo_home.cache / "onboarding_complete"
|
||||
)
|
||||
|
||||
self.paused: bool = False
|
||||
self.paused_ev: anyio.Event = anyio.Event()
|
||||
@@ -292,13 +290,13 @@ class API:
|
||||
self._image_generation_queues: dict[
|
||||
CommandId, Sender[ImageChunk | ErrorChunk]
|
||||
] = {}
|
||||
self._image_store = ImageStore(EXO_IMAGE_CACHE_DIR)
|
||||
self._image_store = ImageStore(config.bootstrap().image_cache_dir)
|
||||
self._tg: TaskGroup = TaskGroup()
|
||||
|
||||
def reset(self, result_clock: int, event_receiver: Receiver[IndexedEvent]):
|
||||
logger.info("Resetting API State")
|
||||
self._event_log.close()
|
||||
self._event_log = DiskEventLog(_API_EVENT_LOG_DIR)
|
||||
self._event_log = DiskEventLog(self._api_event_log_dir)
|
||||
self.state = State()
|
||||
self._system_id = SystemId()
|
||||
self._text_generation_queues = {}
|
||||
@@ -693,10 +691,10 @@ class API:
|
||||
)
|
||||
|
||||
async def get_feature_flags(self) -> dict[str, bool]:
|
||||
return {"disaggregation": ENABLE_DISAGGREGATION}
|
||||
return {"disaggregation": config.app().disaggregation_enabled}
|
||||
|
||||
async def list_instance_links(self) -> list[InstanceLink]:
|
||||
if not ENABLE_DISAGGREGATION:
|
||||
if not config.app().disaggregation_enabled:
|
||||
return []
|
||||
return list(self.state.instance_links.values())
|
||||
|
||||
@@ -2020,7 +2018,7 @@ class API:
|
||||
)
|
||||
for t in event.traces
|
||||
]
|
||||
output_path = EXO_TRACING_CACHE_DIR / f"trace_{event.task_id}.json"
|
||||
output_path = self._tracing_cache_dir / f"trace_{event.task_id}.json"
|
||||
export_trace(traces, output_path)
|
||||
logger.debug(f"Saved merged trace to {output_path}")
|
||||
|
||||
@@ -2082,10 +2080,9 @@ class API:
|
||||
await self._send_download(command)
|
||||
return CancelDownloadResponse(command_id=command.command_id)
|
||||
|
||||
@staticmethod
|
||||
def _get_trace_path(task_id: str) -> Path:
|
||||
trace_path = EXO_TRACING_CACHE_DIR / f"trace_{task_id}.json"
|
||||
if not trace_path.resolve().is_relative_to(EXO_TRACING_CACHE_DIR.resolve()):
|
||||
def _get_trace_path(self, task_id: str) -> Path:
|
||||
trace_path = self._tracing_cache_dir / f"trace_{task_id}.json"
|
||||
if not trace_path.resolve().is_relative_to(self._tracing_cache_dir.resolve()):
|
||||
raise HTTPException(status_code=400, detail=f"Invalid task ID: {task_id}")
|
||||
return trace_path
|
||||
|
||||
@@ -2093,7 +2090,7 @@ class API:
|
||||
traces: list[TraceListItem] = []
|
||||
|
||||
for trace_file in sorted(
|
||||
EXO_TRACING_CACHE_DIR.glob("trace_*.json"),
|
||||
self._tracing_cache_dir.glob("trace_*.json"),
|
||||
key=lambda p: p.stat().st_mtime,
|
||||
reverse=True,
|
||||
):
|
||||
@@ -2199,9 +2196,9 @@ class API:
|
||||
return DeleteTracesResponse(deleted=deleted, not_found=not_found)
|
||||
|
||||
async def get_onboarding(self) -> JSONResponse:
|
||||
return JSONResponse({"completed": ONBOARDING_COMPLETE_FILE.exists()})
|
||||
return JSONResponse({"completed": self._onboarding_complete_file.exists()})
|
||||
|
||||
async def complete_onboarding(self) -> JSONResponse:
|
||||
ONBOARDING_COMPLETE_FILE.parent.mkdir(parents=True, exist_ok=True)
|
||||
ONBOARDING_COMPLETE_FILE.write_text("true")
|
||||
self._onboarding_complete_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._onboarding_complete_file.write_text("true")
|
||||
return JSONResponse({"completed": True})
|
||||
@@ -7,6 +7,7 @@ import anyio
|
||||
from anyio import BrokenResourceError, ClosedResourceError, current_time, to_thread
|
||||
from loguru import logger
|
||||
|
||||
import exo.shared.config as config
|
||||
from exo.download.download_utils import (
|
||||
RepoDownloadProgress,
|
||||
delete_model,
|
||||
@@ -19,7 +20,6 @@ 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
|
||||
from exo.shared.types.commands import (
|
||||
@@ -69,7 +69,9 @@ class DownloadCoordinator:
|
||||
|
||||
@staticmethod
|
||||
def _default_model_dir(model_id: ModelId) -> str:
|
||||
return str(EXO_DEFAULT_MODELS_DIR / model_id.normalize())
|
||||
return str(
|
||||
config.bootstrap().models_dirs.default_models_dir / model_id.normalize()
|
||||
)
|
||||
|
||||
def _completed_from_path(
|
||||
self,
|
||||
@@ -433,7 +435,7 @@ class DownloadCoordinator:
|
||||
NodeDownloadProgress(download_progress=status)
|
||||
)
|
||||
# Scan read-only directories for pre-downloaded models
|
||||
if EXO_MODELS_READ_ONLY_DIRS:
|
||||
if config.bootstrap().models_dirs.models_read_only_dirs:
|
||||
for card in await model_cards.card_cache.list_all():
|
||||
mid = card.model_id
|
||||
if mid in self.active_downloads:
|
||||
|
||||
@@ -24,6 +24,7 @@ from pydantic import (
|
||||
TypeAdapter,
|
||||
)
|
||||
|
||||
import exo.shared.config as config
|
||||
from exo.download.huggingface_utils import (
|
||||
filter_repo_objects,
|
||||
get_allow_patterns,
|
||||
@@ -31,11 +32,6 @@ from exo.download.huggingface_utils import (
|
||||
get_hf_endpoint,
|
||||
get_hf_token,
|
||||
)
|
||||
from exo.shared.constants import (
|
||||
EXO_DEFAULT_MODELS_DIR,
|
||||
EXO_MODELS_DIRS,
|
||||
EXO_MODELS_READ_ONLY_DIRS,
|
||||
)
|
||||
from exo.shared.models.model_cards import ModelCard, ModelTask
|
||||
from exo.shared.types.common import ModelId
|
||||
from exo.shared.types.memory import Memory
|
||||
@@ -49,6 +45,23 @@ from exo.shared.types.worker.downloads import (
|
||||
from exo.shared.types.worker.shards import ShardMetadata
|
||||
|
||||
|
||||
def _default_models_dir() -> Path:
|
||||
return config.bootstrap().models_dirs.default_models_dir
|
||||
|
||||
|
||||
def _writable_models_dirs() -> list[Path]:
|
||||
return config.bootstrap().models_dirs.models_dirs
|
||||
|
||||
|
||||
def _read_only_models_dirs() -> list[Path]:
|
||||
return config.bootstrap().models_dirs.models_read_only_dirs
|
||||
|
||||
|
||||
def _model_search_dirs() -> tuple[Path, ...]:
|
||||
models_dirs = config.bootstrap().models_dirs
|
||||
return (*models_dirs.models_read_only_dirs, *models_dirs.models_dirs)
|
||||
|
||||
|
||||
class HuggingFaceAuthenticationError(Exception):
|
||||
"""Raised when HuggingFace returns 401/403 for a model download."""
|
||||
|
||||
@@ -159,7 +172,7 @@ def resolve_existing_model(
|
||||
all weight files are present.
|
||||
"""
|
||||
normalized = model_id.normalize()
|
||||
for search_dir in (*EXO_MODELS_READ_ONLY_DIRS, *EXO_MODELS_DIRS):
|
||||
for search_dir in _model_search_dirs():
|
||||
candidate = search_dir / normalized
|
||||
if candidate.is_dir() and is_model_directory_complete(candidate, card):
|
||||
return candidate
|
||||
@@ -168,14 +181,14 @@ def resolve_existing_model(
|
||||
|
||||
def is_read_only_model_dir(model_dir: Path) -> bool:
|
||||
"""Check if a model directory lives under a read-only models root."""
|
||||
return any(model_dir.is_relative_to(d) for d in EXO_MODELS_READ_ONLY_DIRS)
|
||||
return any(model_dir.is_relative_to(d) for d in _read_only_models_dirs())
|
||||
|
||||
|
||||
def build_model_path(model_id: ModelId) -> Path:
|
||||
found = resolve_existing_model(model_id)
|
||||
if found is not None:
|
||||
return found
|
||||
return EXO_DEFAULT_MODELS_DIR / model_id.normalize()
|
||||
return _default_models_dir() / model_id.normalize()
|
||||
|
||||
|
||||
def select_download_dir(required_bytes: int) -> Path:
|
||||
@@ -183,7 +196,8 @@ def select_download_dir(required_bytes: int) -> Path:
|
||||
|
||||
Raises ``InsufficientDiskSpaceError`` if none have enough space.
|
||||
"""
|
||||
for candidate_dir in EXO_MODELS_DIRS:
|
||||
writable_models_dirs = _writable_models_dirs()
|
||||
for candidate_dir in writable_models_dirs:
|
||||
if not candidate_dir.exists():
|
||||
continue
|
||||
try:
|
||||
@@ -194,7 +208,7 @@ def select_download_dir(required_bytes: int) -> Path:
|
||||
continue
|
||||
raise InsufficientDiskSpaceError(
|
||||
f"No writable model directory has {required_bytes / (1024**3):.1f} GiB free. "
|
||||
f"Checked: {[str(d) for d in EXO_MODELS_DIRS]}"
|
||||
f"Checked: {[str(d) for d in writable_models_dirs]}"
|
||||
)
|
||||
|
||||
|
||||
@@ -203,7 +217,7 @@ async def select_download_dir_for_shard(
|
||||
filtered_file_list: list[FileListEntry],
|
||||
total_size: int,
|
||||
) -> Path:
|
||||
for candidate_dir in EXO_MODELS_DIRS:
|
||||
for candidate_dir in _writable_models_dirs():
|
||||
if not candidate_dir.exists():
|
||||
continue
|
||||
sub = candidate_dir / model_id.normalize()
|
||||
@@ -234,7 +248,7 @@ async def resolve_model_dir(model_id: ModelId) -> Path:
|
||||
|
||||
async def ensure_cache_dir(model_id: ModelId) -> Path:
|
||||
"""Return the cache directory for a model's metadata, creating it if needed."""
|
||||
target = EXO_DEFAULT_MODELS_DIR / "caches" / model_id.normalize()
|
||||
target = _default_models_dir() / "caches" / model_id.normalize()
|
||||
await aios.makedirs(target, exist_ok=True)
|
||||
return target
|
||||
|
||||
@@ -243,14 +257,14 @@ async def delete_model(model_id: ModelId) -> bool:
|
||||
"""Delete a model from writable directories. Skips read-only dirs."""
|
||||
normalized = model_id.normalize()
|
||||
deleted = False
|
||||
for models_dir in EXO_MODELS_DIRS:
|
||||
for models_dir in _writable_models_dirs():
|
||||
model_dir = models_dir / normalized
|
||||
if await aios.path.exists(model_dir):
|
||||
await asyncio.to_thread(shutil.rmtree, model_dir, ignore_errors=False)
|
||||
deleted = True
|
||||
|
||||
# Clear cache from default dir
|
||||
cache_dir = EXO_DEFAULT_MODELS_DIR / "caches" / normalized
|
||||
cache_dir = _default_models_dir() / "caches" / normalized
|
||||
if await aios.path.exists(cache_dir):
|
||||
await asyncio.to_thread(shutil.rmtree, cache_dir, ignore_errors=False)
|
||||
|
||||
@@ -260,8 +274,8 @@ async def delete_model(model_id: ModelId) -> bool:
|
||||
async def seed_models(seed_dir: str | Path):
|
||||
"""Move models from resources folder to the default models directory."""
|
||||
source_dir = Path(seed_dir)
|
||||
await aios.makedirs(EXO_DEFAULT_MODELS_DIR, exist_ok=True)
|
||||
dest_dir = EXO_DEFAULT_MODELS_DIR
|
||||
await aios.makedirs(_default_models_dir(), exist_ok=True)
|
||||
dest_dir = _default_models_dir()
|
||||
for path in source_dir.iterdir():
|
||||
if path.is_dir() and path.name.startswith("models--"):
|
||||
dest_path = dest_dir / path.name
|
||||
@@ -349,7 +363,7 @@ def is_model_directory_complete(model_dir: Path, card: ModelCard | None = None)
|
||||
):
|
||||
vision_id = ModelId(card.vision.weights_repo)
|
||||
normalized = vision_id.normalize()
|
||||
for search_dir in (*EXO_MODELS_READ_ONLY_DIRS, *EXO_MODELS_DIRS):
|
||||
for search_dir in _model_search_dirs():
|
||||
candidate = search_dir / normalized
|
||||
if candidate.is_dir() and is_model_directory_complete(candidate):
|
||||
return True
|
||||
@@ -368,7 +382,7 @@ async def _build_file_list_from_local_directory(
|
||||
safetensors listed there.
|
||||
"""
|
||||
normalized = model_id.normalize()
|
||||
for search_dir in (*EXO_MODELS_READ_ONLY_DIRS, *EXO_MODELS_DIRS):
|
||||
for search_dir in _model_search_dirs():
|
||||
model_dir = search_dir / normalized
|
||||
if await aios.path.exists(model_dir):
|
||||
file_list = await asyncio.to_thread(
|
||||
@@ -910,7 +924,7 @@ async def download_shard(
|
||||
status="not_started",
|
||||
file_progress={},
|
||||
)
|
||||
return EXO_DEFAULT_MODELS_DIR / model_id.normalize(), not_started_progress
|
||||
return _default_models_dir() / model_id.normalize(), not_started_progress
|
||||
filtered_file_list = list(
|
||||
filter_repo_objects(
|
||||
file_list,
|
||||
@@ -936,7 +950,7 @@ async def download_shard(
|
||||
target_dir = (
|
||||
existing
|
||||
if existing is not None
|
||||
else EXO_DEFAULT_MODELS_DIR / model_id.normalize()
|
||||
else _default_models_dir() / model_id.normalize()
|
||||
)
|
||||
else:
|
||||
models_dir = await select_download_dir_for_shard(
|
||||
|
||||
@@ -3,11 +3,12 @@
|
||||
import time
|
||||
from datetime import timedelta
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
||||
|
||||
import aiofiles
|
||||
import aiofiles.os as aios
|
||||
import pytest
|
||||
from exo_rs import BootstrapSettings
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from exo.download.download_utils import (
|
||||
@@ -24,6 +25,13 @@ def model_id() -> ModelId:
|
||||
return ModelId("test-org/test-model")
|
||||
|
||||
|
||||
def _mock_bootstrap_settings(models_dir: Path) -> BootstrapSettings:
|
||||
cfg = BootstrapSettings.default()
|
||||
cfg.models_dirs.default_models_dir = models_dir
|
||||
cfg.models_dirs.models_dirs = [models_dir]
|
||||
return cfg
|
||||
|
||||
|
||||
class TestFileVerification:
|
||||
"""Tests for file size verification in _download_file."""
|
||||
|
||||
@@ -177,9 +185,14 @@ class TestFileListCache:
|
||||
FileListEntry(type="file", path="config.json", size=100),
|
||||
]
|
||||
|
||||
cfg = _mock_bootstrap_settings(models_dir)
|
||||
|
||||
with (
|
||||
patch("exo.download.download_utils.EXO_MODELS_DIRS", (models_dir,)),
|
||||
patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", models_dir),
|
||||
patch(
|
||||
"exo.download.download_utils.config.bootstrap",
|
||||
new_callable=Mock,
|
||||
return_value=cfg,
|
||||
),
|
||||
patch(
|
||||
"exo.download.download_utils.fetch_file_list_with_retry",
|
||||
new_callable=AsyncMock,
|
||||
@@ -224,9 +237,14 @@ class TestFileListCache:
|
||||
TypeAdapter(list[FileListEntry]).dump_json(cached_file_list).decode()
|
||||
)
|
||||
|
||||
cfg = _mock_bootstrap_settings(models_dir)
|
||||
|
||||
with (
|
||||
patch("exo.download.download_utils.EXO_MODELS_DIRS", (models_dir,)),
|
||||
patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", models_dir),
|
||||
patch(
|
||||
"exo.download.download_utils.config.bootstrap",
|
||||
new_callable=Mock,
|
||||
return_value=cfg,
|
||||
),
|
||||
patch(
|
||||
"exo.download.download_utils.fetch_file_list_with_retry",
|
||||
new_callable=AsyncMock,
|
||||
@@ -243,9 +261,14 @@ class TestFileListCache:
|
||||
"""Test that errors propagate when fetch fails and no cache exists."""
|
||||
models_dir = tmp_path / "models"
|
||||
|
||||
cfg = _mock_bootstrap_settings(models_dir)
|
||||
|
||||
with (
|
||||
patch("exo.download.download_utils.EXO_MODELS_DIRS", (models_dir,)),
|
||||
patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", models_dir),
|
||||
patch(
|
||||
"exo.download.download_utils.config.bootstrap",
|
||||
new_callable=Mock,
|
||||
return_value=cfg,
|
||||
),
|
||||
patch(
|
||||
"exo.download.download_utils.fetch_file_list_with_retry",
|
||||
new_callable=AsyncMock,
|
||||
@@ -277,9 +300,14 @@ class TestModelDeletion:
|
||||
async with aiofiles.open(cache_dir / "file_list.json", "w") as f:
|
||||
await f.write("[]")
|
||||
|
||||
cfg = _mock_bootstrap_settings(models_dir)
|
||||
|
||||
with (
|
||||
patch("exo.download.download_utils.EXO_MODELS_DIRS", (models_dir,)),
|
||||
patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", models_dir),
|
||||
patch(
|
||||
"exo.download.download_utils.config.bootstrap",
|
||||
new_callable=Mock,
|
||||
return_value=cfg,
|
||||
),
|
||||
):
|
||||
result = await delete_model(model_id)
|
||||
|
||||
@@ -299,9 +327,14 @@ class TestModelDeletion:
|
||||
async with aiofiles.open(cache_dir / "file_list.json", "w") as f:
|
||||
await f.write("[]")
|
||||
|
||||
cfg = _mock_bootstrap_settings(models_dir)
|
||||
|
||||
with (
|
||||
patch("exo.download.download_utils.EXO_MODELS_DIRS", (models_dir,)),
|
||||
patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", models_dir),
|
||||
patch(
|
||||
"exo.download.download_utils.config.bootstrap",
|
||||
new_callable=Mock,
|
||||
return_value=cfg,
|
||||
),
|
||||
):
|
||||
result = await delete_model(model_id)
|
||||
|
||||
@@ -317,9 +350,14 @@ class TestModelDeletion:
|
||||
models_dir = tmp_path / "models"
|
||||
await aios.makedirs(models_dir, exist_ok=True)
|
||||
|
||||
cfg = _mock_bootstrap_settings(models_dir)
|
||||
|
||||
with (
|
||||
patch("exo.download.download_utils.EXO_MODELS_DIRS", (models_dir,)),
|
||||
patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", models_dir),
|
||||
patch(
|
||||
"exo.download.download_utils.config.bootstrap",
|
||||
new_callable=Mock,
|
||||
return_value=cfg,
|
||||
),
|
||||
):
|
||||
result = await delete_model(model_id)
|
||||
|
||||
|
||||
@@ -2,13 +2,15 @@
|
||||
|
||||
import json
|
||||
import shutil
|
||||
from collections.abc import AsyncIterator
|
||||
from collections.abc import AsyncIterator, Iterator, Sequence
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import aiofiles
|
||||
import aiofiles.os as aios
|
||||
import pytest
|
||||
from exo_rs import BootstrapSettings
|
||||
|
||||
from exo.download.download_utils import (
|
||||
InsufficientDiskSpaceError,
|
||||
@@ -42,6 +44,37 @@ def _create_incomplete_model(model_dir: Path) -> None:
|
||||
# model.safetensors is missing
|
||||
|
||||
|
||||
def _mock_bootstrap_settings(
|
||||
default: Path,
|
||||
*,
|
||||
writable: list[Path] | None = None,
|
||||
read_only: list[Path] | None = None,
|
||||
) -> BootstrapSettings:
|
||||
cfg = BootstrapSettings.default()
|
||||
cfg.models_dirs.default_models_dir = default
|
||||
cfg.models_dirs.models_dirs = writable or []
|
||||
cfg.models_dirs.models_read_only_dirs = read_only or []
|
||||
return cfg
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patched_model_dirs(
|
||||
default: Path,
|
||||
*,
|
||||
writable: Sequence[Path] = (),
|
||||
read_only: Sequence[Path] = (),
|
||||
) -> Iterator[None]:
|
||||
cfg = _mock_bootstrap_settings(
|
||||
default, writable=list(writable), read_only=list(read_only)
|
||||
)
|
||||
with patch(
|
||||
"exo.download.download_utils.config.bootstrap",
|
||||
new_callable=Mock,
|
||||
return_value=cfg,
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# resolve_existing_model
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -51,19 +84,13 @@ class TestResolveExistingModel:
|
||||
def test_returns_none_when_no_dirs_have_model(self, tmp_path: Path) -> None:
|
||||
writable = tmp_path / "writable"
|
||||
writable.mkdir()
|
||||
with (
|
||||
patch("exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", ()),
|
||||
patch("exo.download.download_utils.EXO_MODELS_DIRS", (writable,)),
|
||||
):
|
||||
with _patched_model_dirs(tmp_path / "default", writable=[writable]):
|
||||
assert resolve_existing_model(MODEL_ID) is None
|
||||
|
||||
def test_finds_model_in_writable_dir(self, tmp_path: Path) -> None:
|
||||
writable = tmp_path / "writable"
|
||||
_create_complete_model(writable / NORMALIZED)
|
||||
with (
|
||||
patch("exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", ()),
|
||||
patch("exo.download.download_utils.EXO_MODELS_DIRS", (writable,)),
|
||||
):
|
||||
with _patched_model_dirs(tmp_path / "default", writable=[writable]):
|
||||
assert resolve_existing_model(MODEL_ID) == writable / NORMALIZED
|
||||
|
||||
def test_finds_model_in_read_only_dir(self, tmp_path: Path) -> None:
|
||||
@@ -71,11 +98,8 @@ class TestResolveExistingModel:
|
||||
_create_complete_model(read_only / NORMALIZED)
|
||||
writable = tmp_path / "writable"
|
||||
writable.mkdir()
|
||||
with (
|
||||
patch(
|
||||
"exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", (read_only,)
|
||||
),
|
||||
patch("exo.download.download_utils.EXO_MODELS_DIRS", (writable,)),
|
||||
with _patched_model_dirs(
|
||||
tmp_path / "default", writable=[writable], read_only=[read_only]
|
||||
):
|
||||
assert resolve_existing_model(MODEL_ID) == read_only / NORMALIZED
|
||||
|
||||
@@ -84,11 +108,8 @@ class TestResolveExistingModel:
|
||||
_create_complete_model(read_only / NORMALIZED)
|
||||
writable = tmp_path / "writable"
|
||||
_create_complete_model(writable / NORMALIZED)
|
||||
with (
|
||||
patch(
|
||||
"exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", (read_only,)
|
||||
),
|
||||
patch("exo.download.download_utils.EXO_MODELS_DIRS", (writable,)),
|
||||
with _patched_model_dirs(
|
||||
tmp_path / "default", writable=[writable], read_only=[read_only]
|
||||
):
|
||||
result = resolve_existing_model(MODEL_ID)
|
||||
assert result == read_only / NORMALIZED
|
||||
@@ -98,11 +119,8 @@ class TestResolveExistingModel:
|
||||
_create_incomplete_model(incomplete / NORMALIZED)
|
||||
complete = tmp_path / "complete"
|
||||
_create_complete_model(complete / NORMALIZED)
|
||||
with (
|
||||
patch(
|
||||
"exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", (incomplete,)
|
||||
),
|
||||
patch("exo.download.download_utils.EXO_MODELS_DIRS", (complete,)),
|
||||
with _patched_model_dirs(
|
||||
tmp_path / "default", writable=[complete], read_only=[incomplete]
|
||||
):
|
||||
result = resolve_existing_model(MODEL_ID)
|
||||
assert result == complete / NORMALIZED
|
||||
@@ -114,9 +132,8 @@ class TestResolveExistingModel:
|
||||
_create_complete_model(ro2 / NORMALIZED)
|
||||
writable = tmp_path / "writable"
|
||||
writable.mkdir()
|
||||
with (
|
||||
patch("exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", (ro1, ro2)),
|
||||
patch("exo.download.download_utils.EXO_MODELS_DIRS", (writable,)),
|
||||
with _patched_model_dirs(
|
||||
tmp_path / "default", writable=[writable], read_only=[ro1, ro2]
|
||||
):
|
||||
assert resolve_existing_model(MODEL_ID) == ro2 / NORMALIZED
|
||||
|
||||
@@ -129,18 +146,18 @@ class TestResolveExistingModel:
|
||||
class TestIsReadOnlyModelDir:
|
||||
def test_path_under_read_only_dir(self, tmp_path: Path) -> None:
|
||||
ro = tmp_path / "readonly"
|
||||
with patch("exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", (ro,)):
|
||||
with _patched_model_dirs(tmp_path / "default", read_only=[ro]):
|
||||
assert is_read_only_model_dir(ro / NORMALIZED) is True
|
||||
|
||||
def test_path_under_writable_dir(self, tmp_path: Path) -> None:
|
||||
writable = tmp_path / "writable"
|
||||
with patch("exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", ()):
|
||||
with _patched_model_dirs(tmp_path / "default", writable=[writable]):
|
||||
assert is_read_only_model_dir(writable / NORMALIZED) is False
|
||||
|
||||
def test_path_not_under_any_read_only_dir(self, tmp_path: Path) -> None:
|
||||
ro = tmp_path / "readonly"
|
||||
other = tmp_path / "other"
|
||||
with patch("exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", (ro,)):
|
||||
with _patched_model_dirs(tmp_path / "default", read_only=[ro]):
|
||||
assert is_read_only_model_dir(other / NORMALIZED) is False
|
||||
|
||||
|
||||
@@ -156,7 +173,7 @@ class TestSelectDownloadDir:
|
||||
dir1.mkdir()
|
||||
dir2.mkdir()
|
||||
# Both exist on same filesystem so both have space; first wins
|
||||
with patch("exo.download.download_utils.EXO_MODELS_DIRS", (dir1, dir2)):
|
||||
with _patched_model_dirs(tmp_path / "default", writable=[dir1, dir2]):
|
||||
assert select_download_dir(1) == dir1
|
||||
|
||||
def test_skips_dir_without_enough_space(self, tmp_path: Path) -> None:
|
||||
@@ -174,7 +191,7 @@ class TestSelectDownloadDir:
|
||||
return real_disk_usage(path)
|
||||
|
||||
with (
|
||||
patch("exo.download.download_utils.EXO_MODELS_DIRS", (dir1, dir2)),
|
||||
_patched_model_dirs(tmp_path / "default", writable=[dir1, dir2]),
|
||||
patch("shutil.disk_usage", side_effect=mock_disk_usage),
|
||||
):
|
||||
assert select_download_dir(1024) == dir2
|
||||
@@ -190,7 +207,7 @@ class TestSelectDownloadDir:
|
||||
return shutil._ntuple_diskusage(real.total, real.total, 0) # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
with (
|
||||
patch("exo.download.download_utils.EXO_MODELS_DIRS", (dir1,)),
|
||||
_patched_model_dirs(tmp_path / "default", writable=[dir1]),
|
||||
patch("shutil.disk_usage", side_effect=mock_disk_usage),
|
||||
pytest.raises(InsufficientDiskSpaceError),
|
||||
):
|
||||
@@ -199,7 +216,7 @@ class TestSelectDownloadDir:
|
||||
def test_skips_nonexistent_dir(self, tmp_path: Path) -> None:
|
||||
nonexistent = tmp_path / "does-not-exist"
|
||||
with (
|
||||
patch("exo.download.download_utils.EXO_MODELS_DIRS", (nonexistent,)),
|
||||
_patched_model_dirs(tmp_path / "default", writable=[nonexistent]),
|
||||
pytest.raises(InsufficientDiskSpaceError),
|
||||
):
|
||||
select_download_dir(1)
|
||||
@@ -218,7 +235,7 @@ class TestSelectDownloadDir:
|
||||
return real_disk_usage(path)
|
||||
|
||||
with (
|
||||
patch("exo.download.download_utils.EXO_MODELS_DIRS", (dir1, dir2)),
|
||||
_patched_model_dirs(tmp_path / "default", writable=[dir1, dir2]),
|
||||
patch("shutil.disk_usage", side_effect=mock_disk_usage),
|
||||
):
|
||||
assert select_download_dir(1) == dir2
|
||||
@@ -238,12 +255,15 @@ class TestDeleteModel:
|
||||
await aios.makedirs(writable1, exist_ok=True)
|
||||
await aios.makedirs(writable2, exist_ok=True)
|
||||
await aios.makedirs(default, exist_ok=True)
|
||||
with (
|
||||
patch(
|
||||
"exo.download.download_utils.EXO_MODELS_DIRS",
|
||||
(writable1, writable2, default),
|
||||
),
|
||||
patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", default),
|
||||
|
||||
cfg = _mock_bootstrap_settings(
|
||||
default, writable=[writable1, writable2, default]
|
||||
)
|
||||
|
||||
with patch(
|
||||
"exo.download.download_utils.config.bootstrap",
|
||||
new_callable=Mock,
|
||||
return_value=cfg,
|
||||
):
|
||||
yield writable1, writable2, default
|
||||
|
||||
|
||||
@@ -4,11 +4,12 @@ import os
|
||||
import time
|
||||
from collections.abc import AsyncIterator
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
import aiofiles
|
||||
import aiofiles.os as aios
|
||||
import pytest
|
||||
from exo_rs import BootstrapSettings
|
||||
|
||||
from exo.download.download_utils import (
|
||||
_download_file, # pyright: ignore[reportPrivateUsage]
|
||||
@@ -24,13 +25,26 @@ def model_id() -> ModelId:
|
||||
return ModelId("test-org/test-model")
|
||||
|
||||
|
||||
def _mock_bootstrap_settings(models_dir: Path) -> BootstrapSettings:
|
||||
cfg = BootstrapSettings.default()
|
||||
cfg.models_dirs.default_models_dir = models_dir
|
||||
cfg.models_dirs.models_dirs = [models_dir]
|
||||
return cfg
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def temp_models_dir(tmp_path: Path) -> AsyncIterator[Path]:
|
||||
models_dir = tmp_path / "models"
|
||||
await aios.makedirs(models_dir, exist_ok=True)
|
||||
|
||||
cfg = _mock_bootstrap_settings(models_dir)
|
||||
|
||||
with (
|
||||
patch("exo.download.download_utils.EXO_MODELS_DIRS", (models_dir,)),
|
||||
patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", models_dir),
|
||||
patch(
|
||||
"exo.download.download_utils.config.bootstrap",
|
||||
new_callable=Mock,
|
||||
return_value=cfg,
|
||||
),
|
||||
):
|
||||
yield models_dir
|
||||
|
||||
|
||||
+41
-154
@@ -1,4 +1,3 @@
|
||||
import argparse
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import resource
|
||||
@@ -10,24 +9,28 @@ from typing import Self
|
||||
import anyio
|
||||
from anyio.lowlevel import checkpoint as anyio_checkpoint
|
||||
from daemon import DaemonContext # pyright: ignore[reportMissingTypeStubs]
|
||||
from exo_pyo3_bindings import Pidfile, PidfileError
|
||||
from exo_rs import (
|
||||
AppSettings,
|
||||
BootstrapSettings,
|
||||
CliArgs,
|
||||
Pidfile,
|
||||
PidfileError,
|
||||
)
|
||||
from loguru import logger
|
||||
from pydantic import PositiveInt
|
||||
|
||||
import exo.routing.topics as topics
|
||||
import exo.shared.config as config
|
||||
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, EXO_PID_FILE
|
||||
from exo.routing.router import Router, get_node_zid
|
||||
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.pydantic_ext import FrozenModel
|
||||
from exo.utils.task_group import TaskGroup
|
||||
from exo.worker.main import Worker
|
||||
|
||||
@@ -49,14 +52,14 @@ class Node:
|
||||
_tg: TaskGroup = field(init=False, default_factory=TaskGroup)
|
||||
|
||||
@classmethod
|
||||
async def create(cls, args: "Args") -> Self:
|
||||
keypair = get_node_id_keypair()
|
||||
node_id = NodeId(keypair.to_node_id())
|
||||
async def create(cls, args: CliArgs) -> Self:
|
||||
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)
|
||||
@@ -72,23 +75,25 @@ class Node:
|
||||
)
|
||||
|
||||
logger.info(f"Starting node {node_id}")
|
||||
offline = config.app().offline
|
||||
|
||||
# Errors the very first time exo is run as dir doesn't exist
|
||||
EXO_DEFAULT_MODELS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
# Create writable model roots before download selection checks disk space.
|
||||
for models_dir in config.bootstrap().models_dirs.models_dirs:
|
||||
models_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Create DownloadCoordinator (unless --no-downloads)
|
||||
if not args.no_downloads:
|
||||
if args.downloads_enabled:
|
||||
download_coordinator = DownloadCoordinator(
|
||||
node_id,
|
||||
exo_shard_downloader(offline=args.offline),
|
||||
exo_shard_downloader(offline=offline),
|
||||
event_sender=event_router.sender(),
|
||||
download_command_receiver=router.receiver(topics.DOWNLOAD_COMMANDS),
|
||||
offline=args.offline,
|
||||
offline=offline,
|
||||
)
|
||||
else:
|
||||
download_coordinator = None
|
||||
|
||||
if args.spawn_api:
|
||||
if args.api_enabled:
|
||||
api = API(
|
||||
node_id,
|
||||
port=args.api_port,
|
||||
@@ -100,7 +105,7 @@ class Node:
|
||||
else:
|
||||
api = None
|
||||
|
||||
if not args.no_worker:
|
||||
if args.worker_enabled:
|
||||
worker = Worker(
|
||||
node_id,
|
||||
event_receiver=event_router.receiver(),
|
||||
@@ -147,7 +152,7 @@ class Node:
|
||||
master,
|
||||
api,
|
||||
node_id,
|
||||
args.offline,
|
||||
offline,
|
||||
args.api_port,
|
||||
)
|
||||
|
||||
@@ -275,12 +280,19 @@ class Node:
|
||||
|
||||
|
||||
def main():
|
||||
# Parse args first => --help or bad args don't require PID-locking
|
||||
args = Args.parse()
|
||||
# Parse args first & resolve/load bootstrap + app settings
|
||||
# => --help or bad args don't require PID-locking
|
||||
args = CliArgs.parse()
|
||||
bootstrap_settings = BootstrapSettings.resolve(args.bootstrap)
|
||||
config.load(
|
||||
bootstrap_settings,
|
||||
AppSettings.resolve(args.app, bootstrap_settings),
|
||||
)
|
||||
|
||||
# Exit early if cannot acquire PID file
|
||||
try:
|
||||
pidfile = Pidfile(EXO_PID_FILE, 0o0600)
|
||||
pidfile_path = config.bootstrap().pid_file
|
||||
pidfile = Pidfile(pidfile_path, 0o0600)
|
||||
except PidfileError as e:
|
||||
print(e, file=sys.stderr)
|
||||
raise SystemExit(1) from e
|
||||
@@ -328,7 +340,7 @@ def main():
|
||||
pidfile.close()
|
||||
|
||||
|
||||
def main_inner(args: "Args"):
|
||||
def main_inner(args: CliArgs):
|
||||
soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
|
||||
target = min(max(soft, 65535), hard)
|
||||
resource.setrlimit(resource.RLIMIT_NOFILE, (target, hard))
|
||||
@@ -336,31 +348,17 @@ def main_inner(args: "Args"):
|
||||
mp.set_start_method("spawn", force=True)
|
||||
|
||||
# TODO: Refactor the current verbosity system
|
||||
logger_setup(EXO_LOG, args.verbosity)
|
||||
logger_setup(config.bootstrap().log_files.exo_log, config.app().verbosity)
|
||||
|
||||
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()}")
|
||||
logger.info(f"Discovery namespace: {args.namespace}")
|
||||
|
||||
if args.offline:
|
||||
if config.app().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}")
|
||||
|
||||
if args.no_batch:
|
||||
os.environ["EXO_NO_BATCH"] = "1"
|
||||
if not config.app().continuous_batching_enabled:
|
||||
logger.info("Continuous batching disabled (--no-batch)")
|
||||
|
||||
# Set FAST_SYNCH override env var for runner subprocesses
|
||||
if args.fast_synch is True:
|
||||
os.environ["EXO_FAST_SYNCH"] = "true"
|
||||
logger.info("FAST_SYNCH forced ON")
|
||||
elif args.fast_synch is False:
|
||||
os.environ["EXO_FAST_SYNCH"] = "false"
|
||||
logger.info("FAST_SYNCH forced OFF")
|
||||
|
||||
node = anyio.run(Node.create, args)
|
||||
try:
|
||||
anyio.run(node.run)
|
||||
@@ -372,114 +370,3 @@ def main_inner(args: "Args"):
|
||||
finally:
|
||||
logger.info("EXO Shutdown complete")
|
||||
logger_cleanup()
|
||||
|
||||
|
||||
class Args(FrozenModel):
|
||||
verbosity: int = 0
|
||||
force_master: bool = False
|
||||
spawn_api: bool = False
|
||||
api_port: PositiveInt = 52415
|
||||
tb_only: bool = False
|
||||
no_worker: bool = False
|
||||
no_downloads: bool = False
|
||||
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
|
||||
legacy_daemon: bool = False
|
||||
bootstrap_peers: list[str] = []
|
||||
libp2p_port: int
|
||||
|
||||
@classmethod
|
||||
def parse(cls) -> Self:
|
||||
parser = argparse.ArgumentParser(prog="EXO")
|
||||
default_verbosity = 0
|
||||
parser.add_argument(
|
||||
"-q",
|
||||
"--quiet",
|
||||
action="store_const",
|
||||
const=-1,
|
||||
dest="verbosity",
|
||||
default=default_verbosity,
|
||||
)
|
||||
parser.add_argument(
|
||||
"-v",
|
||||
"--verbose",
|
||||
action="count",
|
||||
dest="verbosity",
|
||||
default=default_verbosity,
|
||||
)
|
||||
parser.add_argument(
|
||||
"-m",
|
||||
"--force-master",
|
||||
action="store_true",
|
||||
dest="force_master",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-api",
|
||||
action="store_false",
|
||||
dest="spawn_api",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--api-port",
|
||||
type=int,
|
||||
dest="api_port",
|
||||
default=52415,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-worker",
|
||||
action="store_true",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-downloads",
|
||||
action="store_true",
|
||||
help="Disable the download coordinator (node won't download models)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--offline",
|
||||
action="store_true",
|
||||
default=os.getenv("EXO_OFFLINE", "false").lower() == "true",
|
||||
help="Run in offline/air-gapped mode: skip internet checks, use only pre-staged local models",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-batch",
|
||||
action="store_true",
|
||||
help="Disable continuous batching, use sequential generation",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--legacy-daemon",
|
||||
action="store_true",
|
||||
help="Run as a legacy SysV-style background daemon using double-fork daemonization",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--bootstrap-peers",
|
||||
type=lambda s: [p for p in s.split(",") if p],
|
||||
default=os.getenv("EXO_BOOTSTRAP_PEERS", "").split(",")
|
||||
if os.getenv("EXO_BOOTSTRAP_PEERS")
|
||||
else [],
|
||||
dest="bootstrap_peers",
|
||||
help="Comma-separated libp2p multiaddrs to dial on startup (env: EXO_BOOTSTRAP_PEERS)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--libp2p-port",
|
||||
type=int,
|
||||
default=0,
|
||||
dest="libp2p_port",
|
||||
help="Fixed TCP port for libp2p to listen on (0 = OS-assigned).",
|
||||
)
|
||||
fast_synch_group = parser.add_mutually_exclusive_group()
|
||||
fast_synch_group.add_argument(
|
||||
"--fast-synch",
|
||||
action="store_true",
|
||||
dest="fast_synch",
|
||||
default=None,
|
||||
help="Force MLX FAST_SYNCH on (for JACCL backend)",
|
||||
)
|
||||
fast_synch_group.add_argument(
|
||||
"--no-fast-synch",
|
||||
action="store_false",
|
||||
dest="fast_synch",
|
||||
help="Force MLX FAST_SYNCH off",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
return cls(**vars(args)) # pyright: ignore[reportAny] - We are intentionally validating here, we can't do it statically
|
||||
@@ -3,6 +3,7 @@ from datetime import datetime, timedelta, timezone
|
||||
import anyio
|
||||
from loguru import logger
|
||||
|
||||
import exo.shared.config as config
|
||||
from exo.master.placement import (
|
||||
add_instance_to_placements,
|
||||
cancel_unnecessary_downloads,
|
||||
@@ -16,7 +17,6 @@ from exo.routing.event_router import (
|
||||
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 (
|
||||
AddCustomModelCard,
|
||||
CreateInstance,
|
||||
@@ -143,7 +143,8 @@ class Master:
|
||||
self.event_sender = event_sender
|
||||
self._system_id = SystemId()
|
||||
self._multi_buffer = MultiSourceBuffer[SystemId, Event]()
|
||||
self._event_log = DiskEventLog(EXO_EVENT_LOG_DIR / "master")
|
||||
self._master_event_log_dir = config.bootstrap().event_log_dir / "master"
|
||||
self._event_log = DiskEventLog(self._master_event_log_dir)
|
||||
self._pending_traces: dict[TaskId, dict[int, list[TraceEventData]]] = {}
|
||||
self._expected_ranks: dict[TaskId, set[int]] = {}
|
||||
|
||||
@@ -288,7 +289,7 @@ class Master:
|
||||
|
||||
self.command_task_mapping[command.command_id] = task_id
|
||||
|
||||
if EXO_TRACING_ENABLED:
|
||||
if config.app().tracing_enabled:
|
||||
selected_instance = self.state.instances.get(
|
||||
selected_instance_id
|
||||
)
|
||||
@@ -344,7 +345,7 @@ class Master:
|
||||
|
||||
self.command_task_mapping[command.command_id] = task_id
|
||||
|
||||
if EXO_TRACING_ENABLED:
|
||||
if config.app().tracing_enabled:
|
||||
selected_instance = self.state.instances.get(
|
||||
selected_instance_id
|
||||
)
|
||||
@@ -464,7 +465,7 @@ class Master:
|
||||
)
|
||||
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
|
||||
|
||||
@@ -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]()
|
||||
|
||||
@@ -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)
|
||||
+28
-37
@@ -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
|
||||
import exo.shared.config as config
|
||||
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,27 @@ 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 | None = None,
|
||||
) -> 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()
|
||||
path = path or config.bootstrap().node_zid
|
||||
|
||||
# TODO(evan): bring back node id persistence once we figure out how to deal with duplicates
|
||||
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 +263,4 @@ def get_node_id_keypair(
|
||||
keypair = Keypair.generate()
|
||||
f.write(keypair.to_bytes())
|
||||
return keypair
|
||||
"""
|
||||
@@ -58,7 +58,6 @@ from exo.utils.info_gatherer.info_gatherer import (
|
||||
MemoryUsage,
|
||||
MiscData,
|
||||
NodeBackends,
|
||||
NodeConfig,
|
||||
NodeDiskUsage,
|
||||
NodeNetworkInterfaces,
|
||||
RdmaCtlStatus,
|
||||
@@ -372,8 +371,6 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
|
||||
update["node_memory"] = {**state.node_memory, event.node_id: info}
|
||||
case NodeDiskUsage():
|
||||
update["node_disk"] = {**state.node_disk, event.node_id: info.disk_usage}
|
||||
case NodeConfig():
|
||||
pass
|
||||
case MiscData():
|
||||
current_identity = state.node_identities.get(event.node_id, NodeIdentity())
|
||||
new_identity = current_identity.model_copy(
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
from exo_rs import AppSettings, BootstrapSettings
|
||||
|
||||
# TODO: for now we are only going to have mutable globals as our configuration,
|
||||
# so reactive reload/injection still needs a real design.
|
||||
_bootstrap_settings: BootstrapSettings | None = None
|
||||
_app_settings: AppSettings | None = None
|
||||
|
||||
|
||||
def bootstrap() -> BootstrapSettings:
|
||||
global _bootstrap_settings
|
||||
if _bootstrap_settings is None:
|
||||
_bootstrap_settings = BootstrapSettings.from_env_only()
|
||||
return _bootstrap_settings
|
||||
|
||||
|
||||
def app() -> AppSettings:
|
||||
global _app_settings
|
||||
if _app_settings is None:
|
||||
_app_settings = AppSettings.from_env_only()
|
||||
return _app_settings
|
||||
|
||||
|
||||
def load_bootstrap(settings: BootstrapSettings):
|
||||
global _bootstrap_settings
|
||||
_bootstrap_settings = settings
|
||||
|
||||
|
||||
def load_app(settings: AppSettings):
|
||||
global _app_settings
|
||||
_app_settings = settings
|
||||
|
||||
|
||||
def load(bootstrap_settings: BootstrapSettings, app_settings: AppSettings):
|
||||
load_bootstrap(bootstrap_settings)
|
||||
load_app(app_settings)
|
||||
+13
-93
@@ -1,61 +1,23 @@
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from exo.utils.dashboard_path import find_dashboard, find_resources
|
||||
|
||||
_EXO_HOME_ENV = os.environ.get("EXO_HOME", None)
|
||||
# TODO: Remaining to-Rust migration candidates (some make more sense than others):
|
||||
#
|
||||
# EXO_MEMORY_THRESHOLD
|
||||
# EXO_MAX_INSTANCE_RETRIES
|
||||
# EXO_DASHBOARD_DIR + EXO_RESOURCES_DIR
|
||||
# EXO_MACMON_PATH
|
||||
#
|
||||
# OVERRIDE_MEMORY_MB ??
|
||||
# EXO_MAX_CHUNK_SIZE ??
|
||||
|
||||
# NOTE: I will leave these alone here for now - I don't know how I should port raw constants
|
||||
EXO_MAX_CHUNK_SIZE = 512 * 1024
|
||||
EXO_MAX_INSTANCE_RETRIES = 5
|
||||
|
||||
def _get_xdg_dir(env_var: str, fallback: str) -> Path:
|
||||
"""Get XDG directory, prioritising EXO_HOME environment variable if its set. On non-Linux platforms, default to ~/.exo."""
|
||||
|
||||
if _EXO_HOME_ENV is not None:
|
||||
return Path.home() / _EXO_HOME_ENV
|
||||
|
||||
if sys.platform != "linux":
|
||||
return Path.home() / ".exo"
|
||||
|
||||
xdg_value = os.environ.get(env_var, None)
|
||||
if xdg_value is not None:
|
||||
return Path(xdg_value) / "exo"
|
||||
return Path.home() / fallback / "exo"
|
||||
|
||||
|
||||
EXO_CONFIG_HOME = _get_xdg_dir("XDG_CONFIG_HOME", ".config")
|
||||
EXO_DATA_HOME = _get_xdg_dir("XDG_DATA_HOME", ".local/share")
|
||||
EXO_CACHE_HOME = _get_xdg_dir("XDG_CACHE_HOME", ".cache")
|
||||
|
||||
# Default models directory (always included as first entry in writable dirs)
|
||||
_EXO_DEFAULT_MODELS_DIR_ENV = os.environ.get("EXO_DEFAULT_MODELS_DIR", None)
|
||||
EXO_DEFAULT_MODELS_DIR = (
|
||||
Path(_EXO_DEFAULT_MODELS_DIR_ENV).expanduser()
|
||||
if _EXO_DEFAULT_MODELS_DIR_ENV is not None
|
||||
else EXO_DATA_HOME / "models"
|
||||
)
|
||||
|
||||
|
||||
def _parse_colon_dirs(env_var: str) -> tuple[Path, ...]:
|
||||
raw = os.environ.get(env_var, None)
|
||||
if raw is None:
|
||||
return ()
|
||||
return tuple(Path(p).expanduser() for p in raw.split(":") if p)
|
||||
|
||||
|
||||
# Read-only model directories (colon-separated). Never written to or deleted from.
|
||||
_EXO_MODELS_READ_ONLY_DIRS_ENV = _parse_colon_dirs("EXO_MODELS_READ_ONLY_DIRS")
|
||||
# Writable model directories (colon-separated). Default dir is always prepended.
|
||||
_EXO_MODELS_DIRS_ENV = _parse_colon_dirs("EXO_MODELS_DIRS")
|
||||
|
||||
# If a directory appears in both lists, treat it as read-only.
|
||||
_read_only_set = frozenset(_EXO_MODELS_READ_ONLY_DIRS_ENV)
|
||||
EXO_MODELS_DIRS: tuple[Path, ...] = tuple(
|
||||
d
|
||||
for d in (EXO_DEFAULT_MODELS_DIR, *_EXO_MODELS_DIRS_ENV)
|
||||
if d not in _read_only_set
|
||||
)
|
||||
EXO_MODELS_READ_ONLY_DIRS: tuple[Path, ...] = _EXO_MODELS_READ_ONLY_DIRS_ENV
|
||||
|
||||
# NOTE: I will leave these alone here for now - I don't wanna (or know how to) port them to Rust yet
|
||||
_RESOURCES_DIR_ENV = os.environ.get("EXO_RESOURCES_DIR", None)
|
||||
RESOURCES_DIR = (
|
||||
find_resources() if _RESOURCES_DIR_ENV is None else Path.home() / _RESOURCES_DIR_ENV
|
||||
@@ -64,45 +26,3 @@ _DASHBOARD_DIR_ENV = os.environ.get("EXO_DASHBOARD_DIR", None)
|
||||
DASHBOARD_DIR = (
|
||||
find_dashboard() if _DASHBOARD_DIR_ENV is None else Path.home() / _DASHBOARD_DIR_ENV
|
||||
)
|
||||
|
||||
# Log files (data/logs or cache)
|
||||
EXO_LOG_DIR = EXO_CACHE_HOME / "exo_log"
|
||||
EXO_LOG = EXO_LOG_DIR / "exo.log"
|
||||
EXO_RUNNER_LOG_DIR = EXO_LOG_DIR / "runner_log"
|
||||
EXO_RUNNER_STDOUT_LOG = EXO_RUNNER_LOG_DIR / "stdout.log"
|
||||
EXO_RUNNER_STDERR_LOG = EXO_RUNNER_LOG_DIR / "stderr.log"
|
||||
|
||||
EXO_TEST_LOG = EXO_CACHE_HOME / "exo_test.log"
|
||||
EXO_PID_FILE = EXO_CACHE_HOME / "exo.pid"
|
||||
|
||||
# Identity (config)
|
||||
EXO_NODE_ID_KEYPAIR = EXO_CONFIG_HOME / "node_id.keypair"
|
||||
EXO_CONFIG_FILE = EXO_CONFIG_HOME / "config.toml"
|
||||
|
||||
# libp2p topics for event forwarding
|
||||
LIBP2P_LOCAL_EVENTS_TOPIC = "worker_events"
|
||||
LIBP2P_GLOBAL_EVENTS_TOPIC = "global_events"
|
||||
LIBP2P_ELECTION_MESSAGES_TOPIC = "election_message"
|
||||
LIBP2P_COMMANDS_TOPIC = "commands"
|
||||
|
||||
EXO_MAX_CHUNK_SIZE = 512 * 1024
|
||||
|
||||
EXO_CUSTOM_MODEL_CARDS_DIR = EXO_DATA_HOME / "custom_model_cards"
|
||||
|
||||
EXO_EVENT_LOG_DIR = EXO_DATA_HOME / "event_log"
|
||||
EXO_IMAGE_CACHE_DIR = EXO_CACHE_HOME / "images"
|
||||
EXO_TRACING_CACHE_DIR = EXO_CACHE_HOME / "traces"
|
||||
|
||||
EXO_ENABLE_IMAGE_MODELS = (
|
||||
os.getenv("EXO_ENABLE_IMAGE_MODELS", "false").lower() == "true"
|
||||
)
|
||||
|
||||
EXO_OFFLINE = os.getenv("EXO_OFFLINE", "false").lower() == "true"
|
||||
|
||||
EXO_TRACING_ENABLED = os.getenv("EXO_TRACING_ENABLED", "false").lower() == "true"
|
||||
|
||||
ENABLE_DISAGGREGATION = os.getenv("ENABLE_DISAGGREGATION", "false").lower() == "true"
|
||||
|
||||
EXO_MAX_CONCURRENT_REQUESTS = int(os.getenv("EXO_MAX_CONCURRENT_REQUESTS", "8"))
|
||||
|
||||
EXO_MAX_INSTANCE_RETRIES = 5
|
||||
+34
-32
@@ -4,6 +4,7 @@ from collections.abc import Iterator
|
||||
from pathlib import Path
|
||||
|
||||
import zstandard
|
||||
from exo_rs import VerbosityFilter
|
||||
from hypercorn import Config
|
||||
from hypercorn.logging import Logger as HypercornLogger
|
||||
from loguru import logger
|
||||
@@ -43,10 +44,31 @@ class _InterceptHandler(logging.Handler):
|
||||
logger.opt(depth=3, exception=record.exc_info).log(level, record.getMessage())
|
||||
|
||||
|
||||
def logger_setup(log_file: Path | None, verbosity: int = 0):
|
||||
def _loguru_log_level(verbosity: VerbosityFilter):
|
||||
match verbosity:
|
||||
case VerbosityFilter.Off:
|
||||
raise ValueError(
|
||||
"VerbosityFilter.Off does not translate to a loguru log-level"
|
||||
)
|
||||
case VerbosityFilter.Error:
|
||||
return "ERROR"
|
||||
case VerbosityFilter.Warn:
|
||||
return "WARNING"
|
||||
case VerbosityFilter.Info:
|
||||
return "INFO"
|
||||
case VerbosityFilter.Debug:
|
||||
return "DEBUG"
|
||||
case VerbosityFilter.Trace:
|
||||
return "TRACE"
|
||||
|
||||
|
||||
def logger_setup(
|
||||
log_file: Path | None, verbosity: VerbosityFilter = VerbosityFilter.Info
|
||||
):
|
||||
"""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)
|
||||
|
||||
@@ -55,11 +77,17 @@ def logger_setup(log_file: Path | None, verbosity: int = 0):
|
||||
# replace all stdlib loggers with _InterceptHandlers that log to loguru
|
||||
logging.basicConfig(handlers=[_InterceptHandler()], level=0)
|
||||
|
||||
if verbosity == 0:
|
||||
# if Off then no logging - return early
|
||||
if verbosity == VerbosityFilter.Off:
|
||||
return
|
||||
|
||||
# info (or less verbose than info) gets a different formatter
|
||||
level = _loguru_log_level(verbosity)
|
||||
if verbosity <= VerbosityFilter.Info:
|
||||
logger.add(
|
||||
sys.__stderr__, # type: ignore
|
||||
format="[ {time:hh:mm:ss.SSSSA} | <level>{level: <8}</level>] <level>{message}</level>",
|
||||
level="INFO",
|
||||
level=level,
|
||||
colorize=True,
|
||||
enqueue=True,
|
||||
)
|
||||
@@ -67,7 +95,7 @@ def logger_setup(log_file: Path | None, verbosity: int = 0):
|
||||
logger.add(
|
||||
sys.__stderr__, # type: ignore
|
||||
format="[ {time:YYYY-MM-DD HH:mm:ss.SSS} | <level>{level: <8}</level> | {name}:{function}:{line} ] <level>{message}</level>",
|
||||
level="DEBUG",
|
||||
level=level,
|
||||
colorize=True,
|
||||
enqueue=True,
|
||||
)
|
||||
@@ -76,7 +104,7 @@ def logger_setup(log_file: Path | None, verbosity: int = 0):
|
||||
logger.add(
|
||||
log_file,
|
||||
format="[ {time:YYYY-MM-DD HH:mm:ss.SSS} | {level: <8} | {name}:{function}:{line} ] {message}",
|
||||
level="DEBUG" if verbosity > 0 else "INFO",
|
||||
level=level,
|
||||
colorize=False,
|
||||
enqueue=True,
|
||||
rotation=lambda _, __: next(rotate_once),
|
||||
@@ -88,29 +116,3 @@ def logger_setup(log_file: Path | None, verbosity: int = 0):
|
||||
def logger_cleanup():
|
||||
"""Flush all queues before shutting down so any in-flight logs are written to disk"""
|
||||
logger.complete()
|
||||
|
||||
|
||||
""" --- TODO: Capture MLX Log output:
|
||||
import contextlib
|
||||
import sys
|
||||
from loguru import logger
|
||||
|
||||
class StreamToLogger:
|
||||
|
||||
def __init__(self, level="INFO"):
|
||||
self._level = level
|
||||
|
||||
def write(self, buffer):
|
||||
for line in buffer.rstrip().splitlines():
|
||||
logger.opt(depth=1).log(self._level, line.rstrip())
|
||||
|
||||
def flush(self):
|
||||
pass
|
||||
|
||||
logger.remove()
|
||||
logger.add(sys.__stdout__)
|
||||
|
||||
stream = StreamToLogger()
|
||||
with contextlib.redirect_stdout(stream):
|
||||
print("Standard output is sent to added handlers.")
|
||||
"""
|
||||
@@ -20,12 +20,8 @@ from pydantic import (
|
||||
)
|
||||
from tomlkit.exceptions import TOMLKitError
|
||||
|
||||
from exo.shared.constants import (
|
||||
EXO_CUSTOM_MODEL_CARDS_DIR,
|
||||
EXO_ENABLE_IMAGE_MODELS,
|
||||
EXO_MODELS_DIRS,
|
||||
RESOURCES_DIR,
|
||||
)
|
||||
import exo.shared.config as config
|
||||
from exo.shared.constants import RESOURCES_DIR
|
||||
from exo.shared.types.backends import Backend
|
||||
from exo.shared.types.common import ModelId
|
||||
from exo.shared.types.memory import Memory
|
||||
@@ -34,13 +30,16 @@ from exo.utils.pydantic_ext import FrozenModel
|
||||
|
||||
# kinda ugly...
|
||||
# TODO: load search path from config.toml
|
||||
_custom_cards_dir = Path(str(EXO_CUSTOM_MODEL_CARDS_DIR))
|
||||
_BUILTIN_CARD_DIRS = [
|
||||
Path(RESOURCES_DIR) / "inference_model_cards",
|
||||
Path(RESOURCES_DIR) / "image_model_cards",
|
||||
]
|
||||
|
||||
|
||||
def _custom_cards_dir() -> Path:
|
||||
return Path(str(config.bootstrap().custom_model_cards_dir))
|
||||
|
||||
|
||||
class _CardCache:
|
||||
def __init__(self):
|
||||
self.cc: dict[ModelId, "ModelCard"] = {}
|
||||
@@ -57,7 +56,7 @@ class _CardCache:
|
||||
|
||||
async def pop(self, model_id: ModelId) -> "ModelCard | None":
|
||||
"""Delete a user-added custom model card. Returns True if deleted."""
|
||||
card_path = _custom_cards_dir / (ModelId(model_id).normalize() + ".toml")
|
||||
card_path = _custom_cards_dir() / (ModelId(model_id).normalize() + ".toml")
|
||||
try:
|
||||
if await card_path.exists():
|
||||
await card_path.unlink()
|
||||
@@ -68,7 +67,7 @@ class _CardCache:
|
||||
async def list_all(self) -> list["ModelCard"]:
|
||||
if len(self.cc) == 0:
|
||||
await self.refresh()
|
||||
if EXO_ENABLE_IMAGE_MODELS:
|
||||
if config.app().image_models_enabled:
|
||||
return list(self.cc.values())
|
||||
return [c for c in self.cc.values() if not _is_image_card(c)]
|
||||
|
||||
@@ -89,7 +88,7 @@ class _CardCache:
|
||||
async def refresh(self) -> None:
|
||||
for path in _BUILTIN_CARD_DIRS:
|
||||
await self._load_cards_from_dir(path, is_custom=False)
|
||||
await self._load_cards_from_dir(_custom_cards_dir, is_custom=True)
|
||||
await self._load_cards_from_dir(_custom_cards_dir(), is_custom=True)
|
||||
|
||||
|
||||
card_cache = _CardCache()
|
||||
@@ -97,7 +96,9 @@ card_cache = _CardCache()
|
||||
|
||||
def detect_vision_from_config(model_id: ModelId) -> "VisionCardConfig | None":
|
||||
normalized = model_id.normalize()
|
||||
for model_dir in [d / normalized for d in EXO_MODELS_DIRS]:
|
||||
for model_dir in [
|
||||
d / normalized for d in config.bootstrap().models_dirs.models_dirs
|
||||
]:
|
||||
config_path = model_dir / "config.json"
|
||||
if not config_path.exists():
|
||||
continue
|
||||
@@ -211,8 +212,8 @@ class ModelCard(FrozenModel):
|
||||
await f.write(data)
|
||||
|
||||
async def save_to_custom_dir(self) -> None:
|
||||
await aios.makedirs(str(_custom_cards_dir), exist_ok=True)
|
||||
await self.save(_custom_cards_dir / (self.model_id.normalize() + ".toml"))
|
||||
await aios.makedirs(str(_custom_cards_dir()), exist_ok=True)
|
||||
await self.save(_custom_cards_dir() / (self.model_id.normalize() + ".toml"))
|
||||
|
||||
@staticmethod
|
||||
async def load_from_path(path: Path) -> "ModelCard":
|
||||
|
||||
@@ -0,0 +1,174 @@
|
||||
"""Tests for bootstrap path and model-directory resolution."""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from exo_rs import BootstrapSettings
|
||||
|
||||
|
||||
@pytest.mark.skipif(sys.platform != "linux", reason="XDG dirs are Linux-specific")
|
||||
def test_xdg_paths_on_linux(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test that XDG paths are used on Linux when XDG env vars are set."""
|
||||
config_home = tmp_path / "test-config"
|
||||
data_home = tmp_path / "test-data"
|
||||
cache_home = tmp_path / "test-cache"
|
||||
|
||||
monkeypatch.delenv("EXO_HOME", raising=False)
|
||||
monkeypatch.setenv("XDG_CONFIG_HOME", str(config_home))
|
||||
monkeypatch.setenv("XDG_DATA_HOME", str(data_home))
|
||||
monkeypatch.setenv("XDG_CACHE_HOME", str(cache_home))
|
||||
|
||||
exo_home = BootstrapSettings.from_env_only().exo_home
|
||||
|
||||
assert config_home / "exo" == exo_home.config
|
||||
assert data_home / "exo" == exo_home.data
|
||||
assert cache_home / "exo" == exo_home.cache
|
||||
|
||||
|
||||
@pytest.mark.skipif(sys.platform != "darwin", reason="macOS dirs are Darwin-specific")
|
||||
def test_standard_directories_on_macos(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test that macOS standard directories are derived from HOME."""
|
||||
home = tmp_path / "home"
|
||||
|
||||
monkeypatch.delenv("EXO_HOME", raising=False)
|
||||
monkeypatch.setenv("HOME", str(home))
|
||||
monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path / "ignored-config"))
|
||||
monkeypatch.setenv("XDG_DATA_HOME", str(tmp_path / "ignored-data"))
|
||||
monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "ignored-cache"))
|
||||
|
||||
exo_home = BootstrapSettings.from_env_only().exo_home
|
||||
|
||||
assert home / "Library" / "Application Support" / "exo" == exo_home.config
|
||||
assert home / "Library" / "Application Support" / "exo" == exo_home.data
|
||||
assert home / "Library" / "Caches" / "exo" == exo_home.cache
|
||||
|
||||
|
||||
@pytest.mark.skipif(sys.platform != "linux", reason="XDG dirs are Linux-specific")
|
||||
def test_xdg_default_paths_on_linux(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test that XDG default paths are used on Linux when env vars are not set."""
|
||||
home = tmp_path / "home"
|
||||
|
||||
monkeypatch.delenv("EXO_HOME", raising=False)
|
||||
monkeypatch.delenv("XDG_CONFIG_HOME", raising=False)
|
||||
monkeypatch.delenv("XDG_DATA_HOME", raising=False)
|
||||
monkeypatch.delenv("XDG_CACHE_HOME", raising=False)
|
||||
monkeypatch.setenv("HOME", str(home))
|
||||
|
||||
exo_home = BootstrapSettings.from_env_only().exo_home
|
||||
|
||||
assert home / ".config" / "exo" == exo_home.config
|
||||
assert home / ".local" / "share" / "exo" == exo_home.data
|
||||
assert home / ".cache" / "exo" == exo_home.cache
|
||||
|
||||
|
||||
def test_legacy_exo_home_takes_precedence(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
"""Test that EXO_HOME environment variable takes precedence for backward compatibility."""
|
||||
exo_home_path = tmp_path / ".custom-exo"
|
||||
|
||||
monkeypatch.setenv("EXO_HOME", str(exo_home_path))
|
||||
monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path / "ignored-config"))
|
||||
monkeypatch.setenv("XDG_DATA_HOME", str(tmp_path / "ignored-data"))
|
||||
monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "ignored-cache"))
|
||||
|
||||
exo_home = BootstrapSettings.from_env_only().exo_home
|
||||
|
||||
assert exo_home_path == exo_home.config
|
||||
assert exo_home_path == exo_home.data
|
||||
assert exo_home_path == exo_home.cache
|
||||
|
||||
|
||||
@pytest.mark.skipif(sys.platform != "linux", reason="XDG dirs are Linux-specific")
|
||||
def test_models_in_data_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test that default models directory is in the data directory."""
|
||||
data_home = tmp_path / "data"
|
||||
|
||||
monkeypatch.delenv("EXO_HOME", raising=False)
|
||||
monkeypatch.delenv("EXO_DEFAULT_MODELS_DIR", raising=False)
|
||||
monkeypatch.delenv("EXO_MODELS_DIRS", raising=False)
|
||||
monkeypatch.setenv("XDG_DATA_HOME", str(data_home))
|
||||
|
||||
cfg = BootstrapSettings.from_env_only()
|
||||
|
||||
assert cfg.models_dirs.default_models_dir.parent == cfg.exo_home.data
|
||||
|
||||
|
||||
def test_default_dir_always_prepended_to_models_dirs(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
"""Test that the default models dir is always the first entry in EXO_MODELS_DIRS."""
|
||||
custom_models_dir = tmp_path / "custom-models"
|
||||
|
||||
monkeypatch.delenv("EXO_HOME", raising=False)
|
||||
monkeypatch.delenv("EXO_DEFAULT_MODELS_DIR", raising=False)
|
||||
monkeypatch.delenv("EXO_MODELS_READ_ONLY_DIRS", raising=False)
|
||||
monkeypatch.setenv("EXO_MODELS_DIRS", str(custom_models_dir))
|
||||
|
||||
models_dirs = BootstrapSettings.from_env_only().models_dirs
|
||||
|
||||
assert models_dirs.models_dirs[0] == models_dirs.default_models_dir
|
||||
assert custom_models_dir in models_dirs.models_dirs
|
||||
|
||||
|
||||
def test_default_models_dir_override(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test that EXO_DEFAULT_MODELS_DIR can be overridden via env var."""
|
||||
default_models_dir = tmp_path / "exo-models"
|
||||
|
||||
monkeypatch.delenv("EXO_HOME", raising=False)
|
||||
monkeypatch.delenv("EXO_MODELS_DIRS", raising=False)
|
||||
monkeypatch.delenv("EXO_MODELS_READ_ONLY_DIRS", raising=False)
|
||||
monkeypatch.setenv("EXO_DEFAULT_MODELS_DIR", str(default_models_dir))
|
||||
|
||||
models_dirs = BootstrapSettings.from_env_only().models_dirs
|
||||
|
||||
assert default_models_dir == models_dirs.default_models_dir
|
||||
assert models_dirs.models_dirs[0] == models_dirs.default_models_dir
|
||||
|
||||
|
||||
def test_default_dir_only_entry_when_env_unset(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test that EXO_MODELS_DIRS contains only the default when env var is not set."""
|
||||
monkeypatch.delenv("EXO_HOME", raising=False)
|
||||
monkeypatch.delenv("EXO_DEFAULT_MODELS_DIR", raising=False)
|
||||
monkeypatch.delenv("EXO_MODELS_DIRS", raising=False)
|
||||
monkeypatch.delenv("EXO_MODELS_READ_ONLY_DIRS", raising=False)
|
||||
|
||||
models_dirs = BootstrapSettings.from_env_only().models_dirs
|
||||
|
||||
assert models_dirs.models_dirs == [models_dirs.default_models_dir]
|
||||
|
||||
|
||||
def test_overlap_between_dirs_and_read_only_dirs(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
"""Test that a directory in both lists is excluded from writable dirs."""
|
||||
shared = tmp_path / "shared"
|
||||
writable_only = tmp_path / "writable-only"
|
||||
read_only = tmp_path / "read-only"
|
||||
shared.mkdir()
|
||||
read_only.mkdir()
|
||||
|
||||
monkeypatch.delenv("EXO_HOME", raising=False)
|
||||
monkeypatch.delenv("EXO_DEFAULT_MODELS_DIR", raising=False)
|
||||
monkeypatch.setenv("EXO_MODELS_DIRS", f"{shared}:{writable_only}")
|
||||
monkeypatch.setenv("EXO_MODELS_READ_ONLY_DIRS", f"{shared}:{read_only}")
|
||||
|
||||
models_dirs = BootstrapSettings.from_env_only().models_dirs
|
||||
|
||||
assert shared not in models_dirs.models_dirs
|
||||
assert writable_only in models_dirs.models_dirs
|
||||
assert shared in models_dirs.models_read_only_dirs
|
||||
assert read_only in models_dirs.models_read_only_dirs
|
||||
|
||||
|
||||
def test_empty_read_only_dirs_when_unset(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test that EXO_MODELS_READ_ONLY_DIRS is empty when env var is not set."""
|
||||
monkeypatch.delenv("EXO_HOME", raising=False)
|
||||
monkeypatch.delenv("EXO_DEFAULT_MODELS_DIR", raising=False)
|
||||
monkeypatch.delenv("EXO_MODELS_DIRS", raising=False)
|
||||
monkeypatch.delenv("EXO_MODELS_READ_ONLY_DIRS", raising=False)
|
||||
|
||||
models_dirs = BootstrapSettings.from_env_only().models_dirs
|
||||
|
||||
assert models_dirs.models_read_only_dirs == []
|
||||
@@ -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:
|
||||
|
||||
@@ -6,27 +6,27 @@ from multiprocessing.process import BaseProcess
|
||||
from multiprocessing.queues import Queue as QueueT
|
||||
from multiprocessing.synchronize import Event as EventT
|
||||
from multiprocessing.synchronize import Semaphore as SemaphoreT
|
||||
from pathlib import Path
|
||||
|
||||
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
|
||||
|
||||
NUM_CONCURRENT_PROCS = 10
|
||||
|
||||
|
||||
def _get_keypair_concurrent_subprocess_task(
|
||||
sem: SemaphoreT, ev: EventT, queue: QueueT[bytes]
|
||||
sem: SemaphoreT, ev: EventT, queue: QueueT[bytes], node_zid_path: Path
|
||||
) -> None:
|
||||
# synchronise with parent process
|
||||
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(node_zid_path).encode())
|
||||
|
||||
|
||||
def _get_keypair_concurrent(num_procs: int) -> bytes:
|
||||
def _get_keypair_concurrent(num_procs: int, node_zid_path: Path) -> bytes:
|
||||
assert num_procs > 0
|
||||
|
||||
sem = Semaphore(0)
|
||||
@@ -38,7 +38,8 @@ def _get_keypair_concurrent(num_procs: int) -> bytes:
|
||||
ps: list[BaseProcess] = []
|
||||
for _ in range(num_procs):
|
||||
p = multiprocessing.get_context("fork").Process(
|
||||
target=_get_keypair_concurrent_subprocess_task, args=(sem, ev, queue)
|
||||
target=_get_keypair_concurrent_subprocess_task,
|
||||
args=(sem, ev, queue, node_zid_path),
|
||||
)
|
||||
ps.append(p)
|
||||
p.start()
|
||||
@@ -75,19 +76,20 @@ def _delete_if_exists(p: str | bytes | os.PathLike[str] | os.PathLike[bytes]):
|
||||
|
||||
|
||||
@mark.skip(reason="this functionality is currently disabled but may return in future")
|
||||
def test_node_id_fetching(caplog: LogCaptureFixture):
|
||||
def test_node_id_fetching(caplog: LogCaptureFixture, tmp_path: Path):
|
||||
reps = 10
|
||||
node_zid_path = tmp_path / "node_zid"
|
||||
|
||||
# delete current file and write a new one
|
||||
_delete_if_exists(EXO_NODE_ID_KEYPAIR)
|
||||
kp = _get_keypair_concurrent(NUM_CONCURRENT_PROCS)
|
||||
_delete_if_exists(node_zid_path)
|
||||
kp = _get_keypair_concurrent(NUM_CONCURRENT_PROCS, node_zid_path)
|
||||
|
||||
with caplog.at_level(101): # supress logs
|
||||
# make sure that continuous fetches return the same value
|
||||
for _ in range(reps):
|
||||
assert kp == _get_keypair_concurrent(NUM_CONCURRENT_PROCS)
|
||||
assert kp == _get_keypair_concurrent(NUM_CONCURRENT_PROCS, node_zid_path)
|
||||
|
||||
# make sure that after deleting, we are not fetching the same value
|
||||
_delete_if_exists(EXO_NODE_ID_KEYPAIR)
|
||||
_delete_if_exists(node_zid_path)
|
||||
for _ in range(reps):
|
||||
assert kp != _get_keypair_concurrent(NUM_CONCURRENT_PROCS)
|
||||
assert kp != _get_keypair_concurrent(NUM_CONCURRENT_PROCS, node_zid_path)
|
||||
@@ -1,220 +0,0 @@
|
||||
"""Tests for XDG Base Directory Specification compliance."""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
|
||||
def test_xdg_paths_on_linux():
|
||||
"""Test that XDG paths are used on Linux when XDG env vars are set."""
|
||||
with (
|
||||
mock.patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"XDG_CONFIG_HOME": "/tmp/test-config",
|
||||
"XDG_DATA_HOME": "/tmp/test-data",
|
||||
"XDG_CACHE_HOME": "/tmp/test-cache",
|
||||
},
|
||||
clear=False,
|
||||
),
|
||||
mock.patch.object(sys, "platform", "linux"),
|
||||
):
|
||||
# Re-import to pick up mocked values
|
||||
import importlib
|
||||
|
||||
import exo.shared.constants as constants
|
||||
|
||||
importlib.reload(constants)
|
||||
|
||||
assert Path("/tmp/test-config/exo") == constants.EXO_CONFIG_HOME
|
||||
assert Path("/tmp/test-data/exo") == constants.EXO_DATA_HOME
|
||||
assert Path("/tmp/test-cache/exo") == constants.EXO_CACHE_HOME
|
||||
|
||||
|
||||
def test_xdg_default_paths_on_linux():
|
||||
"""Test that XDG default paths are used on Linux when env vars are not set."""
|
||||
# Remove XDG env vars and EXO_HOME
|
||||
env = {
|
||||
k: v
|
||||
for k, v in os.environ.items()
|
||||
if not k.startswith("XDG_") and k != "EXO_HOME"
|
||||
}
|
||||
with (
|
||||
mock.patch.dict(os.environ, env, clear=True),
|
||||
mock.patch.object(sys, "platform", "linux"),
|
||||
):
|
||||
import importlib
|
||||
|
||||
import exo.shared.constants as constants
|
||||
|
||||
importlib.reload(constants)
|
||||
|
||||
home = Path.home()
|
||||
assert home / ".config" / "exo" == constants.EXO_CONFIG_HOME
|
||||
assert home / ".local/share" / "exo" == constants.EXO_DATA_HOME
|
||||
assert home / ".cache" / "exo" == constants.EXO_CACHE_HOME
|
||||
|
||||
|
||||
def test_legacy_exo_home_takes_precedence():
|
||||
"""Test that EXO_HOME environment variable takes precedence for backward compatibility."""
|
||||
with mock.patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"EXO_HOME": ".custom-exo",
|
||||
"XDG_CONFIG_HOME": "/tmp/test-config",
|
||||
},
|
||||
clear=False,
|
||||
):
|
||||
import importlib
|
||||
|
||||
import exo.shared.constants as constants
|
||||
|
||||
importlib.reload(constants)
|
||||
|
||||
home = Path.home()
|
||||
assert home / ".custom-exo" == constants.EXO_CONFIG_HOME
|
||||
assert home / ".custom-exo" == constants.EXO_DATA_HOME
|
||||
|
||||
|
||||
def test_macos_uses_traditional_paths():
|
||||
"""Test that macOS uses traditional ~/.exo directory."""
|
||||
# Remove EXO_HOME to ensure we test the default behavior
|
||||
env = {k: v for k, v in os.environ.items() if k != "EXO_HOME"}
|
||||
with (
|
||||
mock.patch.dict(os.environ, env, clear=True),
|
||||
mock.patch.object(sys, "platform", "darwin"),
|
||||
):
|
||||
import importlib
|
||||
|
||||
import exo.shared.constants as constants
|
||||
|
||||
importlib.reload(constants)
|
||||
|
||||
home = Path.home()
|
||||
assert home / ".exo" == constants.EXO_CONFIG_HOME
|
||||
assert home / ".exo" == constants.EXO_DATA_HOME
|
||||
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
|
||||
env = {k: v for k, v in os.environ.items() if k != "EXO_MODELS_DIRS"}
|
||||
with mock.patch.dict(os.environ, env, clear=True):
|
||||
import importlib
|
||||
|
||||
import exo.shared.constants as constants
|
||||
|
||||
importlib.reload(constants)
|
||||
|
||||
assert constants.EXO_DEFAULT_MODELS_DIR.parent == constants.EXO_DATA_HOME
|
||||
|
||||
|
||||
def test_default_dir_always_prepended_to_models_dirs():
|
||||
"""Test that the default models dir is always the first entry in EXO_MODELS_DIRS."""
|
||||
env = {
|
||||
k: v
|
||||
for k, v in os.environ.items()
|
||||
if k not in ("EXO_MODELS_DIRS", "EXO_MODELS_READ_ONLY_DIRS", "EXO_HOME")
|
||||
}
|
||||
env["EXO_MODELS_DIRS"] = "/tmp/custom-models"
|
||||
with mock.patch.dict(os.environ, env, clear=True):
|
||||
import importlib
|
||||
|
||||
import exo.shared.constants as constants
|
||||
|
||||
importlib.reload(constants)
|
||||
|
||||
assert constants.EXO_MODELS_DIRS[0] == constants.EXO_DEFAULT_MODELS_DIR
|
||||
assert Path("/tmp/custom-models") in constants.EXO_MODELS_DIRS
|
||||
|
||||
|
||||
def test_default_models_dir_override():
|
||||
"""Test that EXO_DEFAULT_MODELS_DIR can be overridden via env var."""
|
||||
env = {
|
||||
k: v
|
||||
for k, v in os.environ.items()
|
||||
if k
|
||||
not in (
|
||||
"EXO_MODELS_DIRS",
|
||||
"EXO_MODELS_READ_ONLY_DIRS",
|
||||
"EXO_HOME",
|
||||
"EXO_DEFAULT_MODELS_DIR",
|
||||
)
|
||||
}
|
||||
env["EXO_DEFAULT_MODELS_DIR"] = "/Volumes/FastSSD/exo-models"
|
||||
with mock.patch.dict(os.environ, env, clear=True):
|
||||
import importlib
|
||||
|
||||
import exo.shared.constants as constants
|
||||
|
||||
importlib.reload(constants)
|
||||
|
||||
assert Path("/Volumes/FastSSD/exo-models") == constants.EXO_DEFAULT_MODELS_DIR
|
||||
assert constants.EXO_MODELS_DIRS[0] == constants.EXO_DEFAULT_MODELS_DIR
|
||||
|
||||
|
||||
def test_default_dir_only_entry_when_env_unset():
|
||||
"""Test that EXO_MODELS_DIRS contains only the default when env var is not set."""
|
||||
env = {
|
||||
k: v
|
||||
for k, v in os.environ.items()
|
||||
if k not in ("EXO_MODELS_DIRS", "EXO_MODELS_READ_ONLY_DIRS", "EXO_HOME")
|
||||
}
|
||||
with mock.patch.dict(os.environ, env, clear=True):
|
||||
import importlib
|
||||
|
||||
import exo.shared.constants as constants
|
||||
|
||||
importlib.reload(constants)
|
||||
|
||||
assert constants.EXO_MODELS_DIRS == (constants.EXO_DEFAULT_MODELS_DIR,)
|
||||
|
||||
|
||||
def test_overlap_between_dirs_and_read_only_dirs():
|
||||
"""Test that a directory in both lists is excluded from writable dirs."""
|
||||
env = {
|
||||
k: v
|
||||
for k, v in os.environ.items()
|
||||
if k not in ("EXO_MODELS_DIRS", "EXO_MODELS_READ_ONLY_DIRS", "EXO_HOME")
|
||||
}
|
||||
env["EXO_MODELS_DIRS"] = "/tmp/shared:/tmp/writable-only"
|
||||
env["EXO_MODELS_READ_ONLY_DIRS"] = "/tmp/shared:/tmp/ro-only"
|
||||
with mock.patch.dict(os.environ, env, clear=True):
|
||||
import importlib
|
||||
|
||||
import exo.shared.constants as constants
|
||||
|
||||
importlib.reload(constants)
|
||||
|
||||
# /tmp/shared should be excluded from writable dirs
|
||||
assert Path("/tmp/shared") not in constants.EXO_MODELS_DIRS
|
||||
assert Path("/tmp/writable-only") in constants.EXO_MODELS_DIRS
|
||||
# /tmp/shared should still be in read-only dirs
|
||||
assert Path("/tmp/shared") in constants.EXO_MODELS_READ_ONLY_DIRS
|
||||
assert Path("/tmp/ro-only") in constants.EXO_MODELS_READ_ONLY_DIRS
|
||||
|
||||
|
||||
def test_empty_read_only_dirs_when_unset():
|
||||
"""Test that EXO_MODELS_READ_ONLY_DIRS is empty when env var is not set."""
|
||||
env = {
|
||||
k: v
|
||||
for k, v in os.environ.items()
|
||||
if k not in ("EXO_MODELS_DIRS", "EXO_MODELS_READ_ONLY_DIRS", "EXO_HOME")
|
||||
}
|
||||
with mock.patch.dict(os.environ, env, clear=True):
|
||||
import importlib
|
||||
|
||||
import exo.shared.constants as constants
|
||||
|
||||
importlib.reload(constants)
|
||||
|
||||
assert constants.EXO_MODELS_READ_ONLY_DIRS == ()
|
||||
@@ -10,7 +10,7 @@ from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import cast, final
|
||||
|
||||
from exo.shared.constants import EXO_TRACING_ENABLED
|
||||
import exo.shared.config as config
|
||||
from exo.worker.runner.bootstrap import logger
|
||||
|
||||
# Context variable to track the current trace category for hierarchical nesting
|
||||
@@ -98,7 +98,7 @@ def trace(
|
||||
# Recorded with category "sync/compute"
|
||||
hidden_states = some_computation(...)
|
||||
"""
|
||||
if not EXO_TRACING_ENABLED:
|
||||
if not config.app().tracing_enabled:
|
||||
yield
|
||||
return
|
||||
|
||||
|
||||
+10
-6
@@ -2,21 +2,25 @@ import logging
|
||||
import os
|
||||
import sys
|
||||
import webbrowser
|
||||
from pathlib import Path
|
||||
|
||||
from exo.shared.constants import EXO_CONFIG_HOME
|
||||
import exo.shared.config as config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_FIRST_RUN_MARKER = EXO_CONFIG_HOME / ".dashboard_opened"
|
||||
|
||||
def _first_run_marker() -> Path:
|
||||
return config.bootstrap().exo_home.config / ".dashboard_opened"
|
||||
|
||||
|
||||
def _is_first_run() -> bool:
|
||||
return not _FIRST_RUN_MARKER.exists()
|
||||
return not _first_run_marker().exists()
|
||||
|
||||
|
||||
def _mark_first_run_done() -> None:
|
||||
_FIRST_RUN_MARKER.parent.mkdir(parents=True, exist_ok=True)
|
||||
_FIRST_RUN_MARKER.touch()
|
||||
first_run_marker = _first_run_marker()
|
||||
first_run_marker.parent.mkdir(parents=True, exist_ok=True)
|
||||
first_run_marker.touch()
|
||||
|
||||
|
||||
def print_startup_banner(port: int) -> None:
|
||||
@@ -38,7 +42,7 @@ def print_startup_banner(port: int) -> None:
|
||||
|
||||
╔═══════════════════════════════════════════════════════════════════════╗
|
||||
║ ║
|
||||
║ 🌐 Dashboard & API Ready ║
|
||||
║ Dashboard & API Ready ║
|
||||
║ ║
|
||||
║ {dashboard_url}{" " * (69 - len(dashboard_url))}║
|
||||
║ ║
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tomllib
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from subprocess import CalledProcessError
|
||||
@@ -13,7 +12,7 @@ from anyio.streams.buffered import BufferedByteReceiveStream
|
||||
from loguru import logger
|
||||
from pydantic import ValidationError
|
||||
|
||||
from exo.shared.constants import EXO_CONFIG_FILE, EXO_DEFAULT_MODELS_DIR
|
||||
import exo.shared.config as config
|
||||
from exo.shared.types.backends import Backend
|
||||
from exo.shared.types.memory import Memory
|
||||
from exo.shared.types.profiling import (
|
||||
@@ -292,24 +291,6 @@ class ThunderboltBridgeInfo(TaggedModel):
|
||||
return None
|
||||
|
||||
|
||||
class NodeConfig(TaggedModel):
|
||||
"""Node configuration from EXO_CONFIG_FILE, reloaded from the file only at startup. Other changes should come in through the API and propagate from there"""
|
||||
|
||||
@classmethod
|
||||
async def gather(cls) -> Self | None:
|
||||
cfg_file = anyio.Path(EXO_CONFIG_FILE)
|
||||
await cfg_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
await cfg_file.touch(exist_ok=True)
|
||||
async with await cfg_file.open("rb") as f:
|
||||
try:
|
||||
contents = (await f.read()).decode("utf-8")
|
||||
data = tomllib.loads(contents)
|
||||
return cls.model_validate(data)
|
||||
except (tomllib.TOMLDecodeError, UnicodeDecodeError, ValidationError):
|
||||
logger.warning("Invalid config file, skipping...")
|
||||
return None
|
||||
|
||||
|
||||
class MiscData(TaggedModel):
|
||||
"""Node information that may slowly change that doesn't fall into the other categories"""
|
||||
|
||||
@@ -329,7 +310,7 @@ class NodeDiskUsage(TaggedModel):
|
||||
async def gather(cls) -> Self:
|
||||
return cls(
|
||||
disk_usage=await to_thread.run_sync(
|
||||
DiskUsage.from_path, EXO_DEFAULT_MODELS_DIR
|
||||
DiskUsage.from_path, config.bootstrap().models_dirs.default_models_dir
|
||||
)
|
||||
)
|
||||
|
||||
@@ -391,7 +372,6 @@ GatheredInfo = (
|
||||
| MacThunderboltConnections
|
||||
| RdmaCtlStatus
|
||||
| ThunderboltBridgeInfo
|
||||
| NodeConfig
|
||||
| MiscData
|
||||
| StaticNodeInformation
|
||||
| NodeDiskUsage
|
||||
@@ -455,10 +435,6 @@ class InfoGatherer:
|
||||
tg.start_soon(self._monitor_static_info, 60)
|
||||
tg.start_soon(self._monitor_disk_usage, 30)
|
||||
|
||||
nc = await NodeConfig.gather()
|
||||
if nc is not None:
|
||||
await self.info_sender.send(nc)
|
||||
|
||||
await self.info_sender.send(await NodeBackends.gather())
|
||||
|
||||
def shutdown(self):
|
||||
|
||||
@@ -8,14 +8,14 @@ import textwrap
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
from exo_pyo3_bindings.exo_pyo3_bindings import Pidfile
|
||||
from exo_rs import Pidfile
|
||||
|
||||
_CHILD_ACQUIRE_PIDFILE_SCRIPT: Final = textwrap.dedent(
|
||||
"""
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from exo_pyo3_bindings.exo_pyo3_bindings import Pidfile, PidfileError
|
||||
from exo_rs import Pidfile, PidfileError
|
||||
|
||||
path = Path(sys.argv[1])
|
||||
try:
|
||||
|
||||
@@ -7,8 +7,8 @@ from typing import BinaryIO
|
||||
import mlx.core as mx
|
||||
from loguru import logger
|
||||
|
||||
import exo.shared.config as config
|
||||
from exo.api.types import ImageEditsTaskParams, ImageGenerationTaskParams
|
||||
from exo.shared.constants import EXO_TRACING_ENABLED
|
||||
from exo.shared.tracing import clear_trace_buffer, get_trace_buffer
|
||||
from exo.shared.types.chunks import Chunk, ErrorChunk
|
||||
from exo.shared.types.events import (
|
||||
@@ -70,7 +70,7 @@ def _send_traces_if_enabled(
|
||||
task_id: TaskId,
|
||||
rank: int,
|
||||
) -> None:
|
||||
if not EXO_TRACING_ENABLED:
|
||||
if not config.app().tracing_enabled:
|
||||
return
|
||||
|
||||
traces = get_trace_buffer()
|
||||
|
||||
@@ -8,7 +8,7 @@ from mflux.models.common.config.config import Config
|
||||
from mflux.utils.exceptions import StopImageGenerationException
|
||||
from tqdm import tqdm
|
||||
|
||||
from exo.shared.constants import EXO_TRACING_ENABLED
|
||||
import exo.shared.config as exo_config
|
||||
from exo.shared.tracing import (
|
||||
clear_trace_buffer,
|
||||
trace,
|
||||
@@ -544,8 +544,7 @@ class DiffusionRunner:
|
||||
num_sync_steps: int,
|
||||
capture_steps: set[int] | None = None,
|
||||
):
|
||||
if capture_steps is None:
|
||||
capture_steps = set()
|
||||
capture_steps = capture_steps or set()
|
||||
|
||||
self._reset_all_caches()
|
||||
clear_trace_buffer()
|
||||
@@ -863,7 +862,7 @@ class DiffusionRunner:
|
||||
rotary_embeddings=image_rotary_embeddings,
|
||||
)
|
||||
|
||||
if EXO_TRACING_ENABLED:
|
||||
if exo_config.app().tracing_enabled:
|
||||
mx.eval(encoder_hidden_states, hidden_states)
|
||||
|
||||
if self.owns_concat_stage:
|
||||
@@ -926,7 +925,7 @@ class DiffusionRunner:
|
||||
rotary_embeddings=image_rotary_embeddings,
|
||||
)
|
||||
|
||||
if EXO_TRACING_ENABLED:
|
||||
if exo_config.app().tracing_enabled:
|
||||
mx.eval(hidden_states)
|
||||
|
||||
if not self.is_last_stage:
|
||||
@@ -1194,7 +1193,7 @@ class DiffusionRunner:
|
||||
rotary_embeddings=image_rotary_embeddings,
|
||||
)
|
||||
|
||||
if EXO_TRACING_ENABLED:
|
||||
if exo_config.app().tracing_enabled:
|
||||
mx.eval(encoder_hidden_states, patch)
|
||||
|
||||
if self.owns_concat_stage:
|
||||
@@ -1264,7 +1263,7 @@ class DiffusionRunner:
|
||||
rotary_embeddings=image_rotary_embeddings,
|
||||
)
|
||||
|
||||
if EXO_TRACING_ENABLED:
|
||||
if exo_config.app().tracing_enabled:
|
||||
mx.eval(patch)
|
||||
|
||||
if not self.is_last_stage:
|
||||
|
||||
@@ -1,11 +1,11 @@
|
||||
import contextlib
|
||||
import os
|
||||
from collections.abc import Generator
|
||||
from dataclasses import dataclass
|
||||
|
||||
import mlx.core as mx
|
||||
from mlx_lm.tokenizer_utils import TokenizerWrapper
|
||||
|
||||
import exo.shared.config as config
|
||||
from exo.shared.types.common import ModelId
|
||||
from exo.shared.types.events import Event
|
||||
from exo.shared.types.tasks import TaskId
|
||||
@@ -83,21 +83,7 @@ class MlxBuilder(Builder):
|
||||
kv_prefix_cache = KVPrefixCache(self.group)
|
||||
|
||||
device_rank = 0 if self.group is None else self.group.rank()
|
||||
if os.environ.get("EXO_NO_BATCH"):
|
||||
logger.info("using SequentialGenerator (batching disabled)")
|
||||
return SequentialGenerator(
|
||||
model=self.inference_model,
|
||||
tokenizer=self.tokenizer,
|
||||
group=self.group,
|
||||
tool_parser=tool_parser,
|
||||
kv_prefix_cache=kv_prefix_cache,
|
||||
model_id=self.model_id,
|
||||
device_rank=device_rank,
|
||||
cancel_receiver=self.cancel_receiver,
|
||||
event_sender=self.event_sender,
|
||||
vision_processor=vision_processor,
|
||||
)
|
||||
else:
|
||||
if config.app().continuous_batching_enabled:
|
||||
logger.info("using BatchGenerator")
|
||||
return BatchGenerator(
|
||||
model=self.inference_model,
|
||||
@@ -111,3 +97,17 @@ class MlxBuilder(Builder):
|
||||
event_sender=self.event_sender,
|
||||
vision_processor=vision_processor,
|
||||
)
|
||||
else:
|
||||
logger.info("using SequentialGenerator (batching disabled)")
|
||||
return SequentialGenerator(
|
||||
model=self.inference_model,
|
||||
tokenizer=self.tokenizer,
|
||||
group=self.group,
|
||||
tool_parser=tool_parser,
|
||||
kv_prefix_cache=kv_prefix_cache,
|
||||
model_id=self.model_id,
|
||||
device_rank=device_rank,
|
||||
cancel_receiver=self.cancel_receiver,
|
||||
event_sender=self.event_sender,
|
||||
vision_processor=vision_processor,
|
||||
)
|
||||
@@ -5,7 +5,9 @@ from dataclasses import dataclass
|
||||
from typing import Self, cast
|
||||
|
||||
import loguru
|
||||
from exo_rs import AppSettings, BootstrapSettings
|
||||
|
||||
import exo.shared.config as config
|
||||
from exo.shared.types.events import Event
|
||||
from exo.shared.types.tasks import Task, TaskId
|
||||
from exo.shared.types.worker.instances import BoundInstance
|
||||
@@ -42,19 +44,29 @@ def entrypoint(
|
||||
event_sender: MpSender[Event | RunnerTerminationError],
|
||||
task_receiver: MpReceiver[Task],
|
||||
cancel_receiver: MpReceiver[TaskId],
|
||||
settings: tuple[BootstrapSettings, AppSettings],
|
||||
_logger: "loguru.Logger",
|
||||
) -> None:
|
||||
# set global configuration
|
||||
# TODO: in the future I really hope that its not going to be a global :)
|
||||
config.load(*settings)
|
||||
|
||||
global logger
|
||||
logger = _logger
|
||||
|
||||
soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
|
||||
resource.setrlimit(resource.RLIMIT_NOFILE, (min(max(soft, 2048), hard), hard))
|
||||
|
||||
fast_synch_override = os.environ.get("EXO_FAST_SYNCH")
|
||||
if fast_synch_override == "false":
|
||||
os.environ["MLX_METAL_FAST_SYNCH"] = "0"
|
||||
else:
|
||||
os.environ["MLX_METAL_FAST_SYNCH"] = "1"
|
||||
match config.app().fast_synch:
|
||||
case False:
|
||||
os.environ["MLX_METAL_FAST_SYNCH"] = "0"
|
||||
logger.info("FAST_SYNCH forced OFF")
|
||||
case True:
|
||||
os.environ["MLX_METAL_FAST_SYNCH"] = "1"
|
||||
logger.info("FAST_SYNCH forced ON")
|
||||
case None:
|
||||
# By default it is on, but it could change..?
|
||||
os.environ["MLX_METAL_FAST_SYNCH"] = "1"
|
||||
|
||||
logger.info(f"Fast synch flag: {os.environ['MLX_METAL_FAST_SYNCH']}")
|
||||
|
||||
|
||||
@@ -8,7 +8,7 @@ from typing import BinaryIO
|
||||
import mlx.core as mx
|
||||
from mlx_lm.tokenizer_utils import TokenizerWrapper
|
||||
|
||||
from exo.shared.constants import EXO_MAX_CONCURRENT_REQUESTS
|
||||
import exo.shared.config as config
|
||||
from exo.shared.types.chunks import ErrorChunk, GenerationChunk, PrefillProgressChunk
|
||||
from exo.shared.types.common import ModelId
|
||||
from exo.shared.types.events import ChunkGenerated, Event
|
||||
@@ -337,6 +337,7 @@ class BatchGenerator(Engine):
|
||||
_all_tasks: dict[TaskId, TextGeneration] = field(default_factory=dict, init=False)
|
||||
_queue: deque[TextGeneration] = field(default_factory=deque, init=False)
|
||||
_gen: ExoBatchGenerator = field(init=False)
|
||||
_max_concurrent_requests: int = field(init=False)
|
||||
_active_tasks: dict[
|
||||
int,
|
||||
tuple[
|
||||
@@ -347,6 +348,7 @@ class BatchGenerator(Engine):
|
||||
] = field(default_factory=dict, init=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self._max_concurrent_requests = config.app().max_concurrent_requests
|
||||
self._gen = ExoBatchGenerator(
|
||||
model=self.model,
|
||||
tokenizer=self.tokenizer,
|
||||
@@ -406,7 +408,7 @@ class BatchGenerator(Engine):
|
||||
self.agree_on_tasks()
|
||||
|
||||
# Submit any queued tasks to the engine
|
||||
while self._queue and len(self._active_tasks) < EXO_MAX_CONCURRENT_REQUESTS:
|
||||
while self._queue and len(self._active_tasks) < self._max_concurrent_requests:
|
||||
task = self._queue.popleft()
|
||||
try:
|
||||
uid = self._start_task(task)
|
||||
|
||||
@@ -7,7 +7,7 @@ from typing import BinaryIO
|
||||
|
||||
from anyio import ClosedResourceError, EndOfStream
|
||||
|
||||
from exo.shared.constants import ENABLE_DISAGGREGATION
|
||||
import exo.shared.config as config
|
||||
from exo.shared.types.chunks import Chunk
|
||||
from exo.shared.types.common import CommandId
|
||||
from exo.shared.types.events import (
|
||||
@@ -126,7 +126,7 @@ class Runner:
|
||||
self.update_status(RunnerIdle())
|
||||
|
||||
def _start_prefill_server(self) -> int | None:
|
||||
if not ENABLE_DISAGGREGATION:
|
||||
if not config.app().disaggregation_enabled:
|
||||
return None
|
||||
if self.device_rank != 0:
|
||||
return None
|
||||
|
||||
@@ -14,7 +14,7 @@ from anyio import (
|
||||
)
|
||||
from loguru import logger
|
||||
|
||||
from exo.shared.constants import EXO_RUNNER_STDERR_LOG, EXO_RUNNER_STDOUT_LOG
|
||||
import exo.shared.config as config
|
||||
from exo.shared.types.chunks import ErrorChunk
|
||||
from exo.shared.types.events import (
|
||||
ChunkGenerated,
|
||||
@@ -76,9 +76,16 @@ class RunnerStdioHandler:
|
||||
*,
|
||||
stdout_rx: Receiver[bytes],
|
||||
stderr_rx: Receiver[bytes],
|
||||
stdout_log_path: PathLike[str] = EXO_RUNNER_STDOUT_LOG,
|
||||
stderr_log_path: PathLike[str] = EXO_RUNNER_STDERR_LOG,
|
||||
stdout_log_path: PathLike[str] | None = None,
|
||||
stderr_log_path: PathLike[str] | None = None,
|
||||
) -> Self:
|
||||
stdout_log_path = (
|
||||
stdout_log_path or config.bootstrap().log_files.exo_runner_stdout_log
|
||||
)
|
||||
stderr_log_path = (
|
||||
stderr_log_path or config.bootstrap().log_files.exo_runner_stderr_log
|
||||
)
|
||||
|
||||
# these are append only logs used to gather data for log template mining
|
||||
#
|
||||
# TODO: in the future use [Drain3](https://github.com/logpai/Drain3)
|
||||
@@ -218,6 +225,7 @@ class RunnerSupervisor:
|
||||
ev_send,
|
||||
task_recv,
|
||||
cancel_recv,
|
||||
(config.bootstrap(), config.app()),
|
||||
logger,
|
||||
),
|
||||
daemon=True,
|
||||
|
||||
@@ -10,7 +10,7 @@ from typing import Any, cast
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from exo.shared.constants import EXO_DEFAULT_MODELS_DIR
|
||||
import exo.shared.config as config
|
||||
from exo.shared.models.model_cards import ModelCard, ModelTask
|
||||
from exo.shared.types.backends import Backend
|
||||
from exo.shared.types.common import ModelId
|
||||
@@ -53,7 +53,8 @@ def create_hostfile(world_size: int, base_port: int) -> tuple[str, list[str]]:
|
||||
# Use GPT OSS 20b to test as it is a model with a lot of strange behaviour
|
||||
|
||||
DEFAULT_GPT_OSS_CONFIG = PipelineTestConfig(
|
||||
model_path=EXO_DEFAULT_MODELS_DIR / "mlx-community--gpt-oss-20b-MXFP4-Q8",
|
||||
model_path=config.bootstrap().models_dirs.default_models_dir
|
||||
/ "mlx-community--gpt-oss-20b-MXFP4-Q8",
|
||||
total_layers=24,
|
||||
base_port=29600,
|
||||
max_tokens=200,
|
||||
|
||||
@@ -15,7 +15,7 @@ from typing import Any, cast
|
||||
|
||||
import pytest
|
||||
|
||||
from exo.shared.constants import EXO_DEFAULT_MODELS_DIR
|
||||
import exo.shared.config as config
|
||||
from exo.shared.models.model_cards import ModelCard, ModelTask
|
||||
from exo.shared.types.backends import Backend
|
||||
from exo.shared.types.common import ModelId
|
||||
@@ -23,7 +23,10 @@ from exo.shared.types.memory import Memory
|
||||
from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
|
||||
|
||||
MODEL_ID = "mlx-community/gpt-oss-20b-MXFP4-Q8"
|
||||
MODEL_PATH = EXO_DEFAULT_MODELS_DIR / "mlx-community--gpt-oss-20b-MXFP4-Q8"
|
||||
MODEL_PATH = (
|
||||
config.bootstrap().models_dirs.default_models_dir
|
||||
/ "mlx-community--gpt-oss-20b-MXFP4-Q8"
|
||||
)
|
||||
TOTAL_LAYERS = 24
|
||||
MAX_TOKENS = 10
|
||||
SEED = 42
|
||||
|
||||
@@ -14,8 +14,8 @@ import pytest
|
||||
from mlx.utils import tree_flatten, tree_unflatten
|
||||
from mlx_lm.tokenizer_utils import TokenizerWrapper
|
||||
|
||||
import exo.shared.config as config
|
||||
from exo.download.download_utils import resolve_existing_model
|
||||
from exo.shared.constants import EXO_MODELS_DIRS, EXO_MODELS_READ_ONLY_DIRS
|
||||
from exo.shared.types.common import ModelId
|
||||
from exo.shared.types.text_generation import (
|
||||
InputMessage,
|
||||
@@ -106,6 +106,11 @@ def _reduce_config(cfg: dict[str, Any]) -> dict[str, Any]:
|
||||
# ── Helpers ───────────────────────────────────────────────────────────────── #
|
||||
|
||||
|
||||
def _search_dirs():
|
||||
m = config.bootstrap().models_dirs
|
||||
return (*m.models_read_only_dirs, *m.models_dirs)
|
||||
|
||||
|
||||
def _find_snapshot(hub_name: str) -> Path | None:
|
||||
"""Locate a model directory under exo's models dirs.
|
||||
|
||||
@@ -117,7 +122,7 @@ def _find_snapshot(hub_name: str) -> Path | None:
|
||||
if found is not None:
|
||||
return found
|
||||
normalized = model_id.normalize()
|
||||
for search_dir in (*EXO_MODELS_READ_ONLY_DIRS, *EXO_MODELS_DIRS):
|
||||
for search_dir in _search_dirs():
|
||||
candidate = search_dir / normalized
|
||||
if candidate.is_dir():
|
||||
return candidate
|
||||
|
||||
@@ -20,7 +20,7 @@ hosts=("$@")
|
||||
|
||||
for host; do
|
||||
ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
|
||||
"EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix build github:exo-explore/exo/$commit" &
|
||||
"EXO_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix build github:exo-explore/exo/$commit" &
|
||||
done
|
||||
wait
|
||||
|
||||
@@ -35,7 +35,7 @@ trap 'cleanup' EXIT INT TERM
|
||||
|
||||
for host; do
|
||||
ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
|
||||
"EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit" &>/dev/null &
|
||||
"EXO_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit" &>/dev/null &
|
||||
done
|
||||
|
||||
for host; do
|
||||
|
||||
@@ -28,12 +28,12 @@ trap 'cleanup' EXIT INT TERM
|
||||
|
||||
for host; do
|
||||
ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
|
||||
"EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix build github:exo-explore/exo/$commit" &
|
||||
"EXO_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix build github:exo-explore/exo/$commit" &
|
||||
done
|
||||
wait
|
||||
for host; do
|
||||
ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
|
||||
"EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit" &>/dev/null &
|
||||
"EXO_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit" &>/dev/null &
|
||||
done
|
||||
|
||||
for host; do
|
||||
|
||||
@@ -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)
|
||||
@@ -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_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
|
||||
|
||||
@@ -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')" },
|
||||
@@ -523,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" },
|
||||
@@ -622,13 +622,13 @@ requires-dist = [
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "exo-pyo3-bindings"
|
||||
version = "0.2.10"
|
||||
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')" },
|
||||
]
|
||||
@@ -637,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" },
|
||||
]
|
||||
|
||||
Reference in new issue
Block a user