mirror of
https://github.com/exo-explore/exo.git
synced 2026-09-10 12:27:32 -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
@@ -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:
|
||||
# ============================================================
|
||||
|
||||
@@ -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
|
||||
|
||||
Generated
+959
-342
File diff suppressed because it is too large.
Load diff
+42
-33
@@ -1,6 +1,6 @@
|
||||
[workspace]
|
||||
resolver = "3"
|
||||
members = ["rust/exo_rs", "rust/networking"]
|
||||
members = ["rust/exo_rs", "rust/networking", "rust/util"]
|
||||
|
||||
[workspace.package]
|
||||
version = "0.0.1"
|
||||
@@ -20,16 +20,18 @@ opt-level = 3
|
||||
[workspace.dependencies]
|
||||
## Crate members as common dependencies
|
||||
networking = { path = "rust/networking" }
|
||||
util = { path = "rust/util" }
|
||||
|
||||
# pyo3
|
||||
pyo3 = "0.27.2"
|
||||
pyo3-async-runtimes = "0.27.0"
|
||||
pyo3-log = "0.13.2"
|
||||
pyo3-stub-gen = "0.22.2"
|
||||
pyo3 = "0.28.3"
|
||||
pyo3-async-runtimes = "0.28.0"
|
||||
pyo3-log = "0.13.3"
|
||||
pyo3-stub-gen = "0.22.3"
|
||||
|
||||
# util
|
||||
extend = "1.2"
|
||||
delegate = "0.13"
|
||||
derive_more = "2.1"
|
||||
|
||||
# Utility dependencies
|
||||
nix = "0.31"
|
||||
@@ -42,7 +44,6 @@ pin-project = "1.1.10"
|
||||
serde_json = "1.0.149"
|
||||
rand = "0.10.1"
|
||||
parking_lot = "0.12.5"
|
||||
pidfile-rs = "0.3.1"
|
||||
|
||||
# Tracing/logging
|
||||
log = "0.4"
|
||||
@@ -54,35 +55,43 @@ 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]
|
||||
zenoh = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-buffers = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-codec = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-collections = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-config = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-core = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-crypto = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-keyexpr = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-link = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-link-commons = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-link-quic = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-link-quic_datagram = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-link-tcp = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-link-tls = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-link-udp = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-link-unixsock_stream = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-link-ws = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-macros = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-plugin-trait = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-protocol = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-result = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-runtime = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-sync = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-task = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-transport = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-util = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh_backend_traits = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
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
|
||||
|
||||
@@ -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_ZENOH_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.
|
||||
@@ -146,7 +146,7 @@
|
||||
config.treefmt.build.wrapper
|
||||
|
||||
# PYTHON
|
||||
#self'.packages.exo.passthru.evenv
|
||||
self'.packages.exo.passthru.evenv
|
||||
uv
|
||||
|
||||
# RUST
|
||||
|
||||
@@ -211,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}
|
||||
'';
|
||||
|
||||
+20
-2
@@ -23,10 +23,12 @@ workspace = true
|
||||
|
||||
[dependencies]
|
||||
networking.workspace = true
|
||||
util.workspace = true
|
||||
extend.workspace = true
|
||||
itertools = "0.14.0"
|
||||
|
||||
# interop
|
||||
pyo3 = { workspace = true, features = ["experimental-async"] }
|
||||
pyo3 = { workspace = true, features = ["serde", "experimental-async"] }
|
||||
pyo3-stub-gen.workspace = true
|
||||
pyo3-async-runtimes = { workspace = true, features = [
|
||||
"attributes",
|
||||
@@ -35,7 +37,23 @@ pyo3-async-runtimes = { workspace = true, features = [
|
||||
] }
|
||||
pyo3-log.workspace = true
|
||||
|
||||
pidfile-rs = { git = "https://github.com/AndreiCravtov/pidfile-rs" }
|
||||
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"] }
|
||||
|
||||
+398
-22
@@ -2,20 +2,358 @@
|
||||
# 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",
|
||||
"PyFromSwarm",
|
||||
"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: bytes, listen_port: builtins.int, discovery_service_port: builtins.int) -> NetworkingHandle: ...
|
||||
def new(identity: builtins.str, namespace: builtins.str, listen_port: builtins.int, discovery_service_port: builtins.int) -> NetworkingHandle: ...
|
||||
async def recv(self) -> FromSwarm: ...
|
||||
async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
|
||||
r"""
|
||||
Subscribe to a `GossipSub` topic.
|
||||
@@ -34,7 +372,6 @@ class NetworkingHandle:
|
||||
|
||||
If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
|
||||
"""
|
||||
async def recv(self) -> PyFromSwarm: ...
|
||||
|
||||
@typing.final
|
||||
class Pidfile:
|
||||
@@ -42,7 +379,7 @@ class Pidfile:
|
||||
A PID file protected with a lock.
|
||||
|
||||
An instance of `Pidfile` can be used to manage a PID file: create it,
|
||||
lock it, detect already running daemons. It is backed by [`pidfile`][]
|
||||
lock it, detect already running daemons. It is backed by [`pidfile`]
|
||||
functions of `libbsd`/`libutil` which use `flopen` to lock the PID
|
||||
file.
|
||||
|
||||
@@ -61,7 +398,6 @@ class Pidfile:
|
||||
def __new__(cls, path: builtins.str | os.PathLike | pathlib.Path, mode: builtins.int) -> Pidfile:
|
||||
r"""
|
||||
Creates a new PID file and locks it.
|
||||
Writes the current process ID to the PID file.
|
||||
|
||||
If the PID file cannot be locked, returns `PidfileError::AlreadyRunning` with
|
||||
a PID of the already running process, or `None` if no PID has been written to
|
||||
@@ -93,22 +429,62 @@ class PidfileError(builtins.Exception):
|
||||
def __repr__(self) -> builtins.str: ...
|
||||
def __str__(self) -> builtins.str: ...
|
||||
|
||||
class PyFromSwarm:
|
||||
@typing.final
|
||||
class Connection(PyFromSwarm):
|
||||
__match_args__ = ("connected",)
|
||||
@property
|
||||
def connected(self) -> builtins.bool: ...
|
||||
def __new__(cls, connected: builtins.bool) -> PyFromSwarm.Connection: ...
|
||||
@typing.final
|
||||
class RejectedArgs:
|
||||
r"""
|
||||
Rejected arguments go here.
|
||||
|
||||
@typing.final
|
||||
class Message(PyFromSwarm):
|
||||
__match_args__ = ("topic", "data",)
|
||||
@property
|
||||
def topic(self) -> builtins.str: ...
|
||||
@property
|
||||
def data(self) -> bytes: ...
|
||||
def __new__(cls, topic: builtins.str, data: bytes) -> PyFromSwarm.Message: ...
|
||||
|
||||
...
|
||||
# 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: ...
|
||||
|
||||
@@ -19,9 +19,15 @@ dev = ["exo_rs", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"]
|
||||
|
||||
[tool.maturin]
|
||||
module-name = "exo_rs"
|
||||
features = ["pyo3/extension-module", "pyo3/experimental-async"]
|
||||
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);
|
||||
|
||||
@@ -1,6 +1,25 @@
|
||||
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()?;
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -1,44 +0,0 @@
|
||||
use crate::ext::ResultExt as _;
|
||||
use pyo3::types::{PyBytes, PyBytesMethods as _};
|
||||
use pyo3::{Bound, PyResult, Python, pyclass, pymethods};
|
||||
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
|
||||
|
||||
/// Identity keypair of a node.
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass(name = "Keypair", frozen)]
|
||||
#[repr(transparent)]
|
||||
pub struct PyKeypair(pub u128);
|
||||
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
#[allow(clippy::needless_pass_by_value)]
|
||||
impl PyKeypair {
|
||||
/// Generate a new Ed25519 keypair.
|
||||
#[staticmethod]
|
||||
fn generate() -> Self {
|
||||
Self(rand::random())
|
||||
}
|
||||
|
||||
/// Construct an Ed25519 keypair from secret key bytes
|
||||
#[staticmethod]
|
||||
fn from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
|
||||
let bytes = Vec::from(bytes.as_bytes());
|
||||
Ok(Self(u128::from_le_bytes(
|
||||
bytes
|
||||
.try_into()
|
||||
.map_err(|_| "passed too many bytes to from_bytes")
|
||||
.pyerr()?,
|
||||
)))
|
||||
}
|
||||
|
||||
/// Get the secret key bytes underlying the keypair
|
||||
fn to_bytes<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
|
||||
let bytes = self.0.to_le_bytes();
|
||||
Ok(PyBytes::new(py, &bytes))
|
||||
}
|
||||
|
||||
/// Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
|
||||
fn to_node_id(&self) -> String {
|
||||
format!("{:x}", self.0)
|
||||
}
|
||||
}
|
||||
+90
-10
@@ -5,14 +5,17 @@
|
||||
//!
|
||||
|
||||
mod allow_threading;
|
||||
mod pidfile;
|
||||
// mod ident;
|
||||
pub mod config;
|
||||
mod networking;
|
||||
pub mod newtype;
|
||||
mod pidfile;
|
||||
|
||||
use crate::config::config_submodule;
|
||||
use crate::networking::networking_submodule;
|
||||
use crate::pidfile::pidfile_submodule;
|
||||
use pyo3::prelude::PyModule;
|
||||
use pyo3::{Bound, PyResult, 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 crate-wide extension traits/methods
|
||||
@@ -45,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,
|
||||
@@ -83,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()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -143,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_rs")]
|
||||
#[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);
|
||||
@@ -158,9 +239,8 @@ fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
// work with maturin, where the types generate correctly, in the right folder, without
|
||||
// too many importing issues...
|
||||
pidfile_submodule(m)?;
|
||||
// m.add_class::<PyKeypair>()?;
|
||||
// networking_submodule(m)?;
|
||||
networking_submodule(m)?;
|
||||
config_submodule(m)?;
|
||||
|
||||
// top-level constructs
|
||||
// TODO: ...
|
||||
|
||||
@@ -1,30 +1,29 @@
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
|
||||
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, ToSwarm, create_swarm};
|
||||
use networking::swarm::{FromSwarm, Swarm, ToSwarm, create_swarm};
|
||||
use networking::{Session, is_valid_zid};
|
||||
use pyo3::exceptions::{PyRuntimeError, PyValueError};
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::PyBytes;
|
||||
use pyo3::{Bound, Py, PyAny, PyErr, PyResult, Python, pymethods};
|
||||
use pyo3_stub_gen::derive::{
|
||||
gen_methods_from_python, gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods,
|
||||
};
|
||||
use pyo3::{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")]
|
||||
struct PyNetworkingHandle {
|
||||
pub struct PyNetworkingHandle {
|
||||
// channels
|
||||
pub to_swarm: mpsc::Sender<ToSwarm>,
|
||||
pub swarm: Arc<Mutex<Pin<Box<dyn Stream<Item = FromSwarm> + Send>>>>,
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass_complex_enum]
|
||||
#[pyclass]
|
||||
enum PyFromSwarm {
|
||||
#[pyclass(name = "FromSwarm")]
|
||||
pub enum PyFromSwarm {
|
||||
Connection { connected: bool },
|
||||
Message { topic: String, data: Py<PyBytes> },
|
||||
}
|
||||
@@ -41,6 +40,20 @@ impl From<FromSwarm> for PyFromSwarm {
|
||||
}
|
||||
}
|
||||
|
||||
impl PyNetworkingHandle {
|
||||
pub fn from_session(session: Session) -> Self {
|
||||
let (to_swarm, from_client) = mpsc::channel(1024);
|
||||
let swarm = Swarm {
|
||||
from_client,
|
||||
session,
|
||||
};
|
||||
PyNetworkingHandle {
|
||||
swarm: Arc::new(Mutex::new(swarm.into_stream())),
|
||||
to_swarm,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl PyNetworkingHandle {
|
||||
@@ -51,30 +64,31 @@ impl PyNetworkingHandle {
|
||||
// ---- Lifecycle management methods ----
|
||||
|
||||
#[staticmethod]
|
||||
fn new<'py>(
|
||||
identity: Bound<'py, PyBytes>,
|
||||
pub fn new(
|
||||
identity: &str,
|
||||
namespace: &str,
|
||||
listen_port: u16,
|
||||
discovery_service_port: u16,
|
||||
) -> PyResult<PyNetworkingHandle> {
|
||||
// todo: zenoh self assigned peers
|
||||
if listen_port == 0 {
|
||||
todo!();
|
||||
todo!("cannot listen on port 0 yet");
|
||||
}
|
||||
// create communication channels
|
||||
let (to_swarm, from_client) = mpsc::channel(1024);
|
||||
|
||||
// get identity
|
||||
let identity = u128::from_le_bytes(
|
||||
identity
|
||||
.extract::<'_, Vec<u8>>()?
|
||||
.try_into()
|
||||
.map_err(|_| PyValueError::new_err("invalid identity bytes"))?,
|
||||
);
|
||||
if !is_valid_zid(identity) {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"{identity} is not a valid zenoh identity"
|
||||
)));
|
||||
}
|
||||
|
||||
// create networking swarm (within tokio context!! or it crashes)
|
||||
let swarm = pyo3_async_runtimes::tokio::get_runtime()
|
||||
.block_on(create_swarm(
|
||||
identity,
|
||||
namespace,
|
||||
from_client,
|
||||
listen_port,
|
||||
discovery_service_port,
|
||||
@@ -87,18 +101,35 @@ impl PyNetworkingHandle {
|
||||
})
|
||||
}
|
||||
|
||||
#[gen_stub(skip)]
|
||||
fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
// #[gen_stub(override_return_type(
|
||||
// type_repr="typing.Awaitable[FromSwarm]", imports=("typing")
|
||||
// ))]
|
||||
// pub fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
// let swarm = Arc::clone(&self.swarm);
|
||||
// pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
// swarm
|
||||
// .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::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)
|
||||
})
|
||||
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 ----
|
||||
@@ -106,7 +137,7 @@ impl PyNetworkingHandle {
|
||||
/// Subscribe to a `GossipSub` topic.
|
||||
///
|
||||
/// Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
|
||||
async fn gossipsub_subscribe(&self, topic: String) -> PyResult<bool> {
|
||||
pub async fn gossipsub_subscribe(&self, topic: String) -> PyResult<bool> {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
|
||||
// send off request to subscribe
|
||||
@@ -128,7 +159,7 @@ impl PyNetworkingHandle {
|
||||
/// Unsubscribes from a `GossipSub` topic.
|
||||
///
|
||||
/// Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
|
||||
async fn gossipsub_unsubscribe(&self, topic: String) -> PyResult<bool> {
|
||||
pub async fn gossipsub_unsubscribe(&self, topic: String) -> PyResult<bool> {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
|
||||
// send off request to unsubscribe
|
||||
@@ -149,7 +180,7 @@ impl PyNetworkingHandle {
|
||||
/// Publishes a message with multiple topics to the `GossipSub` network.
|
||||
///
|
||||
/// If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
|
||||
async fn gossipsub_publish(&self, topic: String, data: Py<PyBytes>) -> PyResult<()> {
|
||||
pub async fn gossipsub_publish(&self, topic: String, data: Py<PyBytes>) -> PyResult<()> {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
|
||||
// send off request to subscribe
|
||||
@@ -173,15 +204,6 @@ impl PyNetworkingHandle {
|
||||
}
|
||||
}
|
||||
|
||||
pyo3_stub_gen::inventory::submit! {
|
||||
gen_methods_from_python! {
|
||||
r#"
|
||||
class PyNetworkingHandle:
|
||||
async def recv() -> PyFromSwarm: ...
|
||||
"#
|
||||
}
|
||||
}
|
||||
|
||||
pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<PyNetworkingHandle>()?;
|
||||
m.add_class::<PyFromSwarm>()?;
|
||||
|
||||
@@ -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.
|
||||
///
|
||||
@@ -77,7 +77,6 @@ impl PyPidfile {
|
||||
#[pymethods]
|
||||
impl PyPidfile {
|
||||
/// Creates a new PID file and locks it.
|
||||
/// Writes the current process ID to the PID file.
|
||||
///
|
||||
/// If the PID file cannot be locked, returns `PidfileError::AlreadyRunning` with
|
||||
/// a PID of the already running process, or `None` if no PID has been written to
|
||||
|
||||
@@ -1,19 +1,22 @@
|
||||
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,
|
||||
PyFromSwarm,
|
||||
FromSwarm,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sleep_on_multiple_items() -> None:
|
||||
print("PYTHON: starting handle")
|
||||
h = NetworkingHandle.new(os.urandom(16), [], 0)
|
||||
h = NetworkingHandle.new(os.urandom(16).hex().lstrip("0"), "default", 52414, 52413)
|
||||
print("PYTHON: handle started")
|
||||
|
||||
rt = asyncio.create_task(_await_recv(h))
|
||||
@@ -25,25 +28,29 @@ async def test_sleep_on_multiple_items() -> None:
|
||||
await h.gossipsub_publish("topic", b"somehting or other")
|
||||
|
||||
|
||||
def test_pidfile(capsys: CaptureFixture[str]):
|
||||
with capsys.disabled():
|
||||
print("\nbefore python")
|
||||
scoped_lock_file()
|
||||
print("after python")
|
||||
|
||||
|
||||
async def _await_recv(h: NetworkingHandle):
|
||||
while True:
|
||||
event = await h.recv()
|
||||
match event:
|
||||
case PyFromSwarm.Connection() as c:
|
||||
case FromSwarm.Connection() as c:
|
||||
print(f"PYTHON: connection update: {c}")
|
||||
case PyFromSwarm.Message() as m:
|
||||
case FromSwarm.Message() as m:
|
||||
print(f"PYTHON: message: {m}")
|
||||
case _:
|
||||
raise Exception("logical error")
|
||||
|
||||
|
||||
def scoped_lock_file():
|
||||
a = Pidfile("/tmp/lock.pid", 0o0600)
|
||||
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__":
|
||||
|
||||
@@ -15,10 +15,13 @@ zenoh-plugin-trait.workspace = true
|
||||
rand.workspace = true
|
||||
log.workspace = true
|
||||
bytemuck = { workspace = true, features = ["derive"] }
|
||||
socket2.workspace = true
|
||||
blake3.workspace = true
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
env_logger.workspace = true
|
||||
tracing = "0.1.44"
|
||||
smol.workspace = true
|
||||
tracing.workspace = true
|
||||
@@ -6,8 +6,8 @@ use zenoh::{Result, Wait};
|
||||
async fn main() -> Result<()> {
|
||||
zenoh::init_log_from_env_or("info");
|
||||
info!("Opening session...");
|
||||
let cfg = networking::cfg(rand::random(), 52414)?;
|
||||
let session = networking::open(cfg, 52414, 52413).await?;
|
||||
let cfg = networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414)?;
|
||||
let session = networking::open(cfg, "exo", 52414, 52413).await?;
|
||||
let _tok = session
|
||||
.z
|
||||
.liveliness()
|
||||
|
||||
@@ -7,8 +7,8 @@ use zenoh::{Result, Wait};
|
||||
async fn main() -> Result<()> {
|
||||
env_logger::try_init_from_env(Env::new().default_filter_or("info")).expect("logger failed");
|
||||
info!("Opening session...");
|
||||
let cfg = networking::cfg(rand::random(), 52414)?;
|
||||
let session = networking::open(cfg, 52414, 52413).await?;
|
||||
let cfg = networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414)?;
|
||||
let session = networking::open(cfg, "exo", 52414, 52413).await?;
|
||||
let _tok = session
|
||||
.z
|
||||
.liveliness()
|
||||
|
||||
@@ -13,8 +13,8 @@ async fn main() -> Result<()> {
|
||||
.and_then(|it| it.parse::<usize>().ok())
|
||||
.expect("USAGE: put_string <n> -- pub a string of n bytes into stream/data");
|
||||
info!("Opening session...");
|
||||
let cfg = networking::cfg(rand::random(), 52414)?;
|
||||
let session = networking::open(cfg, 52414, 52413).await?;
|
||||
let cfg = networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414)?;
|
||||
let session = networking::open(cfg, "exo", 52414, 52413).await?;
|
||||
let _tok = session
|
||||
.z
|
||||
.liveliness()
|
||||
|
||||
@@ -9,8 +9,8 @@ use zenoh::Result;
|
||||
async fn main() -> Result<()> {
|
||||
env_logger::try_init_from_env(Env::new().default_filter_or("info")).expect("logger failed");
|
||||
info!("Opening session...");
|
||||
let cfg = networking::cfg(rand::random(), 52414)?;
|
||||
let session = networking::open(cfg, 52414, 52413).await?;
|
||||
let cfg = networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414)?;
|
||||
let session = networking::open(cfg, "exo", 52414, 52413).await?;
|
||||
let _tok = session
|
||||
.z
|
||||
.liveliness()
|
||||
|
||||
@@ -9,8 +9,8 @@ use zenoh::{Result, Wait};
|
||||
async fn main() -> Result<()> {
|
||||
env_logger::try_init_from_env(Env::new().default_filter_or("info")).expect("logger failed");
|
||||
info!("Opening session...");
|
||||
let cfg = networking::cfg(rand::random(), 52414)?;
|
||||
let session = networking::open(cfg, 52414, 52413).await?;
|
||||
let cfg = networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414)?;
|
||||
let session = networking::open(cfg, "exo", 52414, 52413).await?;
|
||||
let other_live = session
|
||||
.z
|
||||
.liveliness()
|
||||
|
||||
@@ -1,6 +1,4 @@
|
||||
use std::{
|
||||
env,
|
||||
hash::{DefaultHasher, Hash, Hasher},
|
||||
io,
|
||||
net::{Ipv6Addr, SocketAddr, SocketAddrV6},
|
||||
sync::Arc,
|
||||
@@ -39,16 +37,24 @@ pub struct Discovered {
|
||||
}
|
||||
|
||||
impl Discovery {
|
||||
pub async fn new(zid: ZenohId, listen_port: u16, discovery_port: u16) -> io::Result<Self> {
|
||||
let namespace = {
|
||||
let mut hasher = DefaultHasher::new();
|
||||
env::var("EXO_ZENOH_NAMESPACE")
|
||||
.unwrap_or_else(|_| "exo".to_string())
|
||||
.hash(&mut hasher);
|
||||
hasher.finish().to_le_bytes()
|
||||
};
|
||||
let sock = Arc::new(UdpSocket::bind(format!("[::]:{discovery_port}")).await?);
|
||||
//sock.set_multicast_loop_v6(false)?;
|
||||
pub async fn new(
|
||||
zid: ZenohId,
|
||||
namespace: [u8; 8],
|
||||
listen_port: u16,
|
||||
discovery_port: u16,
|
||||
) -> io::Result<Self> {
|
||||
let sock = socket2::Socket::new(
|
||||
socket2::Domain::IPV6,
|
||||
socket2::Type::DGRAM,
|
||||
Some(socket2::Protocol::UDP),
|
||||
)?;
|
||||
sock.set_reuse_address(true)?;
|
||||
#[cfg(unix)]
|
||||
sock.set_reuse_port(true)?;
|
||||
sock.bind(&SocketAddrV6::new(Ipv6Addr::UNSPECIFIED, discovery_port, 0, 0).into())?;
|
||||
sock.set_nonblocking(true)?;
|
||||
sock.set_multicast_loop_v6(true)?;
|
||||
let sock = Arc::new(UdpSocket::from_std(sock.into())?);
|
||||
let ifaces: Arc<Mutex<Vec<SocketAddrV6>>> = Default::default();
|
||||
let _sync = Mutex::new(
|
||||
netwatcher::watch_interfaces_with_callback({
|
||||
@@ -64,9 +70,12 @@ impl Discovery {
|
||||
}
|
||||
|
||||
match sock.join_multicast_v6(&GROUP, *iface_idx) {
|
||||
Ok(()) => ifaces
|
||||
.lock()
|
||||
.push(SocketAddrV6::new(GROUP, 52413, 0, *iface_idx)),
|
||||
Ok(()) => ifaces.lock().push(SocketAddrV6::new(
|
||||
GROUP,
|
||||
discovery_port,
|
||||
0,
|
||||
*iface_idx,
|
||||
)),
|
||||
Err(e) if e.kind() != io::ErrorKind::AddrInUse => {
|
||||
// skip AddrInUse - just means we've already joined the mv6
|
||||
if let Some(iface) = update.interfaces.get(&iface_idx) {
|
||||
|
||||
+32
-13
@@ -12,11 +12,21 @@ use crate::discovery::Discovery;
|
||||
pub mod discovery;
|
||||
pub mod swarm;
|
||||
|
||||
pub fn cfg(identity: u128, listen_port: u16) -> Result<zenoh::Config> {
|
||||
assert!(listen_port != 0, "must used defined listen port port");
|
||||
pub fn is_valid_zid(identity: &str) -> bool {
|
||||
let mut iter = identity.chars();
|
||||
iter.next()
|
||||
.is_some_and(|c| ('1'..='9').contains(&c) || ('a'..='f').contains(&c))
|
||||
&& iter.all(|c| ('0'..='9').contains(&c) || ('a'..='f').contains(&c))
|
||||
&& identity.len() <= 32
|
||||
}
|
||||
|
||||
pub fn cfg(identity: &str, listen_port: u16) -> Result<zenoh::Config> {
|
||||
assert!(is_valid_zid(identity));
|
||||
assert!(identity.len() <= 32);
|
||||
assert!(listen_port != 0, "must used defined listen port");
|
||||
let mut cfg = zenoh::Config::default();
|
||||
// todo: cleanup
|
||||
cfg.insert_json5("id", &format!("\"{identity:x}\""))?;
|
||||
cfg.insert_json5("id", &format!("\"{identity}\""))?;
|
||||
cfg.insert_json5("mode", "\"router\"")?;
|
||||
cfg.insert_json5("listen/endpoints", &format!("[\"tcp/[::]:{listen_port}\"]"))?;
|
||||
cfg.insert_json5("scouting/multicast/enabled", "false")?;
|
||||
@@ -25,7 +35,7 @@ pub fn cfg(identity: u128, listen_port: u16) -> Result<zenoh::Config> {
|
||||
cfg.insert_json5("adminspace/enabled", "true")?;
|
||||
//cfg.insert_json5("transport/link/tx/batch_size", "9216")?;
|
||||
cfg.insert_json5("transport/link/rx/buffer_size", "16777216")?;
|
||||
cfg.insert_json5("timestamping/enabled", "true")?;
|
||||
//cfg.insert_json5("timestamping/enabled", "true")?;
|
||||
cfg.insert_json5("plugins/storage_manager/__required__", "true")?;
|
||||
cfg.insert_json5(
|
||||
"plugins/storage_manager/storages/mem1",
|
||||
@@ -43,10 +53,16 @@ pub fn cfg(identity: u128, listen_port: u16) -> Result<zenoh::Config> {
|
||||
|
||||
pub async fn open(
|
||||
cfg: zenoh::Config,
|
||||
namespace: &str,
|
||||
listen_port: u16,
|
||||
discovery_service_port: u16,
|
||||
) -> Result<Session> {
|
||||
assert!(listen_port != 0, "must used defined listen port");
|
||||
let namespace: [u8; 8] = {
|
||||
blake3::hash(namespace.as_bytes()).as_bytes()[..8]
|
||||
.try_into()
|
||||
.expect("8 is equal to 8")
|
||||
};
|
||||
let mut plugins = PluginsManager::static_plugins_only();
|
||||
plugins.declare_static_plugin::<StoragesPlugin, _>("storage_manager", true);
|
||||
let mut runtime = zenoh::internal::runtime::RuntimeBuilder::new(cfg)
|
||||
@@ -55,8 +71,9 @@ pub async fn open(
|
||||
.await?;
|
||||
let z = zenoh::session::init(runtime.clone().into()).await?;
|
||||
runtime.start().await?;
|
||||
let mut discovery = Discovery::new(z.zid(), listen_port, discovery_service_port).await?;
|
||||
let _jh = Arc::new(tokio::task::spawn(async move {
|
||||
let mut discovery =
|
||||
Discovery::new(z.zid(), namespace, listen_port, discovery_service_port).await?;
|
||||
let _jh = Arc::new(AbortOnDrop(tokio::task::spawn(async move {
|
||||
loop {
|
||||
let Ok(discovered) = discovery.next().await.inspect_err(|e| {
|
||||
log::warn!("discovery error {e}");
|
||||
@@ -81,17 +98,19 @@ pub async fn open(
|
||||
.connect_peer(&discovered.zid.into(), &[locator])
|
||||
.await;
|
||||
}
|
||||
}));
|
||||
})));
|
||||
Ok(Session { z, _jh })
|
||||
}
|
||||
|
||||
struct AbortOnDrop(JoinHandle<()>);
|
||||
impl Drop for AbortOnDrop {
|
||||
fn drop(&mut self) {
|
||||
self.0.abort();
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Session {
|
||||
pub z: ZSession,
|
||||
_jh: Arc<JoinHandle<()>>,
|
||||
}
|
||||
impl Drop for Session {
|
||||
fn drop(&mut self) {
|
||||
self._jh.abort();
|
||||
}
|
||||
_jh: Arc<AbortOnDrop>,
|
||||
}
|
||||
@@ -70,14 +70,14 @@ impl Swarm {
|
||||
token = discovery.recv_async() => {
|
||||
if let Ok(token) = token {
|
||||
let key_expr = token.key_expr().as_str().to_owned();
|
||||
let nid = key_expr.strip_prefix("nodes/").and_then(|s| s.strip_suffix("/live"));
|
||||
let zid = key_expr.strip_prefix("live/");
|
||||
yield match token.kind() {
|
||||
SampleKind::Put => {
|
||||
log::info!("discovered: {nid:?}");
|
||||
log::info!("discovered: {zid:?}");
|
||||
FromSwarm::Discovered {}
|
||||
}
|
||||
SampleKind::Delete => {
|
||||
log::info!("expired: {nid:?}");
|
||||
log::info!("expired: {zid:?}");
|
||||
FromSwarm::Expired {}
|
||||
}
|
||||
}
|
||||
@@ -96,11 +96,11 @@ async fn register_liveness(
|
||||
) -> Result<(LivelinessToken, Subscriber<FifoChannelHandler<Sample>>)> {
|
||||
let token = session
|
||||
.liveliness()
|
||||
.declare_token(format!("nodes/{}/live", session.zid()))
|
||||
.declare_token(format!("live/{}", session.zid()))
|
||||
.await?;
|
||||
let sub = session
|
||||
.liveliness()
|
||||
.declare_subscriber("nodes/*/live")
|
||||
.declare_subscriber("live/*")
|
||||
.history(true)
|
||||
.await?;
|
||||
Ok((token, sub))
|
||||
@@ -193,13 +193,14 @@ async fn on_message(
|
||||
}
|
||||
|
||||
pub async fn create_swarm(
|
||||
identity: u128,
|
||||
identity: &str,
|
||||
namespace: &str,
|
||||
from_client: mpsc::Receiver<ToSwarm>,
|
||||
listen_port: u16,
|
||||
discovery_service_port: u16,
|
||||
) -> Result<Swarm> {
|
||||
let cfg = crate::cfg(identity, listen_port)?;
|
||||
let session = crate::open(cfg, listen_port, discovery_service_port).await?;
|
||||
let session = crate::open(cfg, namespace, listen_port, discovery_service_port).await?;
|
||||
Ok(Swarm {
|
||||
session,
|
||||
from_client,
|
||||
|
||||
+1
-1
@@ -97,7 +97,7 @@
|
||||
--release \
|
||||
--manylinux off \
|
||||
--manifest-path rust/exo_rs/Cargo.toml \
|
||||
--features "pyo3/extension-module,pyo3/experimental-async" \
|
||||
--features "pyo3/extension-module,pyo3/serde,pyo3/experimental-async" \
|
||||
--interpreter ${pkgs.python313}/bin/python \
|
||||
--out dist
|
||||
'';
|
||||
|
||||
@@ -0,0 +1,22 @@
|
||||
[package]
|
||||
name = "util"
|
||||
version = { workspace = true }
|
||||
edition = { workspace = true }
|
||||
publish = false
|
||||
|
||||
[lib]
|
||||
doctest = false
|
||||
name = "util"
|
||||
path = "src/lib.rs"
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
[dependencies]
|
||||
extend.workspace = true
|
||||
|
||||
expanduser = "1.2"
|
||||
path-clean = "1.0"
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = "3.23"
|
||||
@@ -0,0 +1,18 @@
|
||||
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);
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
|
||||
+39
-153
@@ -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_rs 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
|
||||
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,13 +52,14 @@ class Node:
|
||||
_tg: TaskGroup = field(init=False, default_factory=TaskGroup)
|
||||
|
||||
@classmethod
|
||||
async def create(cls, args: "Args") -> Self:
|
||||
keypair = os.urandom(16)
|
||||
node_id = NodeId(keypair.hex())
|
||||
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,
|
||||
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)
|
||||
@@ -71,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,
|
||||
@@ -99,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(),
|
||||
@@ -146,7 +152,7 @@ class Node:
|
||||
master,
|
||||
api,
|
||||
node_id,
|
||||
args.offline,
|
||||
offline,
|
||||
args.api_port,
|
||||
)
|
||||
|
||||
@@ -274,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
|
||||
@@ -327,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))
|
||||
@@ -335,33 +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"pid = {os.getpid()}")
|
||||
if os.getenv("EXO_LIBP2P_NAMESPACE"):
|
||||
raise ValueError(
|
||||
"EXO_LIBP2P_NAMESPACE has been removed - use EXO_ZENOH_NAMESPACE instead"
|
||||
)
|
||||
logger.info(f"EXO_ZENOH_NAMESPACE: {os.getenv('EXO_ZENOH_NAMESPACE')}")
|
||||
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:
|
||||
raise ValueError("Bootstrap peers has been temporarily removed")
|
||||
|
||||
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)
|
||||
@@ -373,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] = []
|
||||
zenoh_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(
|
||||
"--zenoh-port",
|
||||
type=int,
|
||||
default=0,
|
||||
dest="zenoh_port",
|
||||
help="Fixed TCP port for zenoh 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
|
||||
)
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
from exo_rs import PyFromSwarm
|
||||
from exo_rs import FromSwarm
|
||||
|
||||
from exo.utils.pydantic_ext import FrozenModel
|
||||
|
||||
@@ -9,5 +9,5 @@ class ConnectionMessage(FrozenModel):
|
||||
connected: bool
|
||||
|
||||
@classmethod
|
||||
def from_update(cls, update: PyFromSwarm.Connection) -> "ConnectionMessage":
|
||||
def from_update(cls, update: FromSwarm.Connection) -> "ConnectionMessage":
|
||||
return cls(connected=update.connected)
|
||||
+15
-10
@@ -12,12 +12,12 @@ from anyio import (
|
||||
sleep_forever,
|
||||
)
|
||||
from exo_rs import (
|
||||
FromSwarm,
|
||||
NetworkingHandle,
|
||||
PyFromSwarm,
|
||||
)
|
||||
from loguru import logger
|
||||
|
||||
from exo.shared.constants import EXO_NODE_ZID
|
||||
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
|
||||
@@ -100,12 +100,15 @@ class Router:
|
||||
@classmethod
|
||||
def create(
|
||||
cls,
|
||||
identity: bytes,
|
||||
listen_port: int = 52414,
|
||||
discovery_service_port: int = 52413,
|
||||
identity: str,
|
||||
namespace: str,
|
||||
listen_port: int,
|
||||
discovery_service_port: int,
|
||||
) -> "Router":
|
||||
return cls(
|
||||
handle=NetworkingHandle.new(identity, listen_port, discovery_service_port)
|
||||
handle=NetworkingHandle.new(
|
||||
identity, namespace, listen_port, discovery_service_port
|
||||
)
|
||||
)
|
||||
|
||||
def __init__(self, handle: NetworkingHandle):
|
||||
@@ -186,7 +189,7 @@ class Router:
|
||||
from_swarm = await self._net.recv()
|
||||
logger.debug(from_swarm)
|
||||
match from_swarm:
|
||||
case PyFromSwarm.Message(topic, data):
|
||||
case FromSwarm.Message(topic, data):
|
||||
logger.trace(f"Received message on {topic} with payload {data}")
|
||||
if topic not in self.topic_routers:
|
||||
logger.warning(
|
||||
@@ -195,7 +198,7 @@ class Router:
|
||||
continue
|
||||
router = self.topic_routers[topic]
|
||||
await router.publish_bytes(data)
|
||||
case PyFromSwarm.Connection():
|
||||
case FromSwarm.Connection():
|
||||
message = ConnectionMessage.from_update(from_swarm)
|
||||
logger.trace(
|
||||
f"Received message on connection_messages with payload {message}"
|
||||
@@ -227,14 +230,16 @@ class Router:
|
||||
|
||||
|
||||
def get_node_zid(
|
||||
path: Path = EXO_NODE_ZID,
|
||||
path: Path | None = None,
|
||||
) -> NodeId:
|
||||
"""
|
||||
Obtains the :class:`Keypair` associated with this node-ID.
|
||||
Obtain the :class:`PeerId` by from it.
|
||||
"""
|
||||
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())
|
||||
return NodeId(os.urandom(16).hex().lstrip("0"))
|
||||
|
||||
"""
|
||||
def lock_path(path: str | bytes | PathLike[str] | PathLike[bytes]) -> Path:
|
||||
|
||||
@@ -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_ZID = EXO_CACHE_HOME / "node_zid"
|
||||
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
|
||||
+32
-31
@@ -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,7 +44,27 @@ 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_rs").setLevel(logging.INFO)
|
||||
@@ -56,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,
|
||||
)
|
||||
@@ -68,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,
|
||||
)
|
||||
@@ -77,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),
|
||||
@@ -89,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 == []
|
||||
@@ -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_zid
|
||||
from exo.shared.constants import EXO_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_zid().encode())
|
||||
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_ZID)
|
||||
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_ZID)
|
||||
_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,213 +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_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
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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
-1
@@ -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
|
||||
|
||||
|
||||
Reference in new issue
Block a user