mirror of
https://github.com/exo-explore/exo.git
synced 2026-09-12 13:27:43 -04:00
Compare commits
15
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Evan
|
b0f650072c | ||
|
Evan
|
684a53f368 | ||
|
Evan
|
c0beb6fedd | ||
|
Evan
|
f5862a5579 | ||
|
Evan
|
de022a81ba | ||
|
Evan
|
09c064a17a | ||
|
Evan
|
86f278db95 | ||
|
Evan
|
84d4ddfc77 | ||
|
Evan
|
6e01a9e8f5 | ||
|
Evan
|
86835e2ffe | ||
|
Evan
|
240e273628 | ||
|
Evan
|
dc4d325a4a | ||
|
Evan
|
0b783cc583 | ||
|
Evan
|
e5bee75199 | ||
|
Evan
|
0ce96f5afa |
No files matched your search
@@ -29,7 +29,6 @@ To run EXO from source:
|
||||
git clone https://github.com/exo-explore/exo.git
|
||||
cd exo/dashboard
|
||||
npm install && npm run build && cd ..
|
||||
uv sync --extra mlx
|
||||
uv run exo
|
||||
```
|
||||
|
||||
|
||||
Generated
+54
-29
@@ -1007,6 +1007,7 @@ dependencies = [
|
||||
"serde_json",
|
||||
"tokio",
|
||||
"zenoh",
|
||||
"zenoh-ext",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1842,6 +1843,12 @@ dependencies = [
|
||||
"spin 0.9.8",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "leb128"
|
||||
version = "0.2.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6cc46bac87ef8093eed6f272babb833b6443374399985ac8ed28471ee0918545"
|
||||
|
||||
[[package]]
|
||||
name = "leb128fmt"
|
||||
version = "0.1.0"
|
||||
@@ -5107,7 +5114,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"ahash",
|
||||
"arc-swap",
|
||||
@@ -5157,7 +5164,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-buffers"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"zenoh-collections",
|
||||
]
|
||||
@@ -5165,7 +5172,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-codec"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"tracing",
|
||||
"uhlc",
|
||||
@@ -5176,7 +5183,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-collections"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"ahash",
|
||||
]
|
||||
@@ -5184,7 +5191,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-config"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"json5",
|
||||
"nonempty-collections",
|
||||
@@ -5209,7 +5216,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-core"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"lazy_static",
|
||||
"tokio",
|
||||
@@ -5220,7 +5227,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-crypto"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"aes",
|
||||
"hmac",
|
||||
@@ -5230,10 +5237,29 @@ dependencies = [
|
||||
"zenoh-result",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "zenoh-ext"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"bincode",
|
||||
"flume",
|
||||
"futures",
|
||||
"leb128",
|
||||
"serde",
|
||||
"tokio",
|
||||
"tracing",
|
||||
"uhlc",
|
||||
"zenoh",
|
||||
"zenoh-macros",
|
||||
"zenoh-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "zenoh-keyexpr"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"getrandom 0.2.17",
|
||||
"hashbrown 0.16.1",
|
||||
@@ -5248,7 +5274,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-link"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"zenoh-config",
|
||||
"zenoh-link-commons",
|
||||
@@ -5266,7 +5292,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-link-commons"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"base64",
|
||||
@@ -5302,7 +5328,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-link-quic"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"rustls-webpki",
|
||||
@@ -5318,7 +5344,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-link-quic_datagram"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"rustls-webpki",
|
||||
@@ -5334,7 +5360,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-link-tcp"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"socket2 0.5.10",
|
||||
@@ -5351,7 +5377,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-link-tls"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"base64",
|
||||
@@ -5380,7 +5406,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-link-udp"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"libc",
|
||||
@@ -5402,7 +5428,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-link-unixsock_stream"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"nix 0.29.0",
|
||||
@@ -5420,7 +5446,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-link-ws"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"futures-util",
|
||||
@@ -5440,7 +5466,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-macros"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
@@ -5451,8 +5477,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-plugin-storage-manager"
|
||||
version = "1.9.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b1d9fd55c490af4dc4abb6d51b620ecf6f824a9285a1f9446b59eeec6d66ccf7"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"bincode",
|
||||
@@ -5479,7 +5504,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-plugin-trait"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"git-version",
|
||||
"libloading",
|
||||
@@ -5496,7 +5521,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-protocol"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"const_format",
|
||||
"rand 0.8.6",
|
||||
@@ -5511,7 +5536,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-result"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
]
|
||||
@@ -5519,7 +5544,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-runtime"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"lazy_static",
|
||||
"ron",
|
||||
@@ -5533,7 +5558,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-sync"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"event-listener",
|
||||
@@ -5547,7 +5572,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-task"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"futures",
|
||||
"tokio",
|
||||
@@ -5560,7 +5585,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-transport"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"crossbeam-utils",
|
||||
@@ -5594,7 +5619,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh-util"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"const_format",
|
||||
@@ -5620,7 +5645,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "zenoh_backend_traits"
|
||||
version = "1.9.0"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#aa219b84e55dfbbf55ebd827b492626a470b8ac8"
|
||||
source = "git+https://github.com/evanev7/zenoh.git?branch=exo#2fed58ef264a8415b29452cc37054b2e04c10396"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"const_format",
|
||||
|
||||
@@ -49,6 +49,7 @@ env_logger = "0.11.10"
|
||||
|
||||
# networking
|
||||
zenoh = "=1.9.0"
|
||||
zenoh-ext = "=1.9.0"
|
||||
zenoh-plugin-storage-manager = { version = "=1.9.0", default-features = false }
|
||||
zenoh-plugin-trait = "=1.9.0"
|
||||
netwatcher = "0.6.0"
|
||||
@@ -61,6 +62,7 @@ pidfile-rs = { git = "https://github.com/AndreiCravtov/pidfile-rs" }
|
||||
|
||||
[patch.crates-io]
|
||||
zenoh = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh-ext = { 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" }
|
||||
@@ -86,6 +88,7 @@ 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-plugin-storage-manager = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
zenoh_backend_traits = { git = "https://github.com/evanev7/zenoh.git", branch = "exo" }
|
||||
|
||||
[workspace.lints.rust]
|
||||
|
||||
@@ -118,7 +118,7 @@ Then restart the Nix daemon: `sudo launchctl kickstart -k system/org.nixos.nix-d
|
||||
--force
|
||||
```
|
||||
|
||||
Clone the repo, build the dashboard, install the dependencies, and run exo:
|
||||
Clone the repo, build the dashboard, and run exo:
|
||||
|
||||
```bash
|
||||
# Clone exo
|
||||
@@ -127,9 +127,6 @@ git clone https://github.com/exo-explore/exo
|
||||
# Build dashboard
|
||||
cd exo/dashboard && npm install && npm run build && cd ..
|
||||
|
||||
# Install Python dependencies, including the MLX backend
|
||||
uv sync --extra mlx
|
||||
|
||||
# Run exo
|
||||
uv run exo
|
||||
```
|
||||
@@ -179,7 +176,7 @@ rustup toolchain install nightly
|
||||
|
||||
**Note:** The `macmon` package is macOS-only and not required for Linux.
|
||||
|
||||
Clone the repo, build the dashboard, install the dependencies, and run exo:
|
||||
Clone the repo, build the dashboard, and run exo:
|
||||
|
||||
```bash
|
||||
# Clone exo
|
||||
@@ -188,10 +185,6 @@ git clone https://github.com/exo-explore/exo
|
||||
# Build dashboard
|
||||
cd exo/dashboard && npm install && npm run build && cd ..
|
||||
|
||||
# Install Python dependencies with the MLX backend for your hardware
|
||||
# (NVIDIA: --extra mlx-cuda13 or --extra mlx-cuda12)
|
||||
uv sync --extra mlx-cpu
|
||||
|
||||
# Run exo
|
||||
uv run exo
|
||||
```
|
||||
|
||||
Generated
+3
-4
@@ -8,7 +8,6 @@
|
||||
"name": "exo-dashboard",
|
||||
"version": "1.0.0",
|
||||
"dependencies": {
|
||||
"devalue": "^5.6.2",
|
||||
"highlight.js": "^11.11.1",
|
||||
"katex": "^0.16.27",
|
||||
"marked": "^17.0.1",
|
||||
@@ -2332,9 +2331,9 @@
|
||||
}
|
||||
},
|
||||
"node_modules/devalue": {
|
||||
"version": "5.6.2",
|
||||
"resolved": "https://registry.npmjs.org/devalue/-/devalue-5.6.2.tgz",
|
||||
"integrity": "sha512-nPRkjWzzDQlsejL1WVifk5rvcFi/y1onBRxjaFMjZeR9mFpqu2gmAZ9xUB9/IEanEP/vBtGeGganC/GO1fmufg==",
|
||||
"version": "5.5.0",
|
||||
"resolved": "https://registry.npmjs.org/devalue/-/devalue-5.5.0.tgz",
|
||||
"integrity": "sha512-69sM5yrHfFLJt0AZ9QqZXGCPfJ7fQjvpln3Rq5+PS03LD32Ost1Q9N+eEnaQwGRIriKkMImXD56ocjQmfjbV3w==",
|
||||
"license": "MIT"
|
||||
},
|
||||
"node_modules/enhanced-resolve": {
|
||||
|
||||
@@ -11,6 +11,8 @@
|
||||
"check": "svelte-kit sync && svelte-check --tsconfig ./tsconfig.json"
|
||||
},
|
||||
"devDependencies": {
|
||||
"prettier": "^3.4.2",
|
||||
"prettier-plugin-svelte": "^3.3.3",
|
||||
"@sveltejs/adapter-static": "^3.0.10",
|
||||
"@sveltejs/kit": "^2.48.4",
|
||||
"@sveltejs/vite-plugin-svelte": "^5.0.0",
|
||||
@@ -18,8 +20,6 @@
|
||||
"@types/d3": "^7.4.3",
|
||||
"@types/node": "^22",
|
||||
"d3": "^7.9.0",
|
||||
"prettier": "^3.4.2",
|
||||
"prettier-plugin-svelte": "^3.3.3",
|
||||
"svelte": "^5.0.0",
|
||||
"svelte-check": "^4.0.0",
|
||||
"tailwindcss": "^4.0.0",
|
||||
@@ -28,7 +28,6 @@
|
||||
"vite": "^6.0.0"
|
||||
},
|
||||
"dependencies": {
|
||||
"devalue": "^5.6.2",
|
||||
"highlight.js": "^11.11.1",
|
||||
"katex": "^0.16.27",
|
||||
"marked": "^17.0.1",
|
||||
|
||||
@@ -12,6 +12,8 @@
|
||||
toggleDebugMode,
|
||||
topologyOnlyMode,
|
||||
toggleTopologyOnlyMode,
|
||||
getInstanceFirstShard,
|
||||
type Instance,
|
||||
} from "$lib/stores/app.svelte";
|
||||
|
||||
interface Props {
|
||||
@@ -186,7 +188,7 @@
|
||||
function extractInstanceModelId(instanceWrapped: unknown): string | null {
|
||||
const [, instance] = getTaggedValue(instanceWrapped);
|
||||
if (!instance || typeof instance !== "object") return null;
|
||||
const inst = instance as { shardAssignments?: { modelId?: string } };
|
||||
const inst = instance as Instance;
|
||||
return inst.shardAssignments?.modelId ?? null;
|
||||
}
|
||||
|
||||
@@ -204,11 +206,7 @@
|
||||
else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
|
||||
|
||||
let sharding: string | null = null;
|
||||
const inst = instance as {
|
||||
shardAssignments?: { runnerToShard?: Record<string, unknown> };
|
||||
};
|
||||
const runnerToShard = inst.shardAssignments?.runnerToShard || {};
|
||||
const firstShardWrapped = Object.values(runnerToShard)[0];
|
||||
const firstShardWrapped = getInstanceFirstShard(instance as Instance);
|
||||
if (firstShardWrapped) {
|
||||
const [shardTag] = getTaggedValue(firstShardWrapped);
|
||||
if (shardTag === "PipelineShardMetadata") sharding = "Pipeline";
|
||||
|
||||
@@ -9,6 +9,7 @@
|
||||
createInstanceLink,
|
||||
updateInstanceLink,
|
||||
deleteInstanceLink,
|
||||
getInstanceNodeIds,
|
||||
type Instance,
|
||||
} from "$lib/stores/app.svelte";
|
||||
import { deriveBaseModel, deriveFamily } from "$lib/utils/model_family";
|
||||
@@ -16,7 +17,6 @@
|
||||
type InstanceWrapper = {
|
||||
MlxRingInstance?: Instance;
|
||||
MlxJacclInstance?: Instance;
|
||||
VllmInstance?: Instance;
|
||||
};
|
||||
|
||||
let interval: ReturnType<typeof setInterval> | null = null;
|
||||
@@ -43,13 +43,9 @@
|
||||
const ids = nodeIdentities();
|
||||
for (const [id, raw] of Object.entries(instances())) {
|
||||
const wrapper = raw as InstanceWrapper;
|
||||
const inst =
|
||||
wrapper.MlxRingInstance ??
|
||||
wrapper.MlxJacclInstance ??
|
||||
wrapper.VllmInstance;
|
||||
const inst = wrapper.MlxRingInstance ?? wrapper.MlxJacclInstance;
|
||||
const modelId = inst?.shardAssignments?.modelId ?? "";
|
||||
const nodeToRunner = inst?.shardAssignments?.nodeToRunner ?? {};
|
||||
const nodeIds = Object.keys(nodeToRunner);
|
||||
const nodeIds = getInstanceNodeIds(inst);
|
||||
const nodeNames = nodeIds
|
||||
.map((nodeId) => ids[nodeId]?.friendlyName ?? nodeId.slice(0, 6))
|
||||
.filter((name) => !!name);
|
||||
|
||||
@@ -66,12 +66,40 @@ export interface TopologyData {
|
||||
edges: TopologyEdge[];
|
||||
}
|
||||
|
||||
export type InstanceShard = [nodeId: string, runnerId: string, shard: unknown];
|
||||
|
||||
export interface ShardAssignments {
|
||||
modelId: string;
|
||||
shards: InstanceShard[];
|
||||
primaryOutputNode: number;
|
||||
}
|
||||
|
||||
export interface Instance {
|
||||
shardAssignments?: {
|
||||
modelId?: string;
|
||||
runnerToShard?: Record<string, unknown>;
|
||||
nodeToRunner?: Record<string, string>;
|
||||
};
|
||||
shardAssignments: ShardAssignments;
|
||||
}
|
||||
|
||||
export function getInstanceShards(
|
||||
instance: Instance | null | undefined,
|
||||
): InstanceShard[] {
|
||||
return instance?.shardAssignments.shards ?? [];
|
||||
}
|
||||
|
||||
export function getInstanceRunnerIds(
|
||||
instance: Instance | null | undefined,
|
||||
): string[] {
|
||||
return getInstanceShards(instance).map(([, runnerId]) => runnerId);
|
||||
}
|
||||
|
||||
export function getInstanceNodeIds(
|
||||
instance: Instance | null | undefined,
|
||||
): string[] {
|
||||
return [...new Set(getInstanceShards(instance).map(([nodeId]) => nodeId))];
|
||||
}
|
||||
|
||||
export function getInstanceFirstShard(
|
||||
instance: Instance | null | undefined,
|
||||
): unknown {
|
||||
return getInstanceShards(instance)[0]?.[2];
|
||||
}
|
||||
|
||||
export interface RawInstanceLink {
|
||||
@@ -918,7 +946,7 @@ class AppStore {
|
||||
private extractInstanceModelId(instanceWrapped: unknown): string | null {
|
||||
const [, instance] = this.getTaggedValue(instanceWrapped);
|
||||
if (!instance || typeof instance !== "object") return null;
|
||||
const inst = instance as { shardAssignments?: { modelId?: string } };
|
||||
const inst = instance as Instance;
|
||||
return inst.shardAssignments?.modelId ?? null;
|
||||
}
|
||||
|
||||
@@ -936,11 +964,8 @@ class AppStore {
|
||||
else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
|
||||
|
||||
let sharding: string | null = null;
|
||||
const inst = instance as {
|
||||
shardAssignments?: { runnerToShard?: Record<string, unknown> };
|
||||
};
|
||||
const runnerToShard = inst.shardAssignments?.runnerToShard || {};
|
||||
const firstShardWrapped = Object.values(runnerToShard)[0];
|
||||
const inst = instance as Instance;
|
||||
const firstShardWrapped = getInstanceFirstShard(inst);
|
||||
if (firstShardWrapped) {
|
||||
const [shardTag] = this.getTaggedValue(firstShardWrapped);
|
||||
if (shardTag === "PipelineShardMetadata") sharding = "Pipeline";
|
||||
@@ -2262,7 +2287,7 @@ class AppStore {
|
||||
if (keys.length === 1) {
|
||||
const instance = (instanceWrapper as Record<string, unknown>)[
|
||||
keys[0]
|
||||
] as { shardAssignments?: { modelId?: string } };
|
||||
] as Instance;
|
||||
const instanceModelId = instance?.shardAssignments?.modelId;
|
||||
|
||||
// ensure to only return requestedModelId that matches an instance
|
||||
|
||||
@@ -65,6 +65,11 @@
|
||||
nodeThunderboltBridge,
|
||||
nodeIdentities,
|
||||
isConnected,
|
||||
getInstanceFirstShard,
|
||||
getInstanceNodeIds,
|
||||
getInstanceRunnerIds,
|
||||
getInstanceShards,
|
||||
type Instance,
|
||||
type DownloadProgress,
|
||||
type PlacementPreview,
|
||||
} from "$lib/stores/app.svelte";
|
||||
@@ -998,11 +1003,7 @@
|
||||
if (keys.length !== 1) return new Set();
|
||||
const instance = (instanceWrapped as Record<string, unknown>)[keys[0]];
|
||||
if (!instance || typeof instance !== "object") return new Set();
|
||||
const inst = instance as {
|
||||
shardAssignments?: { nodeToRunner?: Record<string, string> };
|
||||
};
|
||||
if (!inst.shardAssignments?.nodeToRunner) return new Set();
|
||||
return new Set(Object.keys(inst.shardAssignments.nodeToRunner));
|
||||
return new Set(getInstanceNodeIds(instance as Instance));
|
||||
}
|
||||
|
||||
function toggleInstanceDownloadDetails(nodeId: string): void {
|
||||
@@ -1784,13 +1785,7 @@
|
||||
};
|
||||
}
|
||||
|
||||
const inst = instance as {
|
||||
shardAssignments?: {
|
||||
nodeToRunner?: Record<string, string>;
|
||||
runnerToShard?: Record<string, unknown>;
|
||||
modelId?: string;
|
||||
};
|
||||
};
|
||||
const inst = instance as Instance;
|
||||
const instanceModelId = inst.shardAssignments?.modelId;
|
||||
|
||||
if (!instanceModelId) {
|
||||
@@ -1805,16 +1800,7 @@
|
||||
};
|
||||
}
|
||||
|
||||
// Get node IDs assigned to this instance
|
||||
const nodeToRunner = inst.shardAssignments?.nodeToRunner || {};
|
||||
const runnerToShard = inst.shardAssignments?.runnerToShard || {};
|
||||
const runnerToNode: Record<string, string> = {};
|
||||
for (const [nodeId, runnerId] of Object.entries(nodeToRunner)) {
|
||||
runnerToNode[runnerId] = nodeId;
|
||||
}
|
||||
const instanceNodeIds = Object.keys(runnerToShard)
|
||||
.map((runnerId) => runnerToNode[runnerId])
|
||||
.filter(Boolean);
|
||||
const instanceNodeIds = getInstanceNodeIds(inst);
|
||||
|
||||
const result = collectDownloadStatus(instanceModelId, instanceNodeIds);
|
||||
|
||||
@@ -1858,6 +1844,7 @@
|
||||
case "FAILED":
|
||||
return "text-red-400";
|
||||
case "SHUTDOWN":
|
||||
case "SHUTTING DOWN":
|
||||
return "text-gray-400";
|
||||
case "DOWNLOADING":
|
||||
return "text-blue-400";
|
||||
@@ -1865,6 +1852,7 @@
|
||||
case "WARMING UP":
|
||||
case "WAITING":
|
||||
case "INITIALIZING":
|
||||
case "CONNECTING":
|
||||
return "text-yellow-400";
|
||||
case "RUNNING":
|
||||
return "text-teal-400";
|
||||
@@ -1887,10 +1875,7 @@
|
||||
return { statusText: "PREPARING", statusClass: "inactive" };
|
||||
}
|
||||
|
||||
const inst = instance as {
|
||||
shardAssignments?: { runnerToShard?: Record<string, unknown> };
|
||||
};
|
||||
const runnerIds = Object.keys(inst.shardAssignments?.runnerToShard || {});
|
||||
const runnerIds = getInstanceRunnerIds(instance as Instance);
|
||||
|
||||
const statuses = runnerIds
|
||||
.map((rid) => {
|
||||
@@ -1898,14 +1883,15 @@
|
||||
if (!r) return null;
|
||||
const [kind] = getTagged(r);
|
||||
const statusMap: Record<string, string> = {
|
||||
RunnerWaitingForInitialization: "WaitingForInitialization",
|
||||
RunnerInitializingBackend: "InitializingBackend",
|
||||
RunnerWaitingForModel: "WaitingForModel",
|
||||
RunnerIdle: "Idle",
|
||||
RunnerConnecting: "Connecting",
|
||||
RunnerConnected: "Connected",
|
||||
RunnerLoading: "Loading",
|
||||
RunnerLoaded: "Loaded",
|
||||
RunnerWarmingUp: "WarmingUp",
|
||||
RunnerReady: "Ready",
|
||||
RunnerRunning: "Running",
|
||||
RunnerShuttingDown: "ShuttingDown",
|
||||
RunnerShutdown: "Shutdown",
|
||||
RunnerFailed: "Failed",
|
||||
};
|
||||
@@ -1959,14 +1945,15 @@
|
||||
return { statusText: "RUNNING", statusClass: "running" };
|
||||
if (has("Ready")) return { statusText: "READY", statusClass: "loaded" };
|
||||
if (has("Loaded")) return { statusText: "LOADED", statusClass: "loaded" };
|
||||
if (has("WaitingForModel"))
|
||||
return { statusText: "WAITING", statusClass: "starting" };
|
||||
if (has("InitializingBackend"))
|
||||
return { statusText: "INITIALIZING", statusClass: "starting" };
|
||||
if (has("WaitingForInitialization"))
|
||||
if (has("Connected"))
|
||||
return { statusText: "INITIALIZING", statusClass: "starting" };
|
||||
if (has("Connecting"))
|
||||
return { statusText: "CONNECTING", statusClass: "starting" };
|
||||
if (has("Idle")) return { statusText: "WAITING", statusClass: "starting" };
|
||||
if (has("ShuttingDown"))
|
||||
return { statusText: "SHUTTING DOWN", statusClass: "inactive" };
|
||||
|
||||
return { statusText: "RUNNING", statusClass: "active" };
|
||||
return { statusText: "PREPARING", statusClass: "inactive" };
|
||||
}
|
||||
|
||||
function getBytes(value: unknown): number {
|
||||
@@ -2039,7 +2026,7 @@
|
||||
function getInstanceModelId(instanceWrapped: unknown): string {
|
||||
const [, instance] = getTagged(instanceWrapped);
|
||||
if (!instance || typeof instance !== "object") return "Unknown";
|
||||
const inst = instance as { shardAssignments?: { modelId?: string } };
|
||||
const inst = instance as Instance;
|
||||
return inst.shardAssignments?.modelId || "Unknown Model";
|
||||
}
|
||||
|
||||
@@ -2067,17 +2054,11 @@
|
||||
if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
|
||||
else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
|
||||
|
||||
const inst = instance as {
|
||||
shardAssignments?: {
|
||||
nodeToRunner?: Record<string, string>;
|
||||
runnerToShard?: Record<string, unknown>;
|
||||
};
|
||||
};
|
||||
const inst = instance as Instance;
|
||||
|
||||
// Sharding strategy from first shard
|
||||
let sharding = "Unknown";
|
||||
const runnerToShard = inst.shardAssignments?.runnerToShard || {};
|
||||
const firstShardWrapped = Object.values(runnerToShard)[0];
|
||||
const firstShardWrapped = getInstanceFirstShard(inst);
|
||||
if (firstShardWrapped) {
|
||||
const [shardTag] = getTagged(firstShardWrapped);
|
||||
if (shardTag === "PipelineShardMetadata") sharding = "Pipeline";
|
||||
@@ -2087,8 +2068,7 @@
|
||||
}
|
||||
|
||||
// Node names from topology
|
||||
const nodeToRunner = inst.shardAssignments?.nodeToRunner || {};
|
||||
const nodeIds = Object.keys(nodeToRunner);
|
||||
const nodeIds = getInstanceNodeIds(inst);
|
||||
const nodeNames = nodeIds.map((nodeId) => {
|
||||
const node = data?.nodes?.[nodeId];
|
||||
return node?.friendly_name || nodeId.slice(0, 8);
|
||||
@@ -2192,35 +2172,19 @@
|
||||
}
|
||||
|
||||
function getOrderedRunnerNodes(
|
||||
instance: Record<string, unknown>,
|
||||
instance: Instance,
|
||||
shardType: "Pipeline" | "Tensor",
|
||||
) {
|
||||
const runnerToShard =
|
||||
(
|
||||
instance.shardAssignments as
|
||||
| { runnerToShard?: Record<string, unknown> }
|
||||
| undefined
|
||||
)?.runnerToShard || {};
|
||||
const nodeToRunner =
|
||||
(
|
||||
instance.shardAssignments as
|
||||
| { nodeToRunner?: Record<string, string> }
|
||||
| undefined
|
||||
)?.nodeToRunner || {};
|
||||
const runnerEntries = Object.entries(runnerToShard).map(
|
||||
([runnerId, shardWrapped]) => {
|
||||
const runnerEntries = getInstanceShards(instance).map(
|
||||
([nodeId, runnerId, shardWrapped]) => {
|
||||
const [tag, shard] = getTagged(shardWrapped);
|
||||
const meta = shard as
|
||||
| {
|
||||
modelMeta?: {
|
||||
worldSize?: number;
|
||||
nLayers?: number;
|
||||
deviceRank?: number;
|
||||
};
|
||||
deviceRank?: number;
|
||||
}
|
||||
| undefined;
|
||||
const deviceRank = meta?.modelMeta?.deviceRank ?? 0;
|
||||
return { runnerId, tag, deviceRank };
|
||||
const deviceRank = meta?.deviceRank ?? 0;
|
||||
return { nodeId, runnerId, tag, deviceRank };
|
||||
},
|
||||
);
|
||||
|
||||
@@ -2231,13 +2195,11 @@
|
||||
: r.tag === "TensorShardMetadata",
|
||||
)
|
||||
.sort((a, b) => a.deviceRank - b.deviceRank)
|
||||
.map((r, idx) => {
|
||||
const nodeId = Object.entries(nodeToRunner).find(
|
||||
([, rid]) => rid === r.runnerId,
|
||||
)?.[0];
|
||||
return { nodeId, runnerId: r.runnerId, order: idx };
|
||||
})
|
||||
.filter((item) => item.nodeId);
|
||||
.map((r, idx) => ({
|
||||
nodeId: r.nodeId,
|
||||
runnerId: r.runnerId,
|
||||
order: idx,
|
||||
}));
|
||||
|
||||
return ordered as Array<{
|
||||
nodeId: string;
|
||||
@@ -2281,10 +2243,7 @@
|
||||
|
||||
// Jaccl (RDMA) – show RDMA interfaces from ibvDevices
|
||||
if (instanceTag === "MlxJacclInstance") {
|
||||
const ordered = getOrderedRunnerNodes(
|
||||
instance as Record<string, unknown>,
|
||||
"Tensor",
|
||||
);
|
||||
const ordered = getOrderedRunnerNodes(instance as Instance, "Tensor");
|
||||
const ibvDevices =
|
||||
(instance as { ibvDevices?: Array<Array<string | null>> }).ibvDevices ||
|
||||
[];
|
||||
@@ -2316,10 +2275,7 @@
|
||||
|
||||
// Ring – derive ring order from pipeline shard ranks and pick host IPs from hostsByNode
|
||||
if (instanceTag === "MlxRingInstance") {
|
||||
const ordered = getOrderedRunnerNodes(
|
||||
instance as Record<string, unknown>,
|
||||
"Pipeline",
|
||||
);
|
||||
const ordered = getOrderedRunnerNodes(instance as Instance, "Pipeline");
|
||||
const hostsByNode =
|
||||
(
|
||||
instance as {
|
||||
@@ -2606,6 +2562,7 @@
|
||||
status.statusText === "WARMING UP" ||
|
||||
status.statusText === "WAITING" ||
|
||||
status.statusText === "INITIALIZING" ||
|
||||
status.statusText === "CONNECTING" ||
|
||||
status.statusText === "PREPARING"
|
||||
) {
|
||||
chatLaunchState = "launching";
|
||||
@@ -5108,7 +5065,10 @@
|
||||
{@const isFailed = statusText === "FAILED"}
|
||||
{@const isLoading = statusText === "LOADING"}
|
||||
{@const isWarmingUp =
|
||||
statusText === "WARMING UP" || statusText === "WAITING"}
|
||||
statusText === "WARMING UP" ||
|
||||
statusText === "WAITING" ||
|
||||
statusText === "INITIALIZING" ||
|
||||
statusText === "CONNECTING"}
|
||||
{@const isReady =
|
||||
statusText === "READY" || statusText === "LOADED"}
|
||||
{@const isRunning = statusText === "RUNNING"}
|
||||
@@ -6244,7 +6204,10 @@
|
||||
{@const isFailed = statusText === "FAILED"}
|
||||
{@const isLoading = statusText === "LOADING"}
|
||||
{@const isWarmingUp =
|
||||
statusText === "WARMING UP" || statusText === "WAITING"}
|
||||
statusText === "WARMING UP" ||
|
||||
statusText === "WAITING" ||
|
||||
statusText === "INITIALIZING" ||
|
||||
statusText === "CONNECTING"}
|
||||
{@const isReady =
|
||||
statusText === "READY" || statusText === "LOADED"}
|
||||
{@const isRunning = statusText === "RUNNING"}
|
||||
|
||||
@@ -3,7 +3,11 @@
|
||||
import { fade } from "svelte/transition";
|
||||
import HeaderNav from "$lib/components/HeaderNav.svelte";
|
||||
import IntegrationCard from "$lib/components/IntegrationCard.svelte";
|
||||
import { instances, refreshState } from "$lib/stores/app.svelte";
|
||||
import {
|
||||
instances,
|
||||
refreshState,
|
||||
type Instance,
|
||||
} from "$lib/stores/app.svelte";
|
||||
import { onMount } from "svelte";
|
||||
|
||||
const apiUrl = browser
|
||||
@@ -24,9 +28,7 @@
|
||||
if (values.length > 0) {
|
||||
const instance = values[0];
|
||||
if (instance && typeof instance === "object") {
|
||||
const inst = instance as {
|
||||
shardAssignments?: { modelId?: string };
|
||||
};
|
||||
const inst = instance as Instance;
|
||||
const modelId = inst.shardAssignments?.modelId;
|
||||
if (modelId && !models.includes(modelId)) {
|
||||
models.push(modelId);
|
||||
|
||||
@@ -16,10 +16,10 @@ check:
|
||||
uv run basedpyright --project pyproject.toml
|
||||
|
||||
sync:
|
||||
uv sync --all-packages --extra mlx
|
||||
uv sync --all-packages
|
||||
|
||||
sync-clean:
|
||||
uv sync --all-packages --extra mlx --force-reinstall --no-cache
|
||||
uv sync --all-packages --force-reinstall --no-cache
|
||||
|
||||
rust-rebuild:
|
||||
PYO3_PYTHON="$(uv run python -c 'import sys; print(sys.executable)')" cargo run --bin stub_gen
|
||||
|
||||
@@ -1,36 +0,0 @@
|
||||
model_id = "moonshotai/Kimi-K2.7-Code"
|
||||
n_layers = 61
|
||||
hidden_size = 7168
|
||||
num_key_value_heads = 64
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "kimi"
|
||||
quantization = ""
|
||||
base_model = "Kimi K2.7 Code"
|
||||
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
|
||||
|
||||
context_length = 262144
|
||||
backends = ["MlxMetal", "MlxCuda", "MlxCpu"]
|
||||
[storage_size]
|
||||
in_bytes = 595204986173
|
||||
|
||||
# Vision tower + mm_projector extracted unmodified (bf16) from the official
|
||||
# repo, in the same format as exolabs/Kimi-K2.6-vision; extraction script
|
||||
# included in the weights repo. Vision config is identical to Kimi-K2.6's.
|
||||
[vision]
|
||||
image_token_id = 163605
|
||||
model_type = "kimi_vl"
|
||||
weights_repo = "aidiffuser/Kimi-K2.7-Code-vision"
|
||||
processor_repo = "moonshotai/Kimi-K2.7-Code"
|
||||
|
||||
# Source: https://huggingface.co/moonshotai/Kimi-K2.7-Code
|
||||
# (recommends temperature 1.0 / top_p 0.95 for thinking mode, same as K2.6)
|
||||
[sampling_defaults]
|
||||
temperature = 1.0
|
||||
top_p = 0.95
|
||||
min_p = 0.01
|
||||
|
||||
[sampling_defaults.non_thinking]
|
||||
temperature = 0.6
|
||||
top_p = 0.95
|
||||
min_p = 0.01
|
||||
@@ -48,6 +48,7 @@ env_logger.workspace = true
|
||||
|
||||
# Networking
|
||||
zenoh.workspace = true
|
||||
zenoh-ext = { workspace = true, features = ["unstable"] }
|
||||
rand.workspace = true
|
||||
serde_json.workspace = true
|
||||
parking_lot.workspace = true
|
||||
@@ -2,16 +2,34 @@
|
||||
# ruff: noqa: E501, F401, F403, F405
|
||||
|
||||
import builtins
|
||||
import collections.abc
|
||||
import os
|
||||
import pathlib
|
||||
import typing
|
||||
__all__ = [
|
||||
"ClearingLVSubscriber",
|
||||
"FromSwarm",
|
||||
"LVAggregator",
|
||||
"LVPublisher",
|
||||
"LVSubscriber",
|
||||
"Mailbox",
|
||||
"NetworkingHandle",
|
||||
"Pidfile",
|
||||
"PidfileError",
|
||||
"SessionHandle",
|
||||
"Storage",
|
||||
"StorageGetter",
|
||||
"TaskChunkSender",
|
||||
"TaskRequest",
|
||||
"TaskRequester",
|
||||
"TaskResponder",
|
||||
"TaskStream",
|
||||
]
|
||||
|
||||
@typing.final
|
||||
class ClearingLVSubscriber:
|
||||
def recv(self) -> collections.abc.Awaitable[tuple[str, str | None] | None]: ...
|
||||
|
||||
class FromSwarm:
|
||||
@typing.final
|
||||
class Connection(FromSwarm):
|
||||
@@ -31,6 +49,23 @@ class FromSwarm:
|
||||
|
||||
...
|
||||
|
||||
@typing.final
|
||||
class LVAggregator:
|
||||
def dump(self) -> builtins.dict[builtins.str, builtins.str]: ...
|
||||
|
||||
@typing.final
|
||||
class LVPublisher:
|
||||
def put(self, data: builtins.str) -> collections.abc.Awaitable[None]: ...
|
||||
def delete(self) -> collections.abc.Awaitable[None]: ...
|
||||
|
||||
@typing.final
|
||||
class LVSubscriber:
|
||||
def recv(self) -> collections.abc.Awaitable[tuple[str, str | None] | None]: ...
|
||||
|
||||
@typing.final
|
||||
class Mailbox:
|
||||
def recv(self) -> collections.abc.Awaitable[str | None]: ...
|
||||
|
||||
@typing.final
|
||||
class NetworkingHandle:
|
||||
@staticmethod
|
||||
@@ -111,3 +146,56 @@ class PidfileError(builtins.Exception):
|
||||
def __repr__(self) -> builtins.str: ...
|
||||
def __str__(self) -> builtins.str: ...
|
||||
|
||||
@typing.final
|
||||
class SessionHandle:
|
||||
@staticmethod
|
||||
def new(identity: builtins.str, namespace: builtins.str, listen_port: builtins.int, discovery_service_port: builtins.int) -> tuple[SessionHandle, NetworkingHandle]: ...
|
||||
def last_value_aggregator(self, prefix: builtins.str) -> LVAggregator: ...
|
||||
def last_value_subscriber(self, kexpr: builtins.str) -> LVSubscriber: ...
|
||||
def last_value_publisher(self, kexpr: builtins.str) -> LVPublisher: ...
|
||||
def clearing_last_value_subscriber(self, kexpr: builtins.str) -> ClearingLVSubscriber:
|
||||
r"""
|
||||
An LV subscriber which synthesizes delete events for offline nodes
|
||||
"""
|
||||
def storage_interface(self) -> Storage: ...
|
||||
def task_requester(self) -> TaskRequester: ...
|
||||
def task_responder(self, instance_id: builtins.str) -> TaskResponder: ...
|
||||
def send_mail(self, node_ids: typing.Sequence[builtins.str], payload: builtins.str) -> collections.abc.Awaitable[None]: ...
|
||||
def mailbox(self, node_id: builtins.str) -> Mailbox: ...
|
||||
|
||||
@typing.final
|
||||
class Storage:
|
||||
def get(self, key: builtins.str) -> collections.abc.Awaitable[str | None]: ...
|
||||
def get_many(self, key: builtins.str) -> StorageGetter: ...
|
||||
def put(self, key: builtins.str, data: builtins.str) -> collections.abc.Awaitable[None]: ...
|
||||
def delete(self, key: builtins.str) -> collections.abc.Awaitable[None]: ...
|
||||
def dump(self, prefix: builtins.str) -> collections.abc.Awaitable[dict[str, str]]: ...
|
||||
|
||||
@typing.final
|
||||
class StorageGetter:
|
||||
def recv(self) -> collections.abc.Awaitable[tuple[str, str] | None]: ...
|
||||
|
||||
@typing.final
|
||||
class TaskChunkSender:
|
||||
def send(self, chunk: builtins.str) -> collections.abc.Awaitable[None]: ...
|
||||
|
||||
@typing.final
|
||||
class TaskRequest:
|
||||
def reply(self, payload: builtins.str) -> None: ...
|
||||
def reply_err(self, payload: builtins.str) -> None: ...
|
||||
|
||||
@typing.final
|
||||
class TaskRequester:
|
||||
def submit(self, instance_id: builtins.str, command_id: builtins.str, command: builtins.str) -> collections.abc.Awaitable[TaskStream]: ...
|
||||
def interrupt(self, instance_id: builtins.str, command_id: builtins.str, command: builtins.str) -> collections.abc.Awaitable[None]: ...
|
||||
|
||||
@typing.final
|
||||
class TaskResponder:
|
||||
def assign_task(self, task_id: builtins.str, task: builtins.str) -> collections.abc.Awaitable[None]: ...
|
||||
def unassign_task(self, task_id: builtins.str) -> collections.abc.Awaitable[None]: ...
|
||||
def recv(self) -> collections.abc.Awaitable[tuple[TaskRequest, TaskChunkSender, str | None] | None]: ...
|
||||
|
||||
@typing.final
|
||||
class TaskStream:
|
||||
def recv(self) -> collections.abc.Awaitable[str | None]: ...
|
||||
|
||||
@@ -0,0 +1,201 @@
|
||||
use networking::{AbortOnDrop, Session, liveliness_aggregator::LivelinessAggregator};
|
||||
use std::{
|
||||
collections::{HashMap, HashSet},
|
||||
sync::Arc,
|
||||
};
|
||||
use tokio::sync::mpsc;
|
||||
use zenoh::{Result as ZResult, Wait};
|
||||
|
||||
use parking_lot::Mutex;
|
||||
use pyo3::{
|
||||
exceptions::{PyConnectionError, PyRuntimeError},
|
||||
prelude::*,
|
||||
};
|
||||
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
|
||||
use zenoh::{
|
||||
handlers::FifoChannelHandler,
|
||||
sample::{Sample, SampleKind},
|
||||
};
|
||||
use zenoh_ext::{
|
||||
AdvancedPublisher, AdvancedSubscriber, AdvancedSubscriberBuilderExt, HistoryConfig,
|
||||
};
|
||||
|
||||
use crate::sample_to_string;
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass]
|
||||
pub struct LVAggregator {
|
||||
pub prefix: Arc<str>,
|
||||
pub store: Arc<Mutex<HashMap<String, String>>>,
|
||||
pub current_live: LivelinessAggregator,
|
||||
}
|
||||
|
||||
pub fn spawn_lv_aggregator_onto(session: &Session, prefix: Arc<str>) -> ZResult<LVAggregator> {
|
||||
// nota bene: config must be kept in line with SessionHandle::last_value_receiver
|
||||
let store = Arc::new(Mutex::new(HashMap::default()));
|
||||
session
|
||||
.z
|
||||
.declare_subscriber(format!("{prefix}/**"))
|
||||
.advanced()
|
||||
.history(
|
||||
HistoryConfig::default()
|
||||
.max_samples(1)
|
||||
.detect_late_publishers(),
|
||||
)
|
||||
.callback({
|
||||
let store = Arc::clone(&store);
|
||||
let prefix = Arc::clone(&prefix);
|
||||
move |sample| {
|
||||
if let Some(s) = sample
|
||||
.key_expr()
|
||||
.to_string()
|
||||
.strip_prefix(&*prefix)
|
||||
.and_then(|it| it.strip_prefix('/'))
|
||||
{
|
||||
let s = s.to_string();
|
||||
match sample.kind() {
|
||||
SampleKind::Put => {
|
||||
store.lock().insert(s, sample_to_string(sample));
|
||||
}
|
||||
SampleKind::Delete => {
|
||||
store.lock().remove(&s);
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
})
|
||||
.background()
|
||||
.wait()?;
|
||||
Ok(LVAggregator {
|
||||
prefix,
|
||||
store,
|
||||
current_live: session.liveliness_aggregator.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl LVAggregator {
|
||||
pub fn dump(&self) -> HashMap<String, String> {
|
||||
let mut store = self.store.lock();
|
||||
let currently_alive: HashSet<String> = self.current_live.dump();
|
||||
// remove any keys that are no longer live
|
||||
store.retain(|key, _| {
|
||||
currently_alive.iter().any(|node_id| {
|
||||
key.strip_prefix(node_id)
|
||||
.is_some_and(|rest| rest.starts_with("/"))
|
||||
})
|
||||
});
|
||||
store.clone()
|
||||
}
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass]
|
||||
pub struct ClearingLVSubscriber {
|
||||
pub receiver: Arc<tokio::sync::Mutex<mpsc::UnboundedReceiver<(String, Option<String>)>>>,
|
||||
pub handle: AbortOnDrop,
|
||||
}
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl ClearingLVSubscriber {
|
||||
#[gen_stub(override_return_type(
|
||||
type_repr="collections.abc.Awaitable[tuple[str, str | None] | None]",
|
||||
imports=("collections.abc")
|
||||
))]
|
||||
pub fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, {
|
||||
let receiver = Arc::clone(&self.receiver);
|
||||
async move { Ok(receiver.lock().await.recv().await) }
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass]
|
||||
pub struct LVSubscriber {
|
||||
pub subscriber: AdvancedSubscriber<FifoChannelHandler<Sample>>,
|
||||
}
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl LVSubscriber {
|
||||
#[gen_stub(override_return_type(
|
||||
type_repr="collections.abc.Awaitable[tuple[str, str | None] | None]",
|
||||
imports=("collections.abc")
|
||||
))]
|
||||
pub fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, {
|
||||
if self.subscriber.receiver_count() != 1 {
|
||||
return Err(PyRuntimeError::new_err(
|
||||
"tried to receive twice on the same receiver",
|
||||
));
|
||||
}
|
||||
let subscriber = self.subscriber.clone();
|
||||
async move {
|
||||
Ok(subscriber.recv_async().await.ok().map(|sample| {
|
||||
(
|
||||
sample.key_expr().to_string(),
|
||||
match sample.kind() {
|
||||
SampleKind::Put => Some(sample_to_string(sample)),
|
||||
SampleKind::Delete => None,
|
||||
},
|
||||
)
|
||||
}))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass]
|
||||
pub struct LVPublisher {
|
||||
pub state: Arc<AdvancedPublisher<'static>>,
|
||||
}
|
||||
impl LVPublisher {
|
||||
pub fn new(publisher: AdvancedPublisher<'static>) -> Self {
|
||||
Self {
|
||||
state: Arc::new(publisher),
|
||||
}
|
||||
}
|
||||
}
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl LVPublisher {
|
||||
#[gen_stub(override_return_type(
|
||||
type_repr="collections.abc.Awaitable[None]",
|
||||
imports=("collections.abc")
|
||||
))]
|
||||
pub fn put<'py>(&'py self, py: Python<'py>, data: String) -> PyResult<Bound<'py, PyAny>> {
|
||||
let state = Arc::clone(&self.state);
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, {
|
||||
// clone the data so py can have it back
|
||||
async move {
|
||||
state
|
||||
.put(data)
|
||||
.await
|
||||
.map_err(|e| PyConnectionError::new_err(e.to_string()))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
#[gen_stub(override_return_type(
|
||||
type_repr="collections.abc.Awaitable[None]",
|
||||
imports=("collections.abc")
|
||||
))]
|
||||
pub fn delete<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
let state = Arc::clone(&self.state);
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
state
|
||||
.delete()
|
||||
.await
|
||||
.map_err(|e| PyConnectionError::new_err(e.to_string()))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub fn lv_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<LVPublisher>()?;
|
||||
m.add_class::<LVSubscriber>()?;
|
||||
m.add_class::<LVAggregator>()?;
|
||||
Ok(())
|
||||
}
|
||||
+28
-11
@@ -4,16 +4,34 @@
|
||||
//!
|
||||
//!
|
||||
|
||||
mod allow_threading;
|
||||
// mod ident;
|
||||
mod networking;
|
||||
mod pidfile;
|
||||
pub mod allow_threading;
|
||||
pub mod last_value;
|
||||
pub mod mailbox;
|
||||
pub mod networking;
|
||||
pub mod pidfile;
|
||||
pub mod session;
|
||||
pub mod storage;
|
||||
pub mod task;
|
||||
|
||||
use crate::last_value::lv_submodule;
|
||||
use crate::mailbox::mailbox_module;
|
||||
use crate::networking::networking_submodule;
|
||||
use crate::pidfile::pidfile_submodule;
|
||||
use crate::session::session_submodule;
|
||||
use crate::storage::storage_submodule;
|
||||
use crate::task::task_submodule;
|
||||
use pyo3::prelude::PyModule;
|
||||
use pyo3::{Bound, PyResult, pymodule};
|
||||
use pyo3_stub_gen::define_stub_info_gatherer;
|
||||
use zenoh::sample::Sample;
|
||||
|
||||
pub(crate) fn sample_to_string(sample: Sample) -> String {
|
||||
sample
|
||||
.payload()
|
||||
.try_to_string()
|
||||
.expect("we only use utf8 encoded strings. someone messed up")
|
||||
.to_string()
|
||||
}
|
||||
|
||||
/// Namespace for crate-wide extension traits/methods
|
||||
pub(crate) mod ext {
|
||||
@@ -154,15 +172,14 @@ fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
builder.enable_all();
|
||||
pyo3_async_runtimes::tokio::init(builder);
|
||||
|
||||
// TODO: for now this is all NOT a submodule, but figure out how to make the submodule system
|
||||
// work with maturin, where the types generate correctly, in the right folder, without
|
||||
// too many importing issues...
|
||||
// TODO: for now this is all NOT a submodule. KISS
|
||||
pidfile_submodule(m)?;
|
||||
// m.add_class::<PyKeypair>()?;
|
||||
networking_submodule(m)?;
|
||||
|
||||
// top-level constructs
|
||||
// TODO: ...
|
||||
lv_submodule(m)?;
|
||||
session_submodule(m)?;
|
||||
storage_submodule(m)?;
|
||||
task_submodule(m)?;
|
||||
mailbox_module(m)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
use pyo3::{exceptions::PyRuntimeError, prelude::*};
|
||||
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
|
||||
use zenoh::pubsub::Subscriber;
|
||||
use zenoh::{
|
||||
handlers::FifoChannelHandler,
|
||||
sample::{Sample, SampleKind},
|
||||
};
|
||||
|
||||
use crate::sample_to_string;
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass]
|
||||
pub struct Mailbox {
|
||||
pub subscriber: Subscriber<FifoChannelHandler<Sample>>,
|
||||
}
|
||||
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl Mailbox {
|
||||
#[gen_stub(override_return_type(
|
||||
type_repr="collections.abc.Awaitable[str | None]",
|
||||
imports=("collections.abc")
|
||||
))]
|
||||
pub fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, {
|
||||
if self.subscriber.receiver_count() != 1 {
|
||||
return Err(PyRuntimeError::new_err(
|
||||
"tried to receive twice on the same receiver",
|
||||
));
|
||||
}
|
||||
let subscriber = self.subscriber.clone();
|
||||
let kexpr = self.subscriber.key_expr().clone();
|
||||
async move {
|
||||
loop {
|
||||
match subscriber.recv_async().await {
|
||||
Ok(sample) if sample.kind() == SampleKind::Delete => continue,
|
||||
Err(_) => {
|
||||
return Ok(None);
|
||||
}
|
||||
Ok(sample) => {
|
||||
if *sample.key_expr() != kexpr {
|
||||
continue;
|
||||
}
|
||||
return Ok(Some(sample_to_string(sample)));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub fn mailbox_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<Mailbox>()?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,275 @@
|
||||
use networking::{AbortOnDrop, Session};
|
||||
use parking_lot::Mutex;
|
||||
use pyo3::{
|
||||
exceptions::{PyConnectionError, PyRuntimeError, PyValueError},
|
||||
prelude::*,
|
||||
};
|
||||
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
|
||||
use std::{
|
||||
collections::{HashMap, HashSet},
|
||||
sync::Arc,
|
||||
};
|
||||
use tokio::sync::mpsc;
|
||||
use zenoh::{Wait, qos::CongestionControl, sample::SampleKind};
|
||||
use zenoh_ext::{
|
||||
AdvancedPublisherBuilderExt, AdvancedSubscriberBuilderExt, CacheConfig, HistoryConfig,
|
||||
MissDetectionConfig,
|
||||
};
|
||||
|
||||
use crate::{
|
||||
last_value::{
|
||||
ClearingLVSubscriber, LVAggregator, LVPublisher, LVSubscriber, spawn_lv_aggregator_onto,
|
||||
},
|
||||
mailbox::Mailbox,
|
||||
networking::PyNetworkingHandle,
|
||||
sample_to_string,
|
||||
storage::Storage,
|
||||
task::{TaskRequester, TaskResponder},
|
||||
};
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass]
|
||||
pub struct SessionHandle {
|
||||
pub session: Session,
|
||||
}
|
||||
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl SessionHandle {
|
||||
#[staticmethod]
|
||||
pub fn new<'py>(
|
||||
identity: &str,
|
||||
namespace: &str,
|
||||
listen_port: u16,
|
||||
discovery_service_port: u16,
|
||||
) -> PyResult<(SessionHandle, PyNetworkingHandle)> {
|
||||
// get identity
|
||||
if !identity
|
||||
.chars()
|
||||
.all(|c| ('0'..='9').contains(&c) || ('a'..='f').contains(&c))
|
||||
|| identity.len() > 32
|
||||
{
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"{identity} is not a valid zenoh identity"
|
||||
)));
|
||||
}
|
||||
|
||||
let cfg = networking::cfg(identity, listen_port).map_err(|e| {
|
||||
PyValueError::new_err(format!("failed to write config: {}", e.to_string()))
|
||||
})?;
|
||||
let session = pyo3_async_runtimes::tokio::get_runtime()
|
||||
.block_on(networking::open(
|
||||
cfg,
|
||||
namespace,
|
||||
listen_port,
|
||||
discovery_service_port,
|
||||
))
|
||||
.map_err(|e| {
|
||||
PyRuntimeError::new_err(format!(
|
||||
"failed to spawn networking on tokio runtime: {}",
|
||||
e.to_string()
|
||||
))
|
||||
})?;
|
||||
let legacy = PyNetworkingHandle::from_session(session.clone());
|
||||
Ok((Self { session }, legacy))
|
||||
}
|
||||
|
||||
pub fn last_value_aggregator(&self, prefix: String) -> PyResult<LVAggregator> {
|
||||
spawn_lv_aggregator_onto(&self.session, prefix.into()).map_err(|e| {
|
||||
PyConnectionError::new_err(format!("failed to spawn liveliness aggregator: {e}"))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn last_value_subscriber(&self, kexpr: &str) -> PyResult<LVSubscriber> {
|
||||
// nota bene: config must be kept in track with the LVAggregator
|
||||
self.session
|
||||
.z
|
||||
.declare_subscriber(kexpr)
|
||||
.advanced()
|
||||
.history(
|
||||
HistoryConfig::default()
|
||||
.max_samples(1)
|
||||
.detect_late_publishers(),
|
||||
)
|
||||
.wait()
|
||||
.map_err(|e| PyConnectionError::new_err(format!("failed to declare subscriber: {e}")))
|
||||
.map(|subscriber| LVSubscriber { subscriber })
|
||||
}
|
||||
|
||||
pub fn last_value_publisher(&self, kexpr: String) -> PyResult<LVPublisher> {
|
||||
self.session
|
||||
.z
|
||||
.declare_publisher(kexpr)
|
||||
.advanced()
|
||||
.publisher_detection()
|
||||
.sample_miss_detection(MissDetectionConfig::default())
|
||||
.cache(CacheConfig::default().max_samples(1))
|
||||
.wait()
|
||||
.map_err(|e| PyConnectionError::new_err(format!("failed to declare publisher: {e}")))
|
||||
.map(LVPublisher::new)
|
||||
}
|
||||
|
||||
/// An LV subscriber which synthesizes delete events for offline nodes
|
||||
pub fn clearing_last_value_subscriber(&self, kexpr: &str) -> PyResult<ClearingLVSubscriber> {
|
||||
// nota bene: config must be kept in track with the LVAggregator
|
||||
let (send, recv) = mpsc::unbounded_channel();
|
||||
let subscriber = self
|
||||
.session
|
||||
.z
|
||||
.declare_subscriber(kexpr)
|
||||
.advanced()
|
||||
.history(
|
||||
HistoryConfig::default()
|
||||
.max_samples(1)
|
||||
.detect_late_publishers(),
|
||||
)
|
||||
.wait()
|
||||
.map_err(|e| {
|
||||
PyConnectionError::new_err(format!("failed to declare subscriber: {e}"))
|
||||
})?;
|
||||
let liveliness_subscriber = self
|
||||
.session
|
||||
.z
|
||||
.liveliness()
|
||||
.declare_subscriber("live/*")
|
||||
.history(true)
|
||||
.wait()
|
||||
.map_err(|e| {
|
||||
PyConnectionError::new_err(format!("failed to declare liveliness subscriber: {e}"))
|
||||
})?;
|
||||
let handle = AbortOnDrop(tokio::task::spawn(async move {
|
||||
let mut seen_keys = HashSet::<String>::new();
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
sample = subscriber.recv_async() => {
|
||||
let Ok(sample) = sample else {
|
||||
break;
|
||||
};
|
||||
|
||||
let key = sample.key_expr().to_string();
|
||||
|
||||
match sample.kind() {
|
||||
SampleKind::Put => {
|
||||
let value = sample_to_string(sample);
|
||||
seen_keys.insert(key.clone());
|
||||
let _ = send.send((key, Some(value)));
|
||||
}
|
||||
SampleKind::Delete => {
|
||||
seen_keys.remove(&key);
|
||||
let _ = send.send((key, None));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
sample = liveliness_subscriber.recv_async() => {
|
||||
let Ok(sample) = sample else {
|
||||
break;
|
||||
};
|
||||
if sample.kind() == SampleKind::Put {
|
||||
continue;
|
||||
}
|
||||
let kexpr = sample.key_expr().to_string();
|
||||
let Some(node_id) = kexpr.strip_prefix("live/") else {
|
||||
continue;
|
||||
};
|
||||
|
||||
let deleted: Vec<String> = seen_keys
|
||||
.iter()
|
||||
.filter(|key| key.contains(node_id))
|
||||
.cloned()
|
||||
.collect();
|
||||
|
||||
for key in deleted {
|
||||
seen_keys.remove(&key);
|
||||
let _ = send.send((key, None));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}));
|
||||
Ok(ClearingLVSubscriber {
|
||||
receiver: Arc::new(tokio::sync::Mutex::new(recv)),
|
||||
handle,
|
||||
})
|
||||
}
|
||||
pub fn storage_interface(&self) -> Storage {
|
||||
Storage {
|
||||
session: self.session.z.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn task_requester(&self) -> TaskRequester {
|
||||
TaskRequester {
|
||||
session: self.session.z.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn task_responder(&self, instance_id: String) -> PyResult<TaskResponder> {
|
||||
let queryable = self
|
||||
.session
|
||||
.z
|
||||
.declare_queryable(format!("task/instances/{instance_id}/tasks/*"))
|
||||
.complete(true)
|
||||
.wait()
|
||||
.map_err(|e| {
|
||||
PyConnectionError::new_err(format!("failed to declare task responder: {e}"))
|
||||
})?;
|
||||
Ok(TaskResponder {
|
||||
instance_id,
|
||||
queryable,
|
||||
session: self.session.z.clone(),
|
||||
assignments: Arc::new(Mutex::new(HashMap::new())),
|
||||
})
|
||||
}
|
||||
|
||||
#[gen_stub(override_return_type(
|
||||
type_repr="collections.abc.Awaitable[None]",
|
||||
imports=("collections.abc")
|
||||
))]
|
||||
pub fn send_mail<'py>(
|
||||
&'py self,
|
||||
py: Python<'py>,
|
||||
node_ids: Vec<String>,
|
||||
payload: String,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let session = self.session.z.clone();
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
tokio::task::spawn_blocking(move || {
|
||||
for node_id in node_ids {
|
||||
// TIL: session.z.put() just. is blocking with CongestionControl::Block. Even through its async apis.
|
||||
// TODO: migrate more(?) zenoh calls to a dedicated thread **if necessary**. I do it here as we spin up N puts simultaneously.
|
||||
session
|
||||
.put(format!("mail/{node_id}"), payload.as_bytes())
|
||||
.congestion_control(CongestionControl::Block)
|
||||
.wait()
|
||||
.map_err(|e| {
|
||||
PyConnectionError::new_err(format!(
|
||||
"failed to declare task responder: {e}"
|
||||
))
|
||||
})?;
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
.expect("panic in worker thread")
|
||||
})
|
||||
}
|
||||
|
||||
pub fn mailbox(&self, node_id: String) -> PyResult<Mailbox> {
|
||||
let subscriber = self
|
||||
.session
|
||||
.z
|
||||
.declare_subscriber(format!("mail/{node_id}"))
|
||||
.wait()
|
||||
.map_err(|e| {
|
||||
PyConnectionError::new_err(format!("failed to declare task responder: {e}"))
|
||||
})?;
|
||||
Ok(Mailbox { subscriber })
|
||||
}
|
||||
}
|
||||
|
||||
pub fn session_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<SessionHandle>()?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,187 @@
|
||||
use std::{collections::HashMap, time::Duration};
|
||||
|
||||
use networking::STORAGE_PREFIX;
|
||||
use pyo3::{
|
||||
exceptions::{PyConnectionError, PyRuntimeError, PyValueError},
|
||||
prelude::*,
|
||||
};
|
||||
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
|
||||
use zenoh::{
|
||||
Session as ZSession, Wait, handlers::FifoChannelHandler, query::Reply, sample::SampleKind,
|
||||
};
|
||||
|
||||
use crate::sample_to_string;
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass]
|
||||
pub struct Storage {
|
||||
pub session: ZSession,
|
||||
}
|
||||
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl Storage {
|
||||
#[gen_stub(override_return_type(
|
||||
type_repr="collections.abc.Awaitable[str | None]",
|
||||
imports=("collections.abc")
|
||||
))]
|
||||
pub fn get<'py>(&'py self, py: Python<'py>, key: String) -> PyResult<Bound<'py, PyAny>> {
|
||||
if key.contains('*') {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"{key} is invalid -- Storage.get only supports fixed keys"
|
||||
)));
|
||||
}
|
||||
let session = self.session.clone();
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
let subscriber = session
|
||||
.get(format!("{STORAGE_PREFIX}/{key}"))
|
||||
//.allowed_destination(Locality::SessionLocal)
|
||||
.await
|
||||
.map_err(|e| PyConnectionError::new_err(format!("failed to query storage: {e}")))?;
|
||||
tokio::select! {
|
||||
_ = tokio::time::sleep(Duration::from_secs(1)) => {
|
||||
Ok(None)
|
||||
}
|
||||
reply = subscriber.recv_async() => {
|
||||
Ok(reply.ok()
|
||||
.and_then(|reply| reply.into_result().ok())
|
||||
.filter(|sample| sample.kind() == SampleKind::Put)
|
||||
.map(sample_to_string)
|
||||
)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
pub fn get_many(&self, key: String) -> PyResult<StorageGetter> {
|
||||
self.session
|
||||
.get(key)
|
||||
.wait()
|
||||
.map_err(|e| PyConnectionError::new_err(format!("failed to query storage: {e}")))
|
||||
.map(StorageGetter)
|
||||
}
|
||||
|
||||
#[gen_stub(override_return_type(
|
||||
type_repr="collections.abc.Awaitable[None]",
|
||||
imports=("collections.abc")
|
||||
))]
|
||||
pub fn put<'py>(
|
||||
&'py self,
|
||||
py: Python<'py>,
|
||||
key: String,
|
||||
data: String,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
if key.contains('*') {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"{key} is invalid -- Storage.put only supports fixed keys"
|
||||
)));
|
||||
}
|
||||
let session = self.session.clone();
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
session
|
||||
.put(format!("{STORAGE_PREFIX}/{key}"), data)
|
||||
.await
|
||||
.map_err(|e| PyConnectionError::new_err(format!("failed to query storage: {e}")))
|
||||
})
|
||||
}
|
||||
|
||||
#[gen_stub(override_return_type(
|
||||
type_repr="collections.abc.Awaitable[None]",
|
||||
imports=("collections.abc")
|
||||
))]
|
||||
pub fn delete<'py>(&'py self, py: Python<'py>, key: String) -> PyResult<Bound<'py, PyAny>> {
|
||||
if key.contains('*') {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"{key} is invalid -- Storage.delete only supports fixed keys"
|
||||
)));
|
||||
}
|
||||
let session = self.session.clone();
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
session
|
||||
.delete(format!("{STORAGE_PREFIX}/{key}"))
|
||||
.await
|
||||
.map_err(|e| PyConnectionError::new_err(format!("failed to query storage: {e}")))
|
||||
})
|
||||
}
|
||||
|
||||
#[gen_stub(override_return_type(
|
||||
type_repr="collections.abc.Awaitable[dict[str, str]]",
|
||||
imports=("collections.abc")
|
||||
))]
|
||||
pub fn dump<'py>(&'py self, py: Python<'py>, prefix: String) -> PyResult<Bound<'py, PyAny>> {
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
Ok(networking::read_raw_memory_storage()
|
||||
.await
|
||||
.into_iter()
|
||||
.filter_map(|(key, value)| {
|
||||
Some((
|
||||
key?.as_str().strip_prefix(prefix.as_str())?.to_string(),
|
||||
value
|
||||
.payload
|
||||
.try_to_string()
|
||||
.expect("we only use utf8 encoded strings. someone messed up")
|
||||
.to_string(),
|
||||
))
|
||||
})
|
||||
.collect::<HashMap<String, String>>())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass]
|
||||
pub struct StorageGetter(FifoChannelHandler<Reply>);
|
||||
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl StorageGetter {
|
||||
#[gen_stub(override_return_type(
|
||||
type_repr="collections.abc.Awaitable[tuple[str, str] | None]",
|
||||
imports=("collections.abc")
|
||||
))]
|
||||
fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
if self.0.receiver_count() != 1 {
|
||||
return Err(PyRuntimeError::new_err(
|
||||
"Tried to call StorageGetter.recv twice concurrently",
|
||||
));
|
||||
}
|
||||
let dupe = self.0.clone();
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
let sample = loop {
|
||||
match dupe.recv_async().await {
|
||||
Err(_) => return Ok(None),
|
||||
Ok(reply) => match reply.into_result() {
|
||||
Err(e) => {
|
||||
log::warn!("Ignoring reply error: {e}");
|
||||
continue;
|
||||
}
|
||||
Ok(sample) => match sample.kind() {
|
||||
SampleKind::Put => break sample,
|
||||
SampleKind::Delete => {
|
||||
log::warn!(
|
||||
"Received unexpected DELETE from queryable: {}",
|
||||
sample.key_expr()
|
||||
);
|
||||
continue;
|
||||
}
|
||||
},
|
||||
},
|
||||
};
|
||||
};
|
||||
|
||||
let key = sample
|
||||
.key_expr()
|
||||
.to_string()
|
||||
.strip_prefix(format!("{STORAGE_PREFIX}/").as_str())
|
||||
.expect("invalid storage format encountered")
|
||||
.to_string();
|
||||
|
||||
Ok(Some((key, sample_to_string(sample))))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub fn storage_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<Storage>()?;
|
||||
m.add_class::<StorageGetter>()?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,352 @@
|
||||
use std::{collections::HashMap, sync::Arc, time::Duration};
|
||||
|
||||
use pyo3::{
|
||||
exceptions::{PyConnectionError, PyRuntimeError, PyTimeoutError, PyValueError},
|
||||
prelude::*,
|
||||
};
|
||||
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
|
||||
use zenoh::{
|
||||
Session as ZSession, Wait,
|
||||
handlers::FifoChannelHandler,
|
||||
pubsub::{Publisher, Subscriber},
|
||||
qos::CongestionControl,
|
||||
query::{ConsolidationMode, Query, Queryable},
|
||||
sample::{Sample, SampleKind},
|
||||
};
|
||||
use zenoh_ext::{AdvancedPublisher, AdvancedPublisherBuilderExt, CacheConfig, MissDetectionConfig};
|
||||
|
||||
use parking_lot::Mutex;
|
||||
|
||||
use crate::sample_to_string;
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass]
|
||||
pub struct TaskRequester {
|
||||
pub session: ZSession,
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass]
|
||||
pub struct TaskResponder {
|
||||
pub instance_id: String,
|
||||
pub queryable: Queryable<FifoChannelHandler<Query>>,
|
||||
pub session: ZSession,
|
||||
pub assignments: Arc<Mutex<HashMap<String, Arc<AdvancedPublisher<'static>>>>>,
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass]
|
||||
pub struct TaskRequest {
|
||||
pub query: Query,
|
||||
pub key: String,
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass]
|
||||
pub struct TaskChunkSender {
|
||||
pub publisher: Arc<Publisher<'static>>,
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass]
|
||||
pub struct TaskStream {
|
||||
pub receiver: Subscriber<FifoChannelHandler<Sample>>,
|
||||
}
|
||||
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl TaskRequester {
|
||||
#[gen_stub(override_return_type(
|
||||
type_repr="collections.abc.Awaitable[TaskStream]",
|
||||
imports=("collections.abc")
|
||||
))]
|
||||
pub fn submit<'py>(
|
||||
&'py self,
|
||||
py: Python<'py>,
|
||||
instance_id: String,
|
||||
command_id: String,
|
||||
command: String,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let session = self.session.clone();
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
let receiver = session
|
||||
.declare_subscriber(task_chunks_key(command_id.as_str()))
|
||||
.wait()
|
||||
.map_err(|e| {
|
||||
PyConnectionError::new_err(format!("failed to declare task stream: {e}"))
|
||||
})?;
|
||||
|
||||
request_task_admission(&session, instance_id, command_id, command).await?;
|
||||
Ok(TaskStream { receiver })
|
||||
})
|
||||
}
|
||||
|
||||
#[gen_stub(override_return_type(
|
||||
type_repr="collections.abc.Awaitable[None]",
|
||||
imports=("collections.abc")
|
||||
))]
|
||||
pub fn interrupt<'py>(
|
||||
&'py self,
|
||||
py: Python<'py>,
|
||||
instance_id: String,
|
||||
command_id: String,
|
||||
command: String,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let session = self.session.clone();
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
request_task_admission(&session, instance_id, command_id, command).await
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn task_chunks_key(command_id: &str) -> String {
|
||||
format!("task/commands/{command_id}/chunks")
|
||||
}
|
||||
|
||||
fn task_assignment_key(instance_id: &str, task_id: &str) -> String {
|
||||
format!("task_assignments/{instance_id}/{task_id}")
|
||||
}
|
||||
|
||||
async fn request_task_admission(
|
||||
session: &ZSession,
|
||||
instance_id: String,
|
||||
command_id: String,
|
||||
command: String,
|
||||
) -> PyResult<()> {
|
||||
let replies = session
|
||||
.get(format!("task/instances/{instance_id}/tasks/{command_id}"))
|
||||
.payload(command)
|
||||
.congestion_control(CongestionControl::Block)
|
||||
.consolidation(ConsolidationMode::None)
|
||||
.timeout(Duration::from_secs(5))
|
||||
.wait()
|
||||
.map_err(|e| PyConnectionError::new_err(format!("failed to submit task: {e}")))?;
|
||||
|
||||
let reply = replies.recv_async().await.map_err(|e| {
|
||||
PyConnectionError::new_err(format!("task admission stream closed early: {e}"))
|
||||
})?;
|
||||
|
||||
match reply.into_result() {
|
||||
Ok(sample) => {
|
||||
if sample.kind() == SampleKind::Delete {
|
||||
Err(PyConnectionError::new_err(
|
||||
"task admission replied with delete",
|
||||
))
|
||||
} else {
|
||||
let _ = sample_to_string(sample);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Err(error) => error.payload().try_to_string().map_or_else(
|
||||
|err| {
|
||||
Err(PyRuntimeError::new_err(format!(
|
||||
"task admission failed: {err}"
|
||||
)))
|
||||
},
|
||||
|ok| {
|
||||
if ok == "Timeout" {
|
||||
Err(PyTimeoutError::new_err("task admission timed out"))
|
||||
} else {
|
||||
Err(PyRuntimeError::new_err(format!(
|
||||
"task admission rejected: {ok}"
|
||||
)))
|
||||
}
|
||||
},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl TaskResponder {
|
||||
#[gen_stub(override_return_type(
|
||||
type_repr="collections.abc.Awaitable[None]",
|
||||
imports=("collections.abc")
|
||||
))]
|
||||
pub fn assign_task<'py>(
|
||||
&'py self,
|
||||
py: Python<'py>,
|
||||
task_id: String,
|
||||
task: String,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let session = self.session.clone();
|
||||
let instance_id = self.instance_id.clone();
|
||||
let assignments = Arc::clone(&self.assignments);
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
let publisher = {
|
||||
let mut assignments = assignments.lock();
|
||||
if let Some(publisher) = assignments.get(&task_id) {
|
||||
publisher.clone()
|
||||
} else {
|
||||
let publisher = Arc::new(
|
||||
session
|
||||
.declare_publisher(task_assignment_key(&instance_id, &task_id))
|
||||
.advanced()
|
||||
.publisher_detection()
|
||||
.sample_miss_detection(MissDetectionConfig::default())
|
||||
.cache(CacheConfig::default().max_samples(1))
|
||||
.wait()
|
||||
.map_err(|e| {
|
||||
PyConnectionError::new_err(format!(
|
||||
"failed to declare task assignment publisher: {e}"
|
||||
))
|
||||
})?,
|
||||
);
|
||||
assignments.insert(task_id.clone(), publisher.clone());
|
||||
publisher
|
||||
}
|
||||
};
|
||||
publisher
|
||||
.put(task)
|
||||
.await
|
||||
.map_err(|e| PyConnectionError::new_err(format!("failed to assign task: {e}")))
|
||||
})
|
||||
}
|
||||
|
||||
#[gen_stub(override_return_type(
|
||||
type_repr="collections.abc.Awaitable[None]",
|
||||
imports=("collections.abc")
|
||||
))]
|
||||
pub fn unassign_task<'py>(
|
||||
&'py self,
|
||||
py: Python<'py>,
|
||||
task_id: String,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let session = self.session.clone();
|
||||
let instance_id = self.instance_id.clone();
|
||||
let assignments = Arc::clone(&self.assignments);
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
let publisher = assignments.lock().remove(&task_id);
|
||||
if let Some(publisher) = publisher {
|
||||
publisher.put("").await.map_err(|e| {
|
||||
PyConnectionError::new_err(format!("failed to unassign task: {e}"))
|
||||
})?;
|
||||
}
|
||||
session
|
||||
.delete(task_assignment_key(&instance_id, &task_id))
|
||||
.await
|
||||
.map_err(|e| PyConnectionError::new_err(format!("failed to unassign task: {e}")))
|
||||
})
|
||||
}
|
||||
|
||||
#[gen_stub(override_return_type(
|
||||
type_repr="collections.abc.Awaitable[tuple[TaskRequest, TaskChunkSender, str | None] | None]",
|
||||
imports=("collections.abc")
|
||||
))]
|
||||
pub fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
if self.queryable.receiver_count() != 1 {
|
||||
return Err(pyo3::exceptions::PyRuntimeError::new_err(
|
||||
"called recv twice concurrently",
|
||||
));
|
||||
}
|
||||
let queryable = self.queryable.clone();
|
||||
let session = self.session.clone();
|
||||
let key_prefix = format!("task/instances/{}/tasks/", self.instance_id);
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
loop {
|
||||
match queryable.recv_async().await {
|
||||
Ok(query) => {
|
||||
let query_key = query.key_expr().as_str();
|
||||
if !query_key.starts_with(&key_prefix) {
|
||||
continue;
|
||||
}
|
||||
let key = query_key.to_owned();
|
||||
let command_id = query_key[key_prefix.len()..].to_string();
|
||||
let payload = query.payload().map(|payload| {
|
||||
payload
|
||||
.try_to_string()
|
||||
.expect("we only use utf8 encoded strings. someone messed up")
|
||||
.to_string()
|
||||
});
|
||||
let publisher = session
|
||||
.declare_publisher(task_chunks_key(&command_id))
|
||||
.wait()
|
||||
.map_err(|e| {
|
||||
PyConnectionError::new_err(format!(
|
||||
"failed to declare task chunk sender: {e}"
|
||||
))
|
||||
})?;
|
||||
return Ok(Some((
|
||||
TaskRequest { query, key },
|
||||
TaskChunkSender {
|
||||
publisher: Arc::new(publisher),
|
||||
},
|
||||
payload,
|
||||
)));
|
||||
}
|
||||
Err(_) => return Ok(None),
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl TaskChunkSender {
|
||||
#[gen_stub(override_return_type(
|
||||
type_repr="collections.abc.Awaitable[None]",
|
||||
imports=("collections.abc")
|
||||
))]
|
||||
pub fn send<'py>(&'py self, py: Python<'py>, chunk: String) -> PyResult<Bound<'py, PyAny>> {
|
||||
let publisher = Arc::clone(&self.publisher);
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
publisher
|
||||
.put(chunk)
|
||||
.await
|
||||
.map_err(|e| PyConnectionError::new_err(format!("failed to send task chunk: {e}")))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl TaskRequest {
|
||||
pub fn reply(&self, payload: String) -> PyResult<()> {
|
||||
self.query
|
||||
.reply(&self.key, payload)
|
||||
.wait()
|
||||
.map_err(|e| PyConnectionError::new_err(format!("failed to reply to task query: {e}")))
|
||||
}
|
||||
|
||||
pub fn reply_err(&self, payload: String) -> PyResult<()> {
|
||||
if payload == "Timeout" {
|
||||
return Err(PyValueError::new_err(
|
||||
"Timeout is reserved for zenoh query timeouts",
|
||||
));
|
||||
}
|
||||
self.query
|
||||
.reply_err(payload)
|
||||
.wait()
|
||||
.map_err(|e| PyConnectionError::new_err(format!("failed to reject task query: {e}")))
|
||||
}
|
||||
}
|
||||
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl TaskStream {
|
||||
#[gen_stub(override_return_type(
|
||||
type_repr="collections.abc.Awaitable[str | None]",
|
||||
imports=("collections.abc")
|
||||
))]
|
||||
pub fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
if self.receiver.receiver_count() != 1 {
|
||||
return Err(pyo3::exceptions::PyRuntimeError::new_err(
|
||||
"called recv twice concurrently",
|
||||
));
|
||||
}
|
||||
let receiver = self.receiver.clone();
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
Ok(receiver.recv_async().await.ok().map(sample_to_string))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub fn task_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<TaskChunkSender>()?;
|
||||
m.add_class::<TaskRequester>()?;
|
||||
m.add_class::<TaskRequest>()?;
|
||||
m.add_class::<TaskResponder>()?;
|
||||
m.add_class::<TaskStream>()?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,168 @@
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use zenoh_ext::{
|
||||
AdvancedPublisherBuilderExt, AdvancedSubscriber, AdvancedSubscriberBuilderExt, CacheConfig,
|
||||
HistoryConfig, MissDetectionConfig,
|
||||
};
|
||||
|
||||
use zenoh::{handlers::FifoChannelHandler, sample::Sample};
|
||||
|
||||
// Adjust these imports to your crate/module paths.
|
||||
use exo_rs::{
|
||||
last_value::{LVPublisher, LVSubscriber},
|
||||
session::SessionHandle,
|
||||
};
|
||||
async fn expect_two_values(
|
||||
sub: &AdvancedSubscriber<FifoChannelHandler<Sample>>,
|
||||
key_a: &str,
|
||||
val_a: &str,
|
||||
key_b: &str,
|
||||
val_b: &str,
|
||||
) {
|
||||
use std::collections::HashMap;
|
||||
use tokio::time::{Duration, Instant, timeout};
|
||||
use zenoh::sample::SampleKind;
|
||||
|
||||
let deadline = Instant::now() + Duration::from_secs(10);
|
||||
let mut seen: HashMap<String, String> = HashMap::new();
|
||||
|
||||
loop {
|
||||
if seen.get(key_a).map(String::as_str) == Some(val_a)
|
||||
&& seen.get(key_b).map(String::as_str) == Some(val_b)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
let remaining = deadline.saturating_duration_since(Instant::now());
|
||||
assert!(
|
||||
!remaining.is_zero(),
|
||||
"timed out waiting for both historical samples; expected {key_a}={val_a}, {key_b}={val_b}; seen = {seen:?}"
|
||||
);
|
||||
|
||||
match timeout(remaining.min(Duration::from_millis(750)), sub.recv_async()).await {
|
||||
Ok(Ok(sample)) => {
|
||||
if sample.kind() == SampleKind::Delete {
|
||||
continue;
|
||||
}
|
||||
|
||||
let key = sample.key_expr().to_string();
|
||||
let value = sample
|
||||
.payload()
|
||||
.try_to_string()
|
||||
.expect("payload should be UTF-8")
|
||||
.to_string();
|
||||
|
||||
if key == key_a || key == key_b {
|
||||
eprintln!("received relevant {key} = {value}");
|
||||
seen.insert(key, value);
|
||||
} else {
|
||||
eprintln!("received unrelated {key} = {value}");
|
||||
}
|
||||
}
|
||||
Ok(Err(e)) => panic!("subscriber receive failed: {e}"),
|
||||
Err(_) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn lv_subscriber_receives_last_value_from_multiple_publishers() {
|
||||
let cfg =
|
||||
networking::cfg(&format!("{:x}", rand::random::<u128>()), 52414).expect("create config");
|
||||
let n_session = networking::open(cfg, "exo", 52414, 52413)
|
||||
.await
|
||||
.expect("open session");
|
||||
|
||||
let session = SessionHandle { session: n_session };
|
||||
|
||||
let run_id = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_nanos();
|
||||
|
||||
let base = format!("zenoh_advanced_history_test/{run_id}");
|
||||
|
||||
let key_a = format!("{base}/a");
|
||||
let key_b = format!("{base}/b");
|
||||
let sub_key = format!("{base}/*");
|
||||
|
||||
let pub1: LVPublisher = session
|
||||
.last_value_publisher(key_a.clone())
|
||||
.expect("declare LV publisher a");
|
||||
|
||||
pub1.state.put("aa").await.expect("publish aa");
|
||||
|
||||
let pub2: LVPublisher = session
|
||||
.last_value_publisher(key_b.clone())
|
||||
.expect("declare LV publisher b");
|
||||
|
||||
pub2.state.put("bb").await.expect("publish bb");
|
||||
|
||||
// Let publisher detection / cache metadata settle before the late subscriber joins.
|
||||
tokio::time::sleep(Duration::from_millis(250)).await;
|
||||
|
||||
let sub: LVSubscriber = session
|
||||
.last_value_subscriber(&*sub_key)
|
||||
.expect("declare LV subscriber");
|
||||
|
||||
expect_two_values(&sub.subscriber, &*key_a, "aa", &*key_b, "bb").await
|
||||
}
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn wildcard_advanced_subscriber_receives_history_from_both_publishers() {
|
||||
let cfg =
|
||||
networking::cfg(&format!("{:x}", rand::random::<u128>()), 52412).expect("create config");
|
||||
let n_session = networking::open(cfg, "exo", 52412, 52411)
|
||||
.await
|
||||
.expect("open session");
|
||||
let session = n_session.z.clone();
|
||||
|
||||
// Unique prefix so the wildcard subscriber cannot accidentally see unrelated traffic.
|
||||
let run_id = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_nanos();
|
||||
|
||||
let base = format!("zenoh_advanced_history_test/{run_id}");
|
||||
let key_a = format!("{base}/a");
|
||||
let key_b = format!("{base}/b");
|
||||
let sub_key = format!("{base}/*");
|
||||
|
||||
let pub1 = session
|
||||
.declare_publisher(key_a.clone())
|
||||
.advanced()
|
||||
.publisher_detection()
|
||||
.sample_miss_detection(MissDetectionConfig::default())
|
||||
.cache(CacheConfig::default().max_samples(1))
|
||||
.await
|
||||
.expect("declare advanced publisher a");
|
||||
|
||||
pub1.put("aa").await.expect("publish aa");
|
||||
|
||||
let pub2 = session
|
||||
.declare_publisher(key_b.clone())
|
||||
.advanced()
|
||||
.sample_miss_detection(MissDetectionConfig::default())
|
||||
.publisher_detection()
|
||||
.cache(CacheConfig::default().max_samples(1))
|
||||
.await
|
||||
.expect("declare advanced publisher b");
|
||||
|
||||
pub2.put("bb").await.expect("publish bb");
|
||||
|
||||
// Give liveliness/cache declarations a brief chance to settle before declaring
|
||||
// the late-joining advanced subscriber.
|
||||
tokio::time::sleep(Duration::from_millis(250)).await;
|
||||
|
||||
let sub = session
|
||||
.declare_subscriber(sub_key)
|
||||
.advanced()
|
||||
.history(
|
||||
HistoryConfig::default()
|
||||
.max_samples(1)
|
||||
.detect_late_publishers(),
|
||||
)
|
||||
.await
|
||||
.expect("declare advanced subscriber");
|
||||
|
||||
expect_two_values(&sub, &*key_a, "aa", &*key_b, "bb").await
|
||||
}
|
||||
+35
-38
@@ -1,54 +1,51 @@
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use core::mem::drop;
|
||||
use core::option::Option::Some;
|
||||
use core::time::Duration;
|
||||
use tokio;
|
||||
use tokio::sync::mpsc;
|
||||
use core::mem::drop;
|
||||
use core::option::Option::Some;
|
||||
use core::time::Duration;
|
||||
use tokio;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_drop_channel() {
|
||||
struct Ping;
|
||||
#[tokio::test]
|
||||
async fn test_drop_channel() {
|
||||
struct Ping;
|
||||
|
||||
let (tx, mut rx) = mpsc::channel::<Ping>(10);
|
||||
let (tx, mut rx) = mpsc::channel::<Ping>(10);
|
||||
|
||||
let _ = tokio::spawn(async move {
|
||||
println!("TASK: entered");
|
||||
let _ = tokio::spawn(async move {
|
||||
println!("TASK: entered");
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
result = rx.recv() => {
|
||||
match result {
|
||||
Some(_) => {
|
||||
println!("TASK: pinged");
|
||||
}
|
||||
None => {
|
||||
println!("TASK: closing channel");
|
||||
break;
|
||||
}
|
||||
loop {
|
||||
tokio::select! {
|
||||
result = rx.recv() => {
|
||||
match result {
|
||||
Some(_) => {
|
||||
println!("TASK: pinged");
|
||||
}
|
||||
None => {
|
||||
println!("TASK: closing channel");
|
||||
break;
|
||||
}
|
||||
}
|
||||
_ = tokio::time::sleep(Duration::from_secs_f32(0.1)) => {
|
||||
println!("TASK: heartbeat");
|
||||
}
|
||||
}
|
||||
_ = tokio::time::sleep(Duration::from_secs_f32(0.1)) => {
|
||||
println!("TASK: heartbeat");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
println!("TASK: exited");
|
||||
});
|
||||
println!("TASK: exited");
|
||||
});
|
||||
|
||||
let tx2 = tx.clone();
|
||||
let tx2 = tx.clone();
|
||||
|
||||
tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
|
||||
tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
|
||||
|
||||
tx.send(Ping).await.expect("Should not fail");
|
||||
drop(tx);
|
||||
tx.send(Ping).await.expect("Should not fail");
|
||||
drop(tx);
|
||||
|
||||
tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
|
||||
tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
|
||||
|
||||
tx2.send(Ping).await.expect("Should not fail");
|
||||
drop(tx2);
|
||||
tx2.send(Ping).await.expect("Should not fail");
|
||||
drop(tx2);
|
||||
|
||||
tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
import os
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
from exo_rs import SessionHandle
|
||||
|
||||
|
||||
ZENOH_PORT = 52414
|
||||
DISCOVERY_PORT = 52413
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def storage():
|
||||
node_id = os.urandom(16).hex().rstrip("0")
|
||||
|
||||
session_handle, _nh = SessionHandle.new(
|
||||
node_id,
|
||||
ZENOH_PORT,
|
||||
DISCOVERY_PORT,
|
||||
)
|
||||
|
||||
return session_handle.storage_interface()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_storage_get_missing_key_returns_none(storage):
|
||||
key = f"tests/storage/{uuid.uuid4().hex}/missing"
|
||||
|
||||
value = await storage.get(key)
|
||||
assert value is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_storage_put_then_get_returns_value(storage):
|
||||
key = f"tests/storage/{uuid.uuid4().hex}/value"
|
||||
expected = "hello storage"
|
||||
|
||||
await storage.put(key, expected)
|
||||
assert await storage.get(key) == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_storage_put_overwrites_value(storage):
|
||||
key = f"tests/storage/{uuid.uuid4().hex}/overwrite"
|
||||
|
||||
await storage.put(key, "old")
|
||||
await storage.put(key, "new")
|
||||
assert await storage.get(key) == "new"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_storage_put_overwrites_value(storage):
|
||||
key = f"tests/storage/{uuid.uuid4().hex}/overwrite"
|
||||
|
||||
await storage.put(key, "old")
|
||||
await storage.delete(key)
|
||||
assert await storage.get(key) == None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_storage_get_rejects_wildcard_key(storage):
|
||||
with pytest.raises(ValueError, match="only supports fixed keys"):
|
||||
await storage.get("tests/storage/*")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_storage_put_rejects_wildcard_key(storage):
|
||||
with pytest.raises(ValueError, match="only supports fixed keys"):
|
||||
await storage.put("tests/storage/*", "value")
|
||||
@@ -0,0 +1,70 @@
|
||||
import asyncio
|
||||
import os
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
from exo_rs import SessionHandle
|
||||
|
||||
ZENOH_PORT = 52416
|
||||
DISCOVERY_PORT = 52413
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def session_handle():
|
||||
node_id = os.urandom(16).hex().lstrip("0")
|
||||
|
||||
session_handle, _nh = SessionHandle.new(
|
||||
node_id,
|
||||
ZENOH_PORT,
|
||||
DISCOVERY_PORT,
|
||||
)
|
||||
|
||||
return session_handle
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_task_requester_responder_round_trip(session_handle):
|
||||
instance_id = f"tests-task-instance-{uuid.uuid4().hex}"
|
||||
command_id = f"tests-task-command-{uuid.uuid4().hex}"
|
||||
command = '{"kind":"command"}'
|
||||
chunk = '{"kind":"chunk","finish_reason":"stop"}'
|
||||
|
||||
requester = session_handle.task_requester()
|
||||
responder = session_handle.task_responder(instance_id)
|
||||
|
||||
async def respond_to_submission():
|
||||
received = await responder.recv()
|
||||
assert received is not None
|
||||
request, chunk_sender, payload = received
|
||||
assert payload == command
|
||||
request.reply(command_id)
|
||||
await chunk_sender.send(chunk)
|
||||
|
||||
stream, _ = await asyncio.gather(
|
||||
requester.submit(instance_id, command_id, command),
|
||||
respond_to_submission(),
|
||||
)
|
||||
|
||||
assert await stream.recv() == chunk
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_task_requester_interrupt_round_trip(session_handle):
|
||||
instance_id = f"tests-task-instance-{uuid.uuid4().hex}"
|
||||
command_id = f"tests-task-command-{uuid.uuid4().hex}"
|
||||
command = '{"kind":"interrupt"}'
|
||||
|
||||
requester = session_handle.task_requester()
|
||||
responder = session_handle.task_responder(instance_id)
|
||||
|
||||
async def respond_to_interrupt():
|
||||
received = await responder.recv()
|
||||
assert received is not None
|
||||
request, _chunk_sender, payload = received
|
||||
assert payload == command
|
||||
request.reply(command_id)
|
||||
|
||||
await asyncio.gather(
|
||||
requester.interrupt(instance_id, command_id, command),
|
||||
respond_to_interrupt(),
|
||||
)
|
||||
+29
-14
@@ -7,9 +7,16 @@ use zenoh_plugin_trait::PluginsManager;
|
||||
|
||||
pub use zenoh::{Config, config::ZenohId};
|
||||
|
||||
use crate::discovery::Discovery;
|
||||
use crate::{
|
||||
discovery::Discovery,
|
||||
liveliness_aggregator::{LivelinessAggregator, spawn_liveliness_aggregator},
|
||||
};
|
||||
|
||||
pub use zenoh_plugin_storage_manager::read_raw_memory_storage;
|
||||
pub const STORAGE_PREFIX: &str = "storage/mem1";
|
||||
|
||||
pub mod discovery;
|
||||
pub mod liveliness_aggregator;
|
||||
pub mod swarm;
|
||||
|
||||
pub fn is_valid_zid(identity: &str) -> bool {
|
||||
@@ -33,20 +40,22 @@ pub fn cfg(identity: &str, listen_port: u16) -> Result<zenoh::Config> {
|
||||
cfg.insert_json5("scouting/multicast/autoconnect", "[]")?;
|
||||
cfg.insert_json5("scouting/gossip/multihop", "true")?;
|
||||
cfg.insert_json5("adminspace/enabled", "true")?;
|
||||
//cfg.insert_json5("transport/link/tx/batch_size", "9216")?;
|
||||
cfg.insert_json5("transport/link/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",
|
||||
r#"{
|
||||
key_expr: "storage/mem1/**",
|
||||
strip_prefix: "storage/mem1",
|
||||
volume: "memory",
|
||||
replication: {
|
||||
interval: 2,
|
||||
}
|
||||
}"#,
|
||||
"plugins/storage_manager/storages/mem1/key_expr",
|
||||
&format!("\"{STORAGE_PREFIX}/**\""),
|
||||
)?;
|
||||
cfg.insert_json5(
|
||||
"plugins/storage_manager/storages/mem1/strip_prefix",
|
||||
&format!("\"{STORAGE_PREFIX}\""),
|
||||
)?;
|
||||
cfg.insert_json5("plugins/storage_manager/storages/mem1/volume", "\"memory\"")?;
|
||||
cfg.insert_json5(
|
||||
"plugins/storage_manager/storages/mem1/replication/interval",
|
||||
"2",
|
||||
)?;
|
||||
Ok(cfg)
|
||||
}
|
||||
@@ -99,10 +108,15 @@ pub async fn open(
|
||||
.await;
|
||||
}
|
||||
})));
|
||||
Ok(Session { z, _jh })
|
||||
let liveliness_aggregator = spawn_liveliness_aggregator(&z)?;
|
||||
Ok(Session {
|
||||
z,
|
||||
liveliness_aggregator,
|
||||
_jh,
|
||||
})
|
||||
}
|
||||
|
||||
struct AbortOnDrop(JoinHandle<()>);
|
||||
pub struct AbortOnDrop(pub JoinHandle<()>);
|
||||
impl Drop for AbortOnDrop {
|
||||
fn drop(&mut self) {
|
||||
self.0.abort();
|
||||
@@ -112,5 +126,6 @@ impl Drop for AbortOnDrop {
|
||||
#[derive(Clone)]
|
||||
pub struct Session {
|
||||
pub z: ZSession,
|
||||
pub liveliness_aggregator: LivelinessAggregator,
|
||||
_jh: Arc<AbortOnDrop>,
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
use std::{collections::HashSet, sync::Arc};
|
||||
|
||||
use parking_lot::Mutex;
|
||||
use zenoh::{Result, Session, Wait, sample::SampleKind};
|
||||
|
||||
pub fn spawn_liveliness_aggregator(session: &Session) -> Result<LivelinessAggregator> {
|
||||
let store = Arc::new(Mutex::new(HashSet::default()));
|
||||
session
|
||||
.liveliness()
|
||||
.declare_subscriber("live/*")
|
||||
.history(true)
|
||||
.callback({
|
||||
let store = Arc::clone(&store);
|
||||
move |sample| {
|
||||
let Some(nid) = sample
|
||||
.key_expr()
|
||||
.to_string()
|
||||
.strip_prefix("live/")
|
||||
.map(str::to_owned)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let mut mg = store.lock();
|
||||
match sample.kind() {
|
||||
SampleKind::Put => mg.insert(nid),
|
||||
SampleKind::Delete => mg.remove(&nid),
|
||||
};
|
||||
}
|
||||
})
|
||||
.background()
|
||||
.wait()?;
|
||||
Ok(LivelinessAggregator { store })
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct LivelinessAggregator {
|
||||
// need two arcs as the sub owns an arc to the store.
|
||||
store: Arc<Mutex<HashSet<String>>>,
|
||||
}
|
||||
impl LivelinessAggregator {
|
||||
pub fn dump(&self) -> HashSet<String> {
|
||||
self.store.lock().clone()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use zenoh::Wait;
|
||||
|
||||
fn unique_key(name: &str) -> String {
|
||||
format!("test/zenoh-runtime-polling/{}/{}", std::process::id(), name)
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn get_and_recv_work_on_tokio_baseline() {
|
||||
let session = zenoh::open(zenoh::Config::default())
|
||||
.await
|
||||
.expect("open session");
|
||||
|
||||
let key = unique_key("tokio-baseline");
|
||||
let reply_key = key.clone();
|
||||
|
||||
let _queryable = session
|
||||
.declare_queryable(key.clone())
|
||||
.callback(move |query| {
|
||||
query
|
||||
.reply(reply_key.clone(), "hello-from-queryable")
|
||||
.wait()
|
||||
.expect("reply from queryable");
|
||||
})
|
||||
.await
|
||||
.expect("declare queryable");
|
||||
|
||||
let replies = session.get(key).await.expect("get");
|
||||
|
||||
let reply = tokio::time::timeout(Duration::from_secs(5), replies.recv_async())
|
||||
.await
|
||||
.expect("timed out waiting for reply")
|
||||
.expect("reply channel closed");
|
||||
|
||||
let sample = reply.result().expect("reply result was error");
|
||||
let payload = sample
|
||||
.payload()
|
||||
.try_to_string()
|
||||
.expect("payload should be utf8");
|
||||
|
||||
assert_eq!(payload.as_ref(), "hello-from-queryable");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn get_and_recv_work_when_polled_by_smol_without_tokio_context() {
|
||||
let session = zenoh::open(zenoh::Config::default())
|
||||
.await
|
||||
.expect("open session under tokio");
|
||||
|
||||
let key = unique_key("smol-no-tokio-context");
|
||||
let reply_key = key.clone();
|
||||
|
||||
let _queryable = session
|
||||
.declare_queryable(key.clone())
|
||||
.callback(move |query| {
|
||||
query
|
||||
.reply(reply_key.clone(), "hello-from-queryable")
|
||||
.wait()
|
||||
.expect("reply from queryable");
|
||||
})
|
||||
.await
|
||||
.expect("declare queryable under tokio");
|
||||
|
||||
let session_for_smol = session.clone();
|
||||
|
||||
let (tx, rx) = std::sync::mpsc::channel();
|
||||
|
||||
std::thread::spawn(move || {
|
||||
// This thread was not entered by Tokio.
|
||||
// If Zenoh's get/recv path requires an ambient Tokio Handle in the polling thread,
|
||||
// this is where it should panic, hang, or error.
|
||||
let result = {
|
||||
smol::block_on(async move {
|
||||
let replies = session_for_smol.get(key).await.expect("get under smol");
|
||||
|
||||
let reply = replies.recv_async().await.expect("reply channel closed");
|
||||
|
||||
let sample = reply.result().expect("reply result was error");
|
||||
let payload = sample
|
||||
.payload()
|
||||
.try_to_string()
|
||||
.expect("payload should be utf8");
|
||||
|
||||
payload.to_string()
|
||||
})
|
||||
};
|
||||
|
||||
tx.send(result).expect("send test result");
|
||||
});
|
||||
|
||||
let result = rx
|
||||
.recv_timeout(Duration::from_secs(5))
|
||||
.expect("smol thread timed out; likely hung waiting for get/reply");
|
||||
|
||||
let payload = result;
|
||||
|
||||
assert_eq!(payload, "hello-from-queryable");
|
||||
}
|
||||
@@ -0,0 +1,200 @@
|
||||
from pydantic import ValidationError
|
||||
|
||||
from exo.master.placement_utils import find_ip_prioritised
|
||||
from exo.shared.models.model_cards import ModelId
|
||||
from exo.shared.types.commands import ImageEdits, ImageGeneration, TextGeneration
|
||||
from exo.shared.types.instance_link import InstanceLink
|
||||
from exo.shared.types.state import State
|
||||
from exo.shared.types.tasks import (
|
||||
ImageEdits as ImageEditsTask,
|
||||
)
|
||||
from exo.shared.types.tasks import (
|
||||
ImageGeneration as ImageGenerationTask,
|
||||
)
|
||||
from exo.shared.types.tasks import (
|
||||
Task,
|
||||
TaskId,
|
||||
TaskStatus,
|
||||
)
|
||||
from exo.shared.types.tasks import (
|
||||
TextGeneration as TextGenerationTask,
|
||||
)
|
||||
from exo.shared.types.worker.instances import InstanceId
|
||||
|
||||
|
||||
def decode_instance_for_text_generation(
|
||||
state: State, model_id: ModelId, instance_links: list[InstanceLink]
|
||||
) -> InstanceId | None:
|
||||
prefill_only: set[InstanceId] = set()
|
||||
for link in instance_links:
|
||||
prefill_only.update(link.prefill_instances)
|
||||
for link in instance_links:
|
||||
prefill_only.difference_update(link.decode_instances)
|
||||
|
||||
instance_task_counts = _instance_task_counts_for_model(state, model_id)
|
||||
for instance_id in prefill_only:
|
||||
instance_task_counts.pop(instance_id, None)
|
||||
|
||||
if not instance_task_counts:
|
||||
return None
|
||||
|
||||
return min(
|
||||
instance_task_counts, key=lambda instance_id: instance_task_counts[instance_id]
|
||||
)
|
||||
|
||||
|
||||
def instance_for_generation(state: State, model_id: ModelId) -> InstanceId | None:
|
||||
instance_task_counts = _instance_task_counts_for_model(state, model_id)
|
||||
if not instance_task_counts:
|
||||
return None
|
||||
|
||||
return min(
|
||||
instance_task_counts, key=lambda instance_id: instance_task_counts[instance_id]
|
||||
)
|
||||
|
||||
|
||||
def text_generation_task(
|
||||
state: State, command: TextGeneration, instance_links: list[InstanceLink]
|
||||
) -> TextGenerationTask:
|
||||
instance_id = decode_instance_for_text_generation(
|
||||
state, command.task_params.model, instance_links
|
||||
)
|
||||
assert instance_id is not None
|
||||
task_params = command.task_params.model_copy(
|
||||
update={
|
||||
"prefill_endpoint": prefill_endpoint_for(
|
||||
state,
|
||||
instance_links,
|
||||
instance_id,
|
||||
),
|
||||
}
|
||||
)
|
||||
return TextGenerationTask(
|
||||
task_id=TaskId(),
|
||||
command_id=command.command_id,
|
||||
instance_id=instance_id,
|
||||
task_status=TaskStatus.Pending,
|
||||
task_params=task_params,
|
||||
)
|
||||
|
||||
|
||||
def image_generation_task(
|
||||
state: State, command: ImageGeneration
|
||||
) -> ImageGenerationTask:
|
||||
instance_id = instance_for_generation(state, ModelId(command.task_params.model))
|
||||
assert instance_id is not None
|
||||
return ImageGenerationTask(
|
||||
task_id=TaskId(),
|
||||
command_id=command.command_id,
|
||||
instance_id=instance_id,
|
||||
task_status=TaskStatus.Pending,
|
||||
task_params=command.task_params,
|
||||
)
|
||||
|
||||
|
||||
def image_edits_task(state: State, command: ImageEdits) -> ImageEditsTask:
|
||||
instance_id = instance_for_generation(state, ModelId(command.task_params.model))
|
||||
assert instance_id is not None
|
||||
return ImageEditsTask(
|
||||
task_id=TaskId(),
|
||||
command_id=command.command_id,
|
||||
instance_id=instance_id,
|
||||
task_status=TaskStatus.Pending,
|
||||
task_params=command.task_params,
|
||||
)
|
||||
|
||||
|
||||
def task_from_command(
|
||||
state: State,
|
||||
command: TextGeneration | ImageGeneration | ImageEdits,
|
||||
instance_links: list[InstanceLink],
|
||||
) -> Task:
|
||||
match command:
|
||||
case TextGeneration():
|
||||
return text_generation_task(state, command, instance_links)
|
||||
case ImageGeneration():
|
||||
return image_generation_task(state, command)
|
||||
case ImageEdits():
|
||||
return image_edits_task(state, command)
|
||||
|
||||
|
||||
def instance_id_for_command(
|
||||
state: State,
|
||||
command: TextGeneration | ImageGeneration | ImageEdits,
|
||||
instance_links: list[InstanceLink],
|
||||
) -> InstanceId | None:
|
||||
match command:
|
||||
case TextGeneration():
|
||||
return decode_instance_for_text_generation(
|
||||
state, command.task_params.model, instance_links
|
||||
)
|
||||
case ImageGeneration() | ImageEdits():
|
||||
return instance_for_generation(state, ModelId(command.task_params.model))
|
||||
|
||||
|
||||
def load_instance_links(values: list[str]) -> list[InstanceLink]:
|
||||
instance_links: list[InstanceLink] = []
|
||||
for value in values:
|
||||
try:
|
||||
instance_links.append(InstanceLink.model_validate_json(value))
|
||||
except ValidationError:
|
||||
continue
|
||||
return instance_links
|
||||
|
||||
|
||||
def prefill_endpoint_for(
|
||||
state: State, instance_links: list[InstanceLink], decode_instance_id: InstanceId
|
||||
) -> str | None:
|
||||
decode = state.instances.get(decode_instance_id)
|
||||
if decode is None:
|
||||
return None
|
||||
decode_node = decode.shard_assignments.shards[
|
||||
decode.shard_assignments.primary_output_node
|
||||
].node_id
|
||||
|
||||
sources: set[InstanceId] = set()
|
||||
for link in instance_links:
|
||||
if decode_instance_id in link.decode_instances:
|
||||
sources.update(link.prefill_instances)
|
||||
sources.discard(decode_instance_id)
|
||||
|
||||
in_flight = {TaskStatus.Pending, TaskStatus.Running}
|
||||
task_counts: dict[InstanceId, int] = {
|
||||
src_id: sum(
|
||||
1
|
||||
for task in state.tasks.values()
|
||||
if task.instance_id == src_id and task.task_status in in_flight
|
||||
)
|
||||
for src_id in sources
|
||||
}
|
||||
for src_id in sorted(sources, key=lambda sid: task_counts[sid]):
|
||||
instance = state.instances.get(src_id)
|
||||
if instance is None:
|
||||
continue
|
||||
for node_id, runner_id, _ in instance.shard_assignments.shards:
|
||||
port = state.prefill_server_ports.get(runner_id)
|
||||
if port is None:
|
||||
continue
|
||||
ip = find_ip_prioritised(
|
||||
decode_node, node_id, state.topology, state.node_network, ring=True
|
||||
)
|
||||
if ip is None:
|
||||
continue
|
||||
return f"{ip}:{port}"
|
||||
return None
|
||||
|
||||
|
||||
def _instance_task_counts_for_model(
|
||||
state: State, model_id: ModelId
|
||||
) -> dict[InstanceId, int]:
|
||||
in_flight = {TaskStatus.Pending, TaskStatus.Running}
|
||||
return {
|
||||
instance.instance_id: sum(
|
||||
1
|
||||
for task in state.tasks.values()
|
||||
if task.instance_id == instance.instance_id
|
||||
and task.task_status in in_flight
|
||||
)
|
||||
for instance in state.instances.values()
|
||||
if instance.shard_assignments.model_id == model_id
|
||||
}
|
||||
+419
-548
File diff suppressed because it is too large.
Load diff
@@ -7,6 +7,7 @@ from fastapi.testclient import TestClient
|
||||
|
||||
from exo.api.main import API
|
||||
from exo.shared.types.common import CommandId
|
||||
from exo.shared.types.worker.instances import InstanceId
|
||||
|
||||
|
||||
def _make_api() -> Any:
|
||||
@@ -15,9 +16,9 @@ def _make_api() -> Any:
|
||||
app = FastAPI()
|
||||
api = object.__new__(API)
|
||||
api.app = app
|
||||
api._text_generation_queues = {} # pyright: ignore[reportPrivateUsage]
|
||||
api._image_generation_queues = {} # pyright: ignore[reportPrivateUsage]
|
||||
api._send = AsyncMock() # pyright: ignore[reportPrivateUsage]
|
||||
api._bridge_command_instances = {} # pyright: ignore[reportPrivateUsage]
|
||||
api.task_requester = MagicMock()
|
||||
api.task_requester.interrupt = AsyncMock()
|
||||
api._setup_exception_handlers() # pyright: ignore[reportPrivateUsage]
|
||||
app.post("/v1/cancel/{command_id}")(api.cancel_command)
|
||||
return api
|
||||
@@ -38,40 +39,42 @@ def test_cancel_nonexistent_command_returns_404() -> None:
|
||||
|
||||
|
||||
def test_cancel_active_text_generation() -> None:
|
||||
"""Cancel an active text generation command: returns 200, sender.close() called."""
|
||||
"""Cancel an active text generation command: returns 200, interrupt sent."""
|
||||
api = _make_api()
|
||||
client = TestClient(api.app)
|
||||
|
||||
cid = CommandId("text-cmd-123")
|
||||
sender = MagicMock()
|
||||
api._text_generation_queues[cid] = sender
|
||||
instance_id = InstanceId("instance-a")
|
||||
api._bridge_command_instances[cid] = instance_id
|
||||
|
||||
response = client.post(f"/v1/cancel/{cid}")
|
||||
assert response.status_code == 200
|
||||
data: dict[str, Any] = response.json()
|
||||
assert data["message"] == "Command cancelled."
|
||||
assert data["command_id"] == str(cid)
|
||||
sender.close.assert_called_once()
|
||||
api._send.assert_called_once()
|
||||
task_cancelled = api._send.call_args[0][0]
|
||||
assert task_cancelled.cancelled_command_id == cid
|
||||
api.task_requester.interrupt.assert_called_once()
|
||||
args = api.task_requester.interrupt.call_args.args
|
||||
assert args[0] == instance_id
|
||||
assert args[1] == cid
|
||||
assert '"cancelled_command_id":"text-cmd-123"' in args[2]
|
||||
|
||||
|
||||
def test_cancel_active_image_generation() -> None:
|
||||
"""Cancel an active image generation command: returns 200, sender.close() called."""
|
||||
"""Cancel an active image generation command: returns 200, interrupt sent."""
|
||||
api = _make_api()
|
||||
client = TestClient(api.app)
|
||||
|
||||
cid = CommandId("img-cmd-456")
|
||||
sender = MagicMock()
|
||||
api._image_generation_queues[cid] = sender
|
||||
instance_id = InstanceId("instance-b")
|
||||
api._bridge_command_instances[cid] = instance_id
|
||||
|
||||
response = client.post(f"/v1/cancel/{cid}")
|
||||
assert response.status_code == 200
|
||||
data: dict[str, Any] = response.json()
|
||||
assert data["message"] == "Command cancelled."
|
||||
assert data["command_id"] == str(cid)
|
||||
sender.close.assert_called_once()
|
||||
api._send.assert_called_once()
|
||||
task_cancelled = api._send.call_args[0][0]
|
||||
assert task_cancelled.cancelled_command_id == cid
|
||||
api.task_requester.interrupt.assert_called_once()
|
||||
args = api.task_requester.interrupt.call_args.args
|
||||
assert args[0] == instance_id
|
||||
assert args[1] == cid
|
||||
assert '"cancelled_command_id":"img-cmd-456"' in args[2]
|
||||
@@ -1,97 +0,0 @@
|
||||
# pyright: reportUnusedFunction=false, reportAny=false
|
||||
"""Tests that InstanceDeleted events close active generation streams."""
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from exo.api.main import API
|
||||
from exo.api.types import ImageGenerationTaskParams
|
||||
from exo.shared.types.common import CommandId, ModelId
|
||||
from exo.shared.types.state import State
|
||||
from exo.shared.types.tasks import ImageGeneration, TextGeneration
|
||||
from exo.shared.types.text_generation import (
|
||||
InputMessage,
|
||||
InputMessageContent,
|
||||
TextGenerationTaskParams,
|
||||
)
|
||||
from exo.shared.types.worker.instances import InstanceId
|
||||
|
||||
|
||||
def _make_api_with_state(state: State) -> API:
|
||||
"""Create a minimal API instance with pre-set state."""
|
||||
api = object.__new__(API)
|
||||
api.state = state
|
||||
api._text_generation_queues = {} # pyright: ignore[reportPrivateUsage]
|
||||
api._image_generation_queues = {} # pyright: ignore[reportPrivateUsage]
|
||||
return api
|
||||
|
||||
|
||||
def _make_text_gen_task(
|
||||
instance_id: InstanceId, command_id: CommandId
|
||||
) -> TextGeneration:
|
||||
return TextGeneration(
|
||||
instance_id=instance_id,
|
||||
command_id=command_id,
|
||||
task_params=TextGenerationTaskParams(
|
||||
model=ModelId("test-model"),
|
||||
input=[InputMessage(role="user", content=InputMessageContent("hello"))],
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_close_streams_for_deleted_instance() -> None:
|
||||
"""Deleting an instance closes the text generation sender for commands on that instance."""
|
||||
instance_id = InstanceId("inst-1")
|
||||
command_id = CommandId("cmd-1")
|
||||
task = _make_text_gen_task(instance_id, command_id)
|
||||
|
||||
state = State(tasks={task.task_id: task})
|
||||
api = _make_api_with_state(state)
|
||||
|
||||
sender = MagicMock()
|
||||
api._text_generation_queues[command_id] = sender # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
api._close_streams_for_instance(instance_id) # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
sender.close.assert_called_once()
|
||||
assert command_id not in api._text_generation_queues # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
def test_close_streams_ignores_unrelated_instances() -> None:
|
||||
"""Deleting an instance does NOT close streams for commands on other instances."""
|
||||
target_id = InstanceId("inst-delete")
|
||||
other_id = InstanceId("inst-keep")
|
||||
other_cmd = CommandId("cmd-keep")
|
||||
other_task = _make_text_gen_task(other_id, other_cmd)
|
||||
|
||||
state = State(tasks={other_task.task_id: other_task})
|
||||
api = _make_api_with_state(state)
|
||||
|
||||
sender = MagicMock()
|
||||
api._text_generation_queues[other_cmd] = sender # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
api._close_streams_for_instance(target_id) # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
sender.close.assert_not_called()
|
||||
assert other_cmd in api._text_generation_queues # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
def test_close_streams_for_deleted_instance_image_generation() -> None:
|
||||
"""Deleting an instance closes the image generation sender for commands on that instance."""
|
||||
instance_id = InstanceId("inst-img")
|
||||
command_id = CommandId("cmd-img")
|
||||
task = ImageGeneration(
|
||||
instance_id=instance_id,
|
||||
command_id=command_id,
|
||||
task_params=ImageGenerationTaskParams(prompt="a cat", model="test-model"),
|
||||
)
|
||||
|
||||
state = State(tasks={task.task_id: task})
|
||||
api = _make_api_with_state(state)
|
||||
|
||||
sender = MagicMock()
|
||||
api._image_generation_queues[command_id] = sender # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
api._close_streams_for_instance(instance_id) # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
sender.close.assert_called_once()
|
||||
assert command_id not in api._image_generation_queues # pyright: ignore[reportPrivateUsage]
|
||||
@@ -297,13 +297,11 @@ class DeleteInstanceTaskParams(BaseModel):
|
||||
|
||||
class CreateInstanceResponse(BaseModel):
|
||||
message: str
|
||||
command_id: CommandId
|
||||
model_card: ModelCard
|
||||
|
||||
|
||||
class DeleteInstanceResponse(BaseModel):
|
||||
message: str
|
||||
command_id: CommandId
|
||||
instance_id: InstanceId
|
||||
|
||||
|
||||
@@ -329,7 +327,6 @@ class InstanceLinkBody(BaseModel):
|
||||
|
||||
class InstanceLinkResponse(BaseModel):
|
||||
message: str
|
||||
command_id: CommandId
|
||||
|
||||
|
||||
ImageSize = Literal[
|
||||
@@ -394,7 +391,6 @@ class ImageEditsTaskParams(BaseModel):
|
||||
"""Internal task params for image-editing requests."""
|
||||
|
||||
image_data: str = "" # Base64-encoded image (empty when using chunked transfer)
|
||||
total_input_chunks: int = 0
|
||||
prompt: str
|
||||
model: str
|
||||
n: int | None = 1
|
||||
|
||||
@@ -5,6 +5,7 @@ from pathlib import Path
|
||||
|
||||
import anyio
|
||||
from anyio import BrokenResourceError, ClosedResourceError, current_time, to_thread
|
||||
from exo_rs import LVPublisher, SessionHandle
|
||||
from loguru import logger
|
||||
|
||||
from exo.download.download_utils import (
|
||||
@@ -31,7 +32,6 @@ from exo.shared.types.commands import (
|
||||
from exo.shared.types.common import NodeId
|
||||
from exo.shared.types.events import (
|
||||
Event,
|
||||
NodeDownloadProgress,
|
||||
)
|
||||
from exo.shared.types.memory import Memory
|
||||
from exo.shared.types.worker.downloads import (
|
||||
@@ -52,6 +52,7 @@ class DownloadCoordinator:
|
||||
shard_downloader: ShardDownloader
|
||||
download_command_receiver: Receiver[ForwarderDownloadCommand]
|
||||
event_sender: Sender[Event]
|
||||
session_handle: SessionHandle
|
||||
offline: bool = False
|
||||
|
||||
# Local state
|
||||
@@ -63,6 +64,7 @@ class DownloadCoordinator:
|
||||
|
||||
# Per-model throttle for download progress events
|
||||
_last_progress_time: dict[ModelId, float] = field(default_factory=dict)
|
||||
_download_publishers: dict[ModelId, LVPublisher] = field(default_factory=dict)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.shard_downloader.on_progress(self._download_progress_callback)
|
||||
@@ -85,6 +87,26 @@ class DownloadCoordinator:
|
||||
read_only=is_read_only_model_dir(found),
|
||||
)
|
||||
|
||||
def _publisher_for_model(self, model_id: ModelId) -> LVPublisher:
|
||||
if (publisher := self._download_publishers.get(model_id)) is None:
|
||||
publisher = self.session_handle.last_value_publisher(
|
||||
f"node_metrics/{self.node_id}/downloads/{model_id}"
|
||||
)
|
||||
self._download_publishers[model_id] = publisher
|
||||
|
||||
return publisher
|
||||
|
||||
async def _set_download_status(self, status: DownloadProgress) -> None:
|
||||
model_id = status.shard_metadata.model_card.model_id
|
||||
self.download_status[model_id] = status
|
||||
publisher = self._publisher_for_model(model_id)
|
||||
await publisher.put(status.model_dump_json())
|
||||
|
||||
async def _delete_download_status(self, model_id: ModelId) -> None:
|
||||
self.download_status.pop(model_id, None)
|
||||
publisher = self._publisher_for_model(model_id)
|
||||
await publisher.delete()
|
||||
|
||||
async def _download_progress_callback(
|
||||
self, callback_shard: ShardMetadata, progress: RepoDownloadProgress
|
||||
) -> None:
|
||||
@@ -107,10 +129,7 @@ class DownloadCoordinator:
|
||||
total=progress.total,
|
||||
model_directory=self._default_model_dir(model_id),
|
||||
)
|
||||
self.download_status[model_id] = completed
|
||||
await self.event_sender.send(
|
||||
NodeDownloadProgress(download_progress=completed)
|
||||
)
|
||||
await self._set_download_status(completed)
|
||||
self._last_progress_time.pop(model_id, None)
|
||||
elif (
|
||||
progress.status == "in_progress"
|
||||
@@ -125,10 +144,7 @@ class DownloadCoordinator:
|
||||
),
|
||||
model_directory=self._default_model_dir(model_id),
|
||||
)
|
||||
self.download_status[model_id] = ongoing
|
||||
await self.event_sender.send(
|
||||
NodeDownloadProgress(download_progress=ongoing)
|
||||
)
|
||||
await self._set_download_status(ongoing)
|
||||
self._last_progress_time[model_id] = current_time()
|
||||
except (BrokenResourceError, ClosedResourceError):
|
||||
logger.debug(
|
||||
@@ -189,10 +205,7 @@ class DownloadCoordinator:
|
||||
downloaded=downloaded,
|
||||
total=total,
|
||||
)
|
||||
self.download_status[model_id] = pending
|
||||
await self.event_sender.send(
|
||||
NodeDownloadProgress(download_progress=pending)
|
||||
)
|
||||
await self._set_download_status(pending)
|
||||
|
||||
async def _start_download(self, shard: ShardMetadata) -> None:
|
||||
model_id = shard.model_card.model_id
|
||||
@@ -215,10 +228,7 @@ class DownloadCoordinator:
|
||||
completed = self._completed_from_path(
|
||||
shard, found_path, shard.model_card.storage_size
|
||||
)
|
||||
self.download_status[model_id] = completed
|
||||
await self.event_sender.send(
|
||||
NodeDownloadProgress(download_progress=completed)
|
||||
)
|
||||
await self._set_download_status(completed)
|
||||
return
|
||||
|
||||
# Emit pending status
|
||||
@@ -227,8 +237,7 @@ class DownloadCoordinator:
|
||||
node_id=self.node_id,
|
||||
model_directory=self._default_model_dir(model_id),
|
||||
)
|
||||
self.download_status[model_id] = progress
|
||||
await self.event_sender.send(NodeDownloadProgress(download_progress=progress))
|
||||
await self._set_download_status(progress)
|
||||
|
||||
# Check initial status from downloader
|
||||
initial_progress = (
|
||||
@@ -250,10 +259,7 @@ class DownloadCoordinator:
|
||||
total=initial_progress.total,
|
||||
model_directory=self._default_model_dir(model_id),
|
||||
)
|
||||
self.download_status[model_id] = completed
|
||||
await self.event_sender.send(
|
||||
NodeDownloadProgress(download_progress=completed)
|
||||
)
|
||||
await self._set_download_status(completed)
|
||||
return
|
||||
|
||||
if self.offline:
|
||||
@@ -266,18 +272,9 @@ class DownloadCoordinator:
|
||||
error_message=f"Model files not found locally in offline mode: {model_id}",
|
||||
model_directory=self._default_model_dir(model_id),
|
||||
)
|
||||
self.download_status[model_id] = failed
|
||||
await self.event_sender.send(NodeDownloadProgress(download_progress=failed))
|
||||
await self._set_download_status(failed)
|
||||
return
|
||||
|
||||
# Start actual download
|
||||
self._start_download_task(shard, initial_progress)
|
||||
|
||||
def _start_download_task(
|
||||
self, shard: ShardMetadata, initial_progress: RepoDownloadProgress
|
||||
) -> None:
|
||||
model_id = shard.model_card.model_id
|
||||
|
||||
# Emit ongoing status
|
||||
status = DownloadOngoing(
|
||||
node_id=self.node_id,
|
||||
@@ -287,8 +284,7 @@ class DownloadCoordinator:
|
||||
),
|
||||
model_directory=self._default_model_dir(model_id),
|
||||
)
|
||||
self.download_status[model_id] = status
|
||||
self.event_sender.send_nowait(NodeDownloadProgress(download_progress=status))
|
||||
await self._set_download_status(status)
|
||||
|
||||
async def download_wrapper(cancel_scope: anyio.CancelScope) -> None:
|
||||
try:
|
||||
@@ -302,10 +298,7 @@ class DownloadCoordinator:
|
||||
error_message=str(e),
|
||||
model_directory=self._default_model_dir(model_id),
|
||||
)
|
||||
self.download_status[model_id] = failed
|
||||
await self.event_sender.send(
|
||||
NodeDownloadProgress(download_progress=failed)
|
||||
)
|
||||
await self._set_download_status(failed)
|
||||
except anyio.get_cancelled_exc_class():
|
||||
# ignore cancellation - let cleanup do its thing
|
||||
pass
|
||||
@@ -338,18 +331,9 @@ class DownloadCoordinator:
|
||||
else:
|
||||
logger.warning(f"Model {model_id} was not found on disk")
|
||||
|
||||
# Emit pending status to reset UI state, then remove from local tracking
|
||||
# Remove local and aggregated download status.
|
||||
if model_id in self.download_status:
|
||||
current_status = self.download_status[model_id]
|
||||
pending = DownloadPending(
|
||||
shard_metadata=current_status.shard_metadata,
|
||||
node_id=self.node_id,
|
||||
model_directory=self._default_model_dir(model_id),
|
||||
)
|
||||
await self.event_sender.send(
|
||||
NodeDownloadProgress(download_progress=pending)
|
||||
)
|
||||
del self.download_status[model_id]
|
||||
await self._delete_download_status(model_id)
|
||||
|
||||
async def _emit_existing_download_progress(self) -> None:
|
||||
while True:
|
||||
@@ -428,10 +412,7 @@ class DownloadCoordinator:
|
||||
else:
|
||||
continue
|
||||
|
||||
self.download_status[progress.shard.model_card.model_id] = status
|
||||
await self.event_sender.send(
|
||||
NodeDownloadProgress(download_progress=status)
|
||||
)
|
||||
await self._set_download_status(status)
|
||||
# Scan read-only directories for pre-downloaded models
|
||||
if EXO_MODELS_READ_ONLY_DIRS:
|
||||
for card in await model_cards.card_cache.list_all():
|
||||
@@ -460,10 +441,7 @@ class DownloadCoordinator:
|
||||
path_shard, found, card.storage_size
|
||||
)
|
||||
)
|
||||
self.download_status[mid] = path_completed
|
||||
await self.event_sender.send(
|
||||
NodeDownloadProgress(download_progress=path_completed)
|
||||
)
|
||||
await self._set_download_status(path_completed)
|
||||
|
||||
logger.debug(
|
||||
"DownloadCoordinator: Done emitting existing download progress."
|
||||
|
||||
@@ -19,11 +19,11 @@ from exo.shared.types.commands import (
|
||||
StartDownload,
|
||||
)
|
||||
from exo.shared.types.common import NodeId, SystemId
|
||||
from exo.shared.types.events import Event, NodeDownloadProgress
|
||||
from exo.shared.types.events import Event
|
||||
from exo.shared.types.memory import Memory
|
||||
from exo.shared.types.worker.downloads import DownloadPending
|
||||
from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
|
||||
from exo.utils.channels import Receiver, Sender, channel
|
||||
from exo.utils.channels import Sender, channel
|
||||
|
||||
NODE_ID = NodeId("aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa")
|
||||
MODEL_ID = ModelId("test-org/test-model")
|
||||
@@ -137,10 +137,9 @@ def _setup_coordinator(
|
||||
) -> tuple[
|
||||
DownloadCoordinator,
|
||||
Sender[ForwarderDownloadCommand],
|
||||
Receiver[Event],
|
||||
]:
|
||||
cmd_send, cmd_recv = channel[ForwarderDownloadCommand]()
|
||||
event_send, event_recv = channel[Event]()
|
||||
event_send, _event_recv = channel[Event]()
|
||||
wrapped = SingletonShardDownloader(downloader)
|
||||
coordinator = DownloadCoordinator(
|
||||
node_id=NODE_ID,
|
||||
@@ -148,24 +147,20 @@ def _setup_coordinator(
|
||||
download_command_receiver=cmd_recv,
|
||||
event_sender=event_send,
|
||||
)
|
||||
return coordinator, cmd_send, event_recv
|
||||
return coordinator, cmd_send
|
||||
|
||||
|
||||
async def _wait_for_pending(
|
||||
event_recv: Receiver[Event], model_id: ModelId, timeout: float = 2.0
|
||||
coordinator: DownloadCoordinator, model_id: ModelId, timeout: float = 2.0
|
||||
) -> DownloadPending | None:
|
||||
"""Drain events until we see a DownloadPending for the given model, or timeout."""
|
||||
"""Wait until coordinator state marks the given model as pending."""
|
||||
try:
|
||||
async with asyncio.timeout(timeout):
|
||||
while True:
|
||||
event = await event_recv.receive()
|
||||
if (
|
||||
isinstance(event, NodeDownloadProgress)
|
||||
and isinstance(event.download_progress, DownloadPending)
|
||||
and event.download_progress.shard_metadata.model_card.model_id
|
||||
== model_id
|
||||
):
|
||||
return event.download_progress
|
||||
progress = coordinator.download_status.get(model_id)
|
||||
if isinstance(progress, DownloadPending):
|
||||
return progress
|
||||
await asyncio.sleep(0.01)
|
||||
except TimeoutError:
|
||||
return None
|
||||
|
||||
@@ -174,7 +169,7 @@ async def test_cancel_active_download_transitions_to_pending() -> None:
|
||||
"""Cancelling an in-progress download should emit a DownloadPending event
|
||||
and remove the model from active_downloads."""
|
||||
slow_downloader = SlowShardDownloader()
|
||||
coordinator, cmd_send, event_recv = _setup_coordinator(slow_downloader)
|
||||
coordinator, cmd_send = _setup_coordinator(slow_downloader)
|
||||
shard = _make_shard()
|
||||
origin = SystemId("test")
|
||||
|
||||
@@ -191,14 +186,6 @@ async def test_cancel_active_download_transitions_to_pending() -> None:
|
||||
# Wait for the download to actually start (blocking in ensure_shard)
|
||||
await asyncio.wait_for(slow_downloader.download_started.wait(), timeout=2.0)
|
||||
|
||||
# Drain any events emitted before the cancel (initial DownloadPending, DownloadOngoing)
|
||||
while True:
|
||||
try:
|
||||
async with asyncio.timeout(0.1):
|
||||
await event_recv.receive()
|
||||
except TimeoutError:
|
||||
break
|
||||
|
||||
# Cancel the download
|
||||
await cmd_send.send(
|
||||
ForwarderDownloadCommand(
|
||||
@@ -207,9 +194,8 @@ async def test_cancel_active_download_transitions_to_pending() -> None:
|
||||
)
|
||||
)
|
||||
|
||||
# Should receive a DownloadPending event with preserved progress
|
||||
pending = await _wait_for_pending(event_recv, MODEL_ID)
|
||||
assert pending is not None, "Cancel should emit DownloadPending"
|
||||
pending = await _wait_for_pending(coordinator, MODEL_ID)
|
||||
assert pending is not None, "Cancel should update state to DownloadPending"
|
||||
assert pending.shard_metadata.model_card.model_id == MODEL_ID
|
||||
assert pending.total == Memory.from_mb(100), "Should preserve total bytes"
|
||||
|
||||
@@ -231,7 +217,7 @@ async def test_cancel_active_download_transitions_to_pending() -> None:
|
||||
async def test_cancel_nonexistent_download_is_noop() -> None:
|
||||
"""Cancelling a model that isn't being downloaded should be a no-op."""
|
||||
slow_downloader = SlowShardDownloader()
|
||||
coordinator, cmd_send, event_recv = _setup_coordinator(slow_downloader)
|
||||
coordinator, cmd_send = _setup_coordinator(slow_downloader)
|
||||
origin = SystemId("test")
|
||||
|
||||
coordinator_task = asyncio.create_task(coordinator.run())
|
||||
@@ -244,9 +230,10 @@ async def test_cancel_nonexistent_download_is_noop() -> None:
|
||||
)
|
||||
)
|
||||
|
||||
# Should NOT receive any DownloadPending event
|
||||
pending = await _wait_for_pending(event_recv, MODEL_ID, timeout=0.5)
|
||||
assert pending is None, "Cancel of non-existent download should not emit events"
|
||||
pending = await _wait_for_pending(coordinator, MODEL_ID, timeout=0.5)
|
||||
assert pending is None, (
|
||||
"Cancel of non-existent download should not update state"
|
||||
)
|
||||
|
||||
# Coordinator state should be empty
|
||||
assert MODEL_ID not in coordinator.active_downloads
|
||||
@@ -261,7 +248,7 @@ async def test_cancel_nonexistent_download_is_noop() -> None:
|
||||
async def test_cancel_then_resume_download() -> None:
|
||||
"""After cancelling, re-issuing StartDownload should restart the download."""
|
||||
slow_downloader = SlowShardDownloader()
|
||||
coordinator, cmd_send, event_recv = _setup_coordinator(slow_downloader)
|
||||
coordinator, cmd_send = _setup_coordinator(slow_downloader)
|
||||
shard = _make_shard()
|
||||
origin = SystemId("test")
|
||||
|
||||
@@ -283,8 +270,8 @@ async def test_cancel_then_resume_download() -> None:
|
||||
command=CancelDownload(target_node_id=NODE_ID, model_id=MODEL_ID),
|
||||
)
|
||||
)
|
||||
pending = await _wait_for_pending(event_recv, MODEL_ID)
|
||||
assert pending is not None, "Cancel should emit DownloadPending"
|
||||
pending = await _wait_for_pending(coordinator, MODEL_ID)
|
||||
assert pending is not None, "Cancel should update state to DownloadPending"
|
||||
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
|
||||
@@ -22,14 +22,14 @@ from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
|
||||
from exo.shared.types.backends import Backend
|
||||
from exo.shared.types.commands import ForwarderDownloadCommand
|
||||
from exo.shared.types.common import NodeId
|
||||
from exo.shared.types.events import Event, NodeDownloadProgress
|
||||
from exo.shared.types.events import Event
|
||||
from exo.shared.types.memory import Memory
|
||||
from exo.shared.types.worker.downloads import (
|
||||
DownloadCompleted,
|
||||
DownloadPending,
|
||||
)
|
||||
from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
|
||||
from exo.utils.channels import Receiver, Sender, channel
|
||||
from exo.utils.channels import Sender, channel
|
||||
|
||||
NODE_ID = NodeId("aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa")
|
||||
MODEL_ID = ModelId("test-org/test-model")
|
||||
@@ -126,10 +126,9 @@ def _setup_coordinator(
|
||||
) -> tuple[
|
||||
DownloadCoordinator,
|
||||
Sender[ForwarderDownloadCommand],
|
||||
Receiver[Event],
|
||||
]:
|
||||
cmd_send, cmd_recv = channel[ForwarderDownloadCommand]()
|
||||
event_send, event_recv = channel[Event]()
|
||||
event_send, _event_recv = channel[Event]()
|
||||
wrapped = SingletonShardDownloader(downloader)
|
||||
coordinator = DownloadCoordinator(
|
||||
node_id=NODE_ID,
|
||||
@@ -137,21 +136,25 @@ def _setup_coordinator(
|
||||
download_command_receiver=cmd_recv,
|
||||
event_sender=event_send,
|
||||
)
|
||||
return coordinator, cmd_send, event_recv
|
||||
return coordinator, cmd_send
|
||||
|
||||
|
||||
async def _collect_events(
|
||||
event_recv: Receiver[Event], timeout: float = 1.0
|
||||
) -> list[Event]:
|
||||
"""Drain events until timeout."""
|
||||
events: list[Event] = []
|
||||
async def _wait_for_status(
|
||||
coordinator: DownloadCoordinator,
|
||||
model_id: ModelId,
|
||||
expected_type: type[DownloadCompleted] | type[DownloadPending],
|
||||
timeout: float = 1.0,
|
||||
) -> DownloadCompleted | DownloadPending | None:
|
||||
"""Wait until coordinator state has the expected status type."""
|
||||
try:
|
||||
async with asyncio.timeout(timeout):
|
||||
while True:
|
||||
events.append(await event_recv.receive())
|
||||
status = coordinator.download_status.get(model_id)
|
||||
if isinstance(status, expected_type):
|
||||
return status
|
||||
await asyncio.sleep(0.01)
|
||||
except TimeoutError:
|
||||
pass
|
||||
return events
|
||||
return None
|
||||
|
||||
|
||||
async def test_completed_status_not_downgraded_by_rescan() -> None:
|
||||
@@ -159,7 +162,7 @@ async def test_completed_status_not_downgraded_by_rescan() -> None:
|
||||
DownloadPending when the periodic rescan reports a non-complete
|
||||
file-size status (regression test for #1918)."""
|
||||
downloader = FakeShardDownloader(status="not_started")
|
||||
coordinator, _cmd_send, event_recv = _setup_coordinator(downloader)
|
||||
coordinator, _cmd_send = _setup_coordinator(downloader)
|
||||
|
||||
# Pre-seed the coordinator with a completed status for the model
|
||||
completed = DownloadCompleted(
|
||||
@@ -173,25 +176,14 @@ async def test_completed_status_not_downgraded_by_rescan() -> None:
|
||||
# Run the coordinator (the rescan loop fires immediately)
|
||||
coordinator_task = asyncio.create_task(coordinator.run())
|
||||
try:
|
||||
# Wait for the rescan to process (it should skip the completed model)
|
||||
events = await _collect_events(event_recv, timeout=1.5)
|
||||
await asyncio.sleep(1.5)
|
||||
|
||||
# The model must still be DownloadCompleted — not downgraded
|
||||
assert isinstance(coordinator.download_status[MODEL_ID], DownloadCompleted), (
|
||||
f"Expected DownloadCompleted but got {type(coordinator.download_status[MODEL_ID]).__name__}"
|
||||
)
|
||||
|
||||
# No DownloadPending event should have been emitted for this model
|
||||
pending_events = [
|
||||
e
|
||||
for e in events
|
||||
if isinstance(e, NodeDownloadProgress)
|
||||
and isinstance(e.download_progress, DownloadPending)
|
||||
and e.download_progress.shard_metadata.model_card.model_id == MODEL_ID
|
||||
]
|
||||
assert len(pending_events) == 0, (
|
||||
f"Expected no DownloadPending events for completed model, got {len(pending_events)}"
|
||||
)
|
||||
assert not isinstance(coordinator.download_status[MODEL_ID], DownloadPending)
|
||||
finally:
|
||||
await coordinator.shutdown()
|
||||
coordinator_task.cancel()
|
||||
@@ -204,7 +196,7 @@ async def test_incomplete_model_with_files_present_detected_as_complete() -> Non
|
||||
confirms the model directory is complete, the model should be marked
|
||||
DownloadCompleted (regression test for #1918 — initial scan case)."""
|
||||
downloader = FakeShardDownloader(status="not_started")
|
||||
coordinator, _cmd_send, event_recv = _setup_coordinator(downloader)
|
||||
coordinator, _cmd_send = _setup_coordinator(downloader)
|
||||
|
||||
# Mock resolve_existing_model to return a valid path (model is on disk)
|
||||
with patch(
|
||||
@@ -213,26 +205,14 @@ async def test_incomplete_model_with_files_present_detected_as_complete() -> Non
|
||||
):
|
||||
coordinator_task = asyncio.create_task(coordinator.run())
|
||||
try:
|
||||
events = await _collect_events(event_recv, timeout=1.5)
|
||||
|
||||
# The model should be DownloadCompleted (resolve_existing_model confirmed it)
|
||||
assert isinstance(
|
||||
coordinator.download_status.get(MODEL_ID), DownloadCompleted
|
||||
), (
|
||||
f"Expected DownloadCompleted but got "
|
||||
f"{type(coordinator.download_status.get(MODEL_ID)).__name__}"
|
||||
status = await _wait_for_status(
|
||||
coordinator, MODEL_ID, DownloadCompleted, timeout=1.5
|
||||
)
|
||||
|
||||
# Should have emitted a DownloadCompleted event
|
||||
completed_events = [
|
||||
e
|
||||
for e in events
|
||||
if isinstance(e, NodeDownloadProgress)
|
||||
and isinstance(e.download_progress, DownloadCompleted)
|
||||
and e.download_progress.shard_metadata.model_card.model_id == MODEL_ID
|
||||
]
|
||||
assert len(completed_events) > 0, (
|
||||
"Expected at least one DownloadCompleted event"
|
||||
# The model should be DownloadCompleted (resolve_existing_model confirmed it)
|
||||
assert isinstance(status, DownloadCompleted), (
|
||||
f"Expected DownloadCompleted but got "
|
||||
f"{type(coordinator.download_status.get(MODEL_ID)).__name__}"
|
||||
)
|
||||
finally:
|
||||
await coordinator.shutdown()
|
||||
@@ -246,7 +226,7 @@ async def test_genuinely_incomplete_model_stays_pending() -> None:
|
||||
returns None (model truly not complete), the model should correctly be
|
||||
DownloadPending."""
|
||||
downloader = FakeShardDownloader(status="not_started")
|
||||
coordinator, _cmd_send, event_recv = _setup_coordinator(downloader)
|
||||
coordinator, _cmd_send = _setup_coordinator(downloader)
|
||||
|
||||
# Mock resolve_existing_model to return None (model not on disk)
|
||||
with patch(
|
||||
@@ -255,26 +235,14 @@ async def test_genuinely_incomplete_model_stays_pending() -> None:
|
||||
):
|
||||
coordinator_task = asyncio.create_task(coordinator.run())
|
||||
try:
|
||||
events = await _collect_events(event_recv, timeout=1.5)
|
||||
|
||||
# The model should be DownloadPending
|
||||
assert isinstance(
|
||||
coordinator.download_status.get(MODEL_ID), DownloadPending
|
||||
), (
|
||||
f"Expected DownloadPending but got "
|
||||
f"{type(coordinator.download_status.get(MODEL_ID)).__name__}"
|
||||
status = await _wait_for_status(
|
||||
coordinator, MODEL_ID, DownloadPending, timeout=1.5
|
||||
)
|
||||
|
||||
# Should have emitted a DownloadPending event
|
||||
pending_events = [
|
||||
e
|
||||
for e in events
|
||||
if isinstance(e, NodeDownloadProgress)
|
||||
and isinstance(e.download_progress, DownloadPending)
|
||||
and e.download_progress.shard_metadata.model_card.model_id == MODEL_ID
|
||||
]
|
||||
assert len(pending_events) > 0, (
|
||||
"Expected at least one DownloadPending event"
|
||||
# The model should be DownloadPending
|
||||
assert isinstance(status, DownloadPending), (
|
||||
f"Expected DownloadPending but got "
|
||||
f"{type(coordinator.download_status.get(MODEL_ID)).__name__}"
|
||||
)
|
||||
finally:
|
||||
await coordinator.shutdown()
|
||||
|
||||
@@ -20,11 +20,11 @@ from exo.shared.types.commands import (
|
||||
StartDownload,
|
||||
)
|
||||
from exo.shared.types.common import NodeId, SystemId
|
||||
from exo.shared.types.events import Event, NodeDownloadProgress
|
||||
from exo.shared.types.events import Event
|
||||
from exo.shared.types.memory import Memory
|
||||
from exo.shared.types.worker.downloads import DownloadCompleted
|
||||
from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
|
||||
from exo.utils.channels import Receiver, Sender, channel
|
||||
from exo.utils.channels import Sender, channel
|
||||
|
||||
NODE_ID = NodeId("aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa")
|
||||
MODEL_ID = ModelId("test-org/test-model")
|
||||
@@ -133,7 +133,7 @@ async def test_re_download_after_delete_completes() -> None:
|
||||
the DownloadCoordinator."""
|
||||
cmd_send: Sender[ForwarderDownloadCommand]
|
||||
cmd_send, cmd_recv = channel[ForwarderDownloadCommand]()
|
||||
event_send, event_recv = channel[Event]()
|
||||
event_send, _event_recv = channel[Event]()
|
||||
|
||||
fake_downloader = FakeShardDownloader()
|
||||
wrapped_downloader = SingletonShardDownloader(fake_downloader)
|
||||
@@ -160,8 +160,7 @@ async def test_re_download_after_delete_completes() -> None:
|
||||
)
|
||||
)
|
||||
|
||||
# Wait for DownloadCompleted
|
||||
first_completed = await _wait_for_download_completed(event_recv, MODEL_ID)
|
||||
first_completed = await _wait_for_download_completed(coordinator, MODEL_ID)
|
||||
assert first_completed is not None, "First download should complete"
|
||||
|
||||
# 2. Delete the model
|
||||
@@ -182,8 +181,7 @@ async def test_re_download_after_delete_completes() -> None:
|
||||
)
|
||||
)
|
||||
|
||||
# Wait for second DownloadCompleted — this is the bug: it never arrives
|
||||
second_completed = await _wait_for_download_completed(event_recv, MODEL_ID)
|
||||
second_completed = await _wait_for_download_completed(coordinator, MODEL_ID)
|
||||
assert second_completed is not None, (
|
||||
"Re-download after deletion should complete"
|
||||
)
|
||||
@@ -195,19 +193,15 @@ async def test_re_download_after_delete_completes() -> None:
|
||||
|
||||
|
||||
async def _wait_for_download_completed(
|
||||
event_recv: Receiver[Event], model_id: ModelId, timeout: float = 2.0
|
||||
coordinator: DownloadCoordinator, model_id: ModelId, timeout: float = 2.0
|
||||
) -> DownloadCompleted | None:
|
||||
"""Drain events until we see a DownloadCompleted for the given model, or timeout."""
|
||||
"""Wait until coordinator state marks the given model as complete."""
|
||||
try:
|
||||
async with asyncio.timeout(timeout):
|
||||
while True:
|
||||
event = await event_recv.receive()
|
||||
if (
|
||||
isinstance(event, NodeDownloadProgress)
|
||||
and isinstance(event.download_progress, DownloadCompleted)
|
||||
and event.download_progress.shard_metadata.model_card.model_id
|
||||
== model_id
|
||||
):
|
||||
return event.download_progress
|
||||
progress = coordinator.download_status.get(model_id)
|
||||
if isinstance(progress, DownloadCompleted):
|
||||
return progress
|
||||
await asyncio.sleep(0.01)
|
||||
except TimeoutError:
|
||||
return None
|
||||
+23
-11
@@ -10,7 +10,7 @@ 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 Pidfile, PidfileError, SessionHandle
|
||||
from loguru import logger
|
||||
from pydantic import PositiveInt
|
||||
|
||||
@@ -47,18 +47,21 @@ class Node:
|
||||
node_id: NodeId
|
||||
offline: bool
|
||||
_api_port: int
|
||||
_sh: SessionHandle
|
||||
_tg: TaskGroup = field(init=False, default_factory=TaskGroup)
|
||||
|
||||
@classmethod
|
||||
async def create(cls, args: "Args") -> Self:
|
||||
node_id = get_node_zid()
|
||||
session_id = SessionId(master_node_id=node_id, election_clock=0)
|
||||
router = Router.create(
|
||||
session_handle, _nh = SessionHandle.new(
|
||||
node_id,
|
||||
namespace=args.namespace,
|
||||
listen_port=args.zenoh_port,
|
||||
discovery_service_port=args.discovery_port,
|
||||
)
|
||||
router = Router(_nh)
|
||||
|
||||
await router.register_topic(topics.GLOBAL_EVENTS)
|
||||
await router.register_topic(topics.LOCAL_EVENTS)
|
||||
await router.register_topic(topics.COMMANDS)
|
||||
@@ -84,6 +87,7 @@ class Node:
|
||||
exo_shard_downloader(offline=args.offline),
|
||||
event_sender=event_router.sender(),
|
||||
download_command_receiver=router.receiver(topics.DOWNLOAD_COMMANDS),
|
||||
session_handle=session_handle,
|
||||
offline=args.offline,
|
||||
)
|
||||
else:
|
||||
@@ -97,6 +101,7 @@ class Node:
|
||||
command_sender=router.sender(topics.COMMANDS),
|
||||
download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
|
||||
election_receiver=router.receiver(topics.ELECTION_MESSAGES),
|
||||
session_handle=session_handle,
|
||||
)
|
||||
else:
|
||||
api = None
|
||||
@@ -108,6 +113,7 @@ class Node:
|
||||
event_sender=event_router.sender(),
|
||||
command_sender=router.sender(topics.COMMANDS),
|
||||
download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
|
||||
session_handle=session_handle,
|
||||
api_port=args.api_port,
|
||||
)
|
||||
else:
|
||||
@@ -122,6 +128,8 @@ class Node:
|
||||
local_event_receiver=router.receiver(topics.LOCAL_EVENTS),
|
||||
command_receiver=router.receiver(topics.COMMANDS),
|
||||
download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
|
||||
aggregator=session_handle.last_value_aggregator("node_metrics"),
|
||||
storage=session_handle.storage_interface(),
|
||||
)
|
||||
|
||||
er_send, er_recv = channel[ElectionResult]()
|
||||
@@ -150,6 +158,7 @@ class Node:
|
||||
node_id,
|
||||
args.offline,
|
||||
args.api_port,
|
||||
session_handle,
|
||||
)
|
||||
|
||||
async def run(self):
|
||||
@@ -225,6 +234,8 @@ class Node:
|
||||
download_command_sender=self.router.sender(
|
||||
topics.DOWNLOAD_COMMANDS
|
||||
),
|
||||
aggregator=self._sh.last_value_aggregator("node_metrics"),
|
||||
storage=self._sh.storage_interface(),
|
||||
)
|
||||
self._tg.start_soon(self.master.run)
|
||||
elif (
|
||||
@@ -250,6 +261,7 @@ class Node:
|
||||
download_command_receiver=self.router.receiver(
|
||||
topics.DOWNLOAD_COMMANDS
|
||||
),
|
||||
session_handle=self._sh,
|
||||
offline=self.offline,
|
||||
)
|
||||
self._tg.start_soon(self.download_coordinator.run)
|
||||
@@ -264,6 +276,7 @@ class Node:
|
||||
download_command_sender=self.router.sender(
|
||||
topics.DOWNLOAD_COMMANDS
|
||||
),
|
||||
session_handle=self._sh,
|
||||
api_port=self._api_port,
|
||||
)
|
||||
self._tg.start_soon(self.worker.run)
|
||||
@@ -378,17 +391,16 @@ def main_inner(args: "Args"):
|
||||
|
||||
|
||||
class Args(FrozenModel):
|
||||
verbosity: int = 0
|
||||
force_master: bool = False
|
||||
spawn_api: bool = False
|
||||
api_port: PositiveInt = 52415
|
||||
tb_only: bool = False
|
||||
verbosity: int
|
||||
force_master: bool
|
||||
spawn_api: bool
|
||||
api_port: PositiveInt
|
||||
no_worker: bool = False
|
||||
no_downloads: bool = False
|
||||
offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true"
|
||||
no_batch: bool = False
|
||||
offline: bool
|
||||
no_batch: bool
|
||||
fast_synch: bool | None = None # None = auto, True = force on, False = force off
|
||||
legacy_daemon: bool = False
|
||||
legacy_daemon: bool
|
||||
bootstrap_peers: list[str] = []
|
||||
namespace: str
|
||||
zenoh_port: int
|
||||
@@ -476,7 +488,7 @@ class Args(FrozenModel):
|
||||
type=int,
|
||||
default=52414,
|
||||
dest="zenoh_port",
|
||||
help="Fixed TCP port for zenoh to listen.",
|
||||
help="Fixed port for zenoh to listen on.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--discovery-port",
|
||||
|
||||
+52
-117
@@ -1,14 +1,13 @@
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import anyio
|
||||
from exo_rs import LVAggregator, Storage
|
||||
from loguru import logger
|
||||
from pydantic import ValidationError
|
||||
|
||||
from exo.master.placement import (
|
||||
add_instance_to_placements,
|
||||
cancel_unnecessary_downloads,
|
||||
delete_instance,
|
||||
get_transition_events,
|
||||
place_instance,
|
||||
)
|
||||
from exo.master.placement_utils import find_ip_prioritised
|
||||
from exo.routing.event_router import (
|
||||
@@ -18,40 +17,26 @@ from exo.routing.event_router import (
|
||||
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,
|
||||
DeleteCustomModelCard,
|
||||
DeleteInstance,
|
||||
DeleteInstanceLink,
|
||||
ForwarderCommand,
|
||||
ForwarderDownloadCommand,
|
||||
ImageEdits,
|
||||
ImageGeneration,
|
||||
PlaceInstance,
|
||||
RequestEventLog,
|
||||
SendInputChunk,
|
||||
SetInstanceLink,
|
||||
TaskCancelled,
|
||||
TaskFinished,
|
||||
TestCommand,
|
||||
TextGeneration,
|
||||
)
|
||||
from exo.shared.types.common import CommandId, NodeId, SessionId, SystemId
|
||||
from exo.shared.types.events import (
|
||||
CustomModelCardAdded,
|
||||
CustomModelCardDeleted,
|
||||
Event,
|
||||
GlobalForwarderEvent,
|
||||
IndexedEvent,
|
||||
InputChunkReceived,
|
||||
InstanceDeleted,
|
||||
InstanceLinkCreated,
|
||||
InstanceLinkDeleted,
|
||||
LocalForwarderEvent,
|
||||
NodeGatheredInfo,
|
||||
NodeTimedOut,
|
||||
TaskCreated,
|
||||
TaskDeleted,
|
||||
TaskStatusUpdated,
|
||||
TraceEventData,
|
||||
TracesCollected,
|
||||
@@ -79,16 +64,18 @@ from exo.utils.event_buffer import MultiSourceBuffer
|
||||
from exo.utils.task_group import TaskGroup
|
||||
|
||||
|
||||
def _prefill_endpoint_for(state: State, decode_instance_id: InstanceId) -> str | None:
|
||||
def _prefill_endpoint_for(
|
||||
state: State, instance_links: list[InstanceLink], decode_instance_id: InstanceId
|
||||
) -> str | None:
|
||||
decode = state.instances.get(decode_instance_id)
|
||||
if decode is None:
|
||||
return None
|
||||
decode_node = next(iter(decode.shard_assignments.node_to_runner.keys()), None)
|
||||
if decode_node is None:
|
||||
return None
|
||||
decode_node = decode.shard_assignments.shards[
|
||||
decode.shard_assignments.primary_output_node
|
||||
].node_id
|
||||
|
||||
sources: set[InstanceId] = set()
|
||||
for link in state.instance_links.values():
|
||||
for link in instance_links:
|
||||
if decode_instance_id in link.decode_instances:
|
||||
sources.update(link.prefill_instances)
|
||||
sources.discard(decode_instance_id)
|
||||
@@ -106,7 +93,7 @@ def _prefill_endpoint_for(state: State, decode_instance_id: InstanceId) -> str |
|
||||
instance = state.instances.get(src_id)
|
||||
if instance is None:
|
||||
continue
|
||||
for node_id, runner_id in instance.shard_assignments.node_to_runner.items():
|
||||
for node_id, runner_id, _ in instance.shard_assignments.shards:
|
||||
port = state.prefill_server_ports.get(runner_id)
|
||||
if port is None:
|
||||
continue
|
||||
@@ -130,6 +117,8 @@ class Master:
|
||||
local_event_receiver: Receiver[LocalForwarderEvent],
|
||||
global_event_sender: Sender[GlobalForwarderEvent],
|
||||
download_command_sender: Sender[ForwarderDownloadCommand],
|
||||
aggregator: LVAggregator,
|
||||
storage: Storage,
|
||||
):
|
||||
self.node_id = node_id
|
||||
self.session_id = session_id
|
||||
@@ -145,7 +134,9 @@ class Master:
|
||||
self._multi_buffer = MultiSourceBuffer[SystemId, Event]()
|
||||
self._event_log = DiskEventLog(EXO_EVENT_LOG_DIR / "master")
|
||||
self._pending_traces: dict[TaskId, dict[int, list[TraceEventData]]] = {}
|
||||
self._expected_ranks: dict[TaskId, set[int]] = {}
|
||||
self._world_sizes: dict[TaskId, int] = {}
|
||||
self.aggregator = aggregator
|
||||
self.storage = storage
|
||||
|
||||
async def run(self):
|
||||
logger.info("Starting Master")
|
||||
@@ -182,10 +173,21 @@ class Master:
|
||||
pass
|
||||
case TextGeneration():
|
||||
# set-difference => prefill-only nodes
|
||||
instance_links: list[InstanceLink] = []
|
||||
prefill_only: set[InstanceId] = set()
|
||||
for link in self.state.instance_links.values():
|
||||
for _, link in (
|
||||
await self.storage.dump("instance_links/")
|
||||
).items():
|
||||
try:
|
||||
instance_links.append(
|
||||
InstanceLink.model_validate_json(link)
|
||||
)
|
||||
except ValidationError:
|
||||
continue
|
||||
|
||||
for link in instance_links:
|
||||
prefill_only.update(link.prefill_instances)
|
||||
for link in self.state.instance_links.values():
|
||||
for link in instance_links:
|
||||
prefill_only.difference_update(link.decode_instances)
|
||||
|
||||
for instance in self.state.instances.values():
|
||||
@@ -225,7 +227,9 @@ class Master:
|
||||
params = command.task_params.model_copy(
|
||||
update={
|
||||
"prefill_endpoint": _prefill_endpoint_for(
|
||||
self.state, decode_instance_id
|
||||
self.state.with_aggregator(self.aggregator),
|
||||
instance_links,
|
||||
decode_instance_id,
|
||||
),
|
||||
}
|
||||
)
|
||||
@@ -293,11 +297,9 @@ class Master:
|
||||
selected_instance_id
|
||||
)
|
||||
if selected_instance:
|
||||
ranks = set(
|
||||
shard.device_rank
|
||||
for shard in selected_instance.shard_assignments.runner_to_shard.values()
|
||||
self._world_sizes[task_id] = len(
|
||||
selected_instance.shard_assignments.shards
|
||||
)
|
||||
self._expected_ranks[task_id] = ranks
|
||||
case ImageEdits():
|
||||
for instance in self.state.instances.values():
|
||||
if (
|
||||
@@ -349,57 +351,20 @@ class Master:
|
||||
selected_instance_id
|
||||
)
|
||||
if selected_instance:
|
||||
ranks = set(
|
||||
shard.device_rank
|
||||
for shard in selected_instance.shard_assignments.runner_to_shard.values()
|
||||
self._world_sizes[task_id] = len(
|
||||
selected_instance.shard_assignments.shards
|
||||
)
|
||||
self._expected_ranks[task_id] = ranks
|
||||
case DeleteInstance():
|
||||
placement = delete_instance(command, self.state.instances)
|
||||
transition_events = get_transition_events(
|
||||
self.state.instances, placement, self.state.tasks
|
||||
)
|
||||
for cmd in cancel_unnecessary_downloads(
|
||||
placement, self.state.downloads
|
||||
):
|
||||
await self.download_command_sender.send(
|
||||
ForwarderDownloadCommand(
|
||||
origin=self._system_id, command=cmd
|
||||
)
|
||||
)
|
||||
generated_events.extend(transition_events)
|
||||
case PlaceInstance():
|
||||
placement = place_instance(
|
||||
command,
|
||||
self.state.topology,
|
||||
self.state.instances,
|
||||
self.state.node_memory,
|
||||
self.state.node_network,
|
||||
self.state.node_backends,
|
||||
download_status=self.state.downloads,
|
||||
node_rdma_ctl=self.state.node_rdma_ctl,
|
||||
)
|
||||
transition_events = get_transition_events(
|
||||
self.state.instances, placement, self.state.tasks
|
||||
)
|
||||
generated_events.extend(transition_events)
|
||||
case CreateInstance():
|
||||
state = self.state.with_aggregator(self.aggregator)
|
||||
placement = add_instance_to_placements(
|
||||
command,
|
||||
self.state.topology,
|
||||
self.state.instances,
|
||||
state.topology,
|
||||
state.instances,
|
||||
)
|
||||
transition_events = get_transition_events(
|
||||
self.state.instances, placement, self.state.tasks
|
||||
state.instances, placement, state.tasks
|
||||
)
|
||||
generated_events.extend(transition_events)
|
||||
case SendInputChunk(chunk=chunk):
|
||||
generated_events.append(
|
||||
InputChunkReceived(
|
||||
command_id=chunk.command_id,
|
||||
chunk=chunk,
|
||||
)
|
||||
)
|
||||
case TaskCancelled():
|
||||
if (
|
||||
task_id := self.command_task_mapping.get(
|
||||
@@ -416,41 +381,6 @@ class Master:
|
||||
logger.warning(
|
||||
f"Nonexistent command {command.cancelled_command_id} cancelled"
|
||||
)
|
||||
case TaskFinished():
|
||||
if (
|
||||
task_id := self.command_task_mapping.pop(
|
||||
command.finished_command_id, None
|
||||
)
|
||||
) is not None:
|
||||
generated_events.append(TaskDeleted(task_id=task_id))
|
||||
else:
|
||||
logger.warning(
|
||||
f"Finished command {command.finished_command_id} finished"
|
||||
)
|
||||
|
||||
case AddCustomModelCard():
|
||||
generated_events.append(
|
||||
CustomModelCardAdded(model_card=command.model_card)
|
||||
)
|
||||
case DeleteCustomModelCard():
|
||||
generated_events.append(
|
||||
CustomModelCardDeleted(model_id=command.model_id)
|
||||
)
|
||||
case SetInstanceLink():
|
||||
link = InstanceLink(
|
||||
link_id=command.link_id,
|
||||
prefill_instances=list(
|
||||
dict.fromkeys(command.prefill_instances)
|
||||
),
|
||||
decode_instances=list(
|
||||
dict.fromkeys(command.decode_instances)
|
||||
),
|
||||
)
|
||||
generated_events.append(InstanceLinkCreated(link=link))
|
||||
case DeleteInstanceLink():
|
||||
generated_events.append(
|
||||
InstanceLinkDeleted(link_id=command.link_id)
|
||||
)
|
||||
case RequestEventLog():
|
||||
# We should just be able to send everything, since other buffers will ignore old messages
|
||||
# rate limit to 1000 at a time
|
||||
@@ -462,6 +392,8 @@ class Master:
|
||||
await self._send_indexed_event(
|
||||
IndexedEvent(idx=i, event=event)
|
||||
)
|
||||
case other:
|
||||
logger.warning(f"ONE SLIPPED THROUGH {other}")
|
||||
for event in generated_events:
|
||||
await self.event_sender.send(event)
|
||||
except Exception as e:
|
||||
@@ -471,9 +403,11 @@ class Master:
|
||||
async def _plan(self) -> None:
|
||||
while True:
|
||||
# kill broken instances
|
||||
connected_node_ids = set(self.state.topology.list_nodes())
|
||||
connected_node_ids = set(
|
||||
self.state.with_aggregator(self.aggregator).topology.list_nodes()
|
||||
)
|
||||
for instance_id, instance in self.state.instances.items():
|
||||
for node_id in instance.shard_assignments.node_to_runner:
|
||||
for node_id, _, _ in instance.shard_assignments.shards:
|
||||
if node_id not in connected_node_ids:
|
||||
await self.event_sender.send(
|
||||
InstanceDeleted(instance_id=instance_id)
|
||||
@@ -481,7 +415,9 @@ class Master:
|
||||
break
|
||||
|
||||
# time out dead nodes
|
||||
for node_id, time in self.state.last_seen.items():
|
||||
for node_id, time in self.state.with_aggregator(
|
||||
self.aggregator
|
||||
).last_seen.items():
|
||||
now = datetime.now(tz=timezone.utc)
|
||||
if now - time > timedelta(seconds=30):
|
||||
logger.info(f"Manually removing node {node_id} due to inactivity")
|
||||
@@ -540,9 +476,8 @@ class Master:
|
||||
self._pending_traces[task_id][event.rank] = event.traces
|
||||
|
||||
if (
|
||||
task_id in self._expected_ranks
|
||||
and set(self._pending_traces[task_id].keys())
|
||||
>= self._expected_ranks[task_id]
|
||||
task_id in self._world_sizes
|
||||
and len(self._pending_traces[task_id]) >= self._world_sizes[task_id]
|
||||
):
|
||||
await self._merge_and_save_traces(task_id)
|
||||
|
||||
@@ -556,5 +491,5 @@ class Master:
|
||||
)
|
||||
|
||||
del self._pending_traces[task_id]
|
||||
if task_id in self._expected_ranks:
|
||||
del self._expected_ranks[task_id]
|
||||
if task_id in self._world_sizes:
|
||||
del self._world_sizes[task_id]
|
||||
@@ -1,5 +1,4 @@
|
||||
from collections.abc import Mapping
|
||||
from copy import deepcopy
|
||||
from typing import Sequence
|
||||
|
||||
from exo.master.placement_utils import (
|
||||
@@ -17,7 +16,6 @@ from exo.shared.types.backends import Backend
|
||||
from exo.shared.types.commands import (
|
||||
CancelDownload,
|
||||
CreateInstance,
|
||||
DeleteInstance,
|
||||
DownloadCommand,
|
||||
PlaceInstance,
|
||||
)
|
||||
@@ -106,14 +104,13 @@ def _cycle_download_score(
|
||||
def place_instance(
|
||||
command: PlaceInstance,
|
||||
topology: Topology,
|
||||
current_instances: Mapping[InstanceId, Instance],
|
||||
node_memory: Mapping[NodeId, MemoryUsage],
|
||||
node_network: Mapping[NodeId, NodeNetworkInfo],
|
||||
node_backends: Mapping[NodeId, list[Backend]],
|
||||
required_nodes: set[NodeId] | None = None,
|
||||
download_status: Mapping[NodeId, Sequence[DownloadProgress]] | None = None,
|
||||
node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus] | None = None,
|
||||
) -> dict[InstanceId, Instance]:
|
||||
) -> Instance:
|
||||
cycles = topology.get_cycles()
|
||||
candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles))
|
||||
|
||||
@@ -258,24 +255,10 @@ def place_instance(
|
||||
cycle_digraph: Topology = topology.get_subgraph_from_nodes(selected_cycle.node_ids)
|
||||
|
||||
instance_id = InstanceId()
|
||||
target_instances = dict(deepcopy(current_instances))
|
||||
|
||||
match command.instance_meta:
|
||||
case InstanceMeta.MlxJaccl:
|
||||
# TODO(evan): shard assignments should contain information about ranks, this is ugly
|
||||
def get_device_rank(node_id: NodeId) -> int:
|
||||
runner_id = shard_assignments.node_to_runner[node_id]
|
||||
shard_metadata = shard_assignments.runner_to_shard.get(runner_id)
|
||||
assert shard_metadata is not None
|
||||
return shard_metadata.device_rank
|
||||
|
||||
zero_node_ids = [
|
||||
node_id
|
||||
for node_id in selected_cycle.node_ids
|
||||
if get_device_rank(node_id) == 0
|
||||
]
|
||||
assert len(zero_node_ids) == 1
|
||||
coordinator_node_id = zero_node_ids[0]
|
||||
coordinator_node_id = shard_assignments.shards[0].node_id
|
||||
|
||||
mlx_jaccl_devices = get_mlx_jaccl_devices_matrix(
|
||||
[node_id for node_id in selected_cycle],
|
||||
@@ -287,7 +270,7 @@ def place_instance(
|
||||
cycle_digraph=cycle_digraph,
|
||||
node_network=node_network,
|
||||
)
|
||||
target_instances[instance_id] = MlxJacclInstance(
|
||||
return MlxJacclInstance(
|
||||
instance_id=instance_id,
|
||||
shard_assignments=shard_assignments,
|
||||
jaccl_devices=mlx_jaccl_devices,
|
||||
@@ -301,26 +284,13 @@ def place_instance(
|
||||
ephemeral_port=ephemeral_port,
|
||||
node_network=node_network,
|
||||
)
|
||||
target_instances[instance_id] = MlxRingInstance(
|
||||
return MlxRingInstance(
|
||||
instance_id=instance_id,
|
||||
shard_assignments=shard_assignments,
|
||||
hosts_by_node=hosts_by_node,
|
||||
ephemeral_port=ephemeral_port,
|
||||
)
|
||||
|
||||
return target_instances
|
||||
|
||||
|
||||
def delete_instance(
|
||||
command: DeleteInstance,
|
||||
current_instances: Mapping[InstanceId, Instance],
|
||||
) -> dict[InstanceId, Instance]:
|
||||
target_instances = dict(deepcopy(current_instances))
|
||||
if command.instance_id in target_instances:
|
||||
del target_instances[command.instance_id]
|
||||
return target_instances
|
||||
raise ValueError(f"Instance {command.instance_id} not found")
|
||||
|
||||
|
||||
def get_transition_events(
|
||||
current_instances: Mapping[InstanceId, Instance],
|
||||
@@ -376,10 +346,10 @@ def cancel_unnecessary_downloads(
|
||||
active_models = set(
|
||||
(
|
||||
node_id,
|
||||
instance.shard_assignments.runner_to_shard[runner_id].model_card.model_id,
|
||||
instance.shard_assignments.model_id,
|
||||
)
|
||||
for instance in instances.values()
|
||||
for node_id, runner_id in instance.shard_assignments.node_to_runner.items()
|
||||
for node_id, _, _ in instance.shard_assignments.shards
|
||||
)
|
||||
for pair in currently_downloading:
|
||||
if pair not in active_models:
|
||||
|
||||
@@ -8,12 +8,11 @@ from exo.shared.types.common import Host, NodeId
|
||||
from exo.shared.types.memory import Memory
|
||||
from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
|
||||
from exo.shared.types.topology import Cycle, RDMAConnection, SocketConnection
|
||||
from exo.shared.types.worker.runners import RunnerId, ShardAssignments
|
||||
from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardWithId
|
||||
from exo.shared.types.worker.shards import (
|
||||
CfgShardMetadata,
|
||||
PipelineShardMetadata,
|
||||
Sharding,
|
||||
ShardMetadata,
|
||||
TensorShardMetadata,
|
||||
)
|
||||
|
||||
@@ -152,27 +151,27 @@ def _get_shard_assignments_for_cfg_parallel(
|
||||
_validate_cycle(cycle)
|
||||
|
||||
world_size = len(cycle)
|
||||
cfg_world_size = 2
|
||||
pipeline_world_size = world_size // cfg_world_size
|
||||
pipeline_world_size = world_size // 2
|
||||
|
||||
# Allocate layers for one pipeline group (both groups run the same layers)
|
||||
pipeline_node_ids = cycle.node_ids[:pipeline_world_size]
|
||||
pipeline_memory = _compute_total_memory(pipeline_node_ids, node_memory)
|
||||
|
||||
# nb: only validates the forward path...
|
||||
layer_allocations = _allocate_and_validate_layers(
|
||||
pipeline_node_ids, node_memory, pipeline_memory, model_card
|
||||
)
|
||||
|
||||
# Ring topology: group 0 ascending [0,1,2,...], group 1 descending [...,2,1,0]
|
||||
# This places both last stages as neighbors for CFG exchange.
|
||||
position_to_cfg_pipeline = [(0, r) for r in range(pipeline_world_size)] + [
|
||||
(1, r) for r in reversed(range(pipeline_world_size))
|
||||
]
|
||||
position_to_cfg_pipeline = list(range(pipeline_world_size)) + list(
|
||||
reversed(range(pipeline_world_size))
|
||||
)
|
||||
|
||||
runner_to_shard: dict[RunnerId, ShardMetadata] = {}
|
||||
node_to_runner: dict[NodeId, RunnerId] = {}
|
||||
shards: list[ShardWithId] = []
|
||||
|
||||
for device_rank, node_id in enumerate(cycle.node_ids):
|
||||
cfg_rank, pipeline_rank = position_to_cfg_pipeline[device_rank]
|
||||
pipeline_rank = position_to_cfg_pipeline[device_rank]
|
||||
layers_before = sum(layer_allocations[:pipeline_rank])
|
||||
node_layers = layer_allocations[pipeline_rank]
|
||||
|
||||
@@ -183,20 +182,15 @@ def _get_shard_assignments_for_cfg_parallel(
|
||||
start_layer=layers_before,
|
||||
end_layer=layers_before + node_layers,
|
||||
n_layers=model_card.n_layers,
|
||||
cfg_rank=cfg_rank,
|
||||
cfg_world_size=cfg_world_size,
|
||||
pipeline_rank=pipeline_rank,
|
||||
pipeline_world_size=pipeline_world_size,
|
||||
)
|
||||
|
||||
runner_id = RunnerId()
|
||||
runner_to_shard[runner_id] = shard
|
||||
node_to_runner[node_id] = runner_id
|
||||
shards.append(ShardWithId(node_id, runner_id, shard))
|
||||
|
||||
return ShardAssignments(
|
||||
model_id=model_card.model_id,
|
||||
runner_to_shard=runner_to_shard,
|
||||
node_to_runner=node_to_runner,
|
||||
shards=shards,
|
||||
primary_output_node=pipeline_world_size - 1,
|
||||
)
|
||||
|
||||
|
||||
@@ -208,13 +202,13 @@ def _get_shard_assignments_for_pure_pipeline(
|
||||
"""Create shard assignments for pure pipeline execution."""
|
||||
_validate_cycle(cycle)
|
||||
total_memory = _compute_total_memory(cycle.node_ids, node_memory)
|
||||
world_size = len(cycle)
|
||||
|
||||
layer_allocations = _allocate_and_validate_layers(
|
||||
cycle.node_ids, node_memory, total_memory, model_card
|
||||
)
|
||||
|
||||
runner_to_shard: dict[RunnerId, ShardMetadata] = {}
|
||||
node_to_runner: dict[NodeId, RunnerId] = {}
|
||||
shards: list[ShardWithId] = []
|
||||
|
||||
for pipeline_rank, node_id in enumerate(cycle.node_ids):
|
||||
layers_before = sum(layer_allocations[:pipeline_rank])
|
||||
@@ -223,20 +217,17 @@ def _get_shard_assignments_for_pure_pipeline(
|
||||
shard = PipelineShardMetadata(
|
||||
model_card=model_card,
|
||||
device_rank=pipeline_rank,
|
||||
world_size=len(cycle),
|
||||
world_size=world_size,
|
||||
start_layer=layers_before,
|
||||
end_layer=layers_before + node_layers,
|
||||
n_layers=model_card.n_layers,
|
||||
)
|
||||
|
||||
runner_id = RunnerId()
|
||||
runner_to_shard[runner_id] = shard
|
||||
node_to_runner[node_id] = runner_id
|
||||
shards.append(ShardWithId(node_id, runner_id, shard))
|
||||
|
||||
return ShardAssignments(
|
||||
model_id=model_card.model_id,
|
||||
runner_to_shard=runner_to_shard,
|
||||
node_to_runner=node_to_runner,
|
||||
model_id=model_card.model_id, shards=shards, primary_output_node=world_size - 1
|
||||
)
|
||||
|
||||
|
||||
@@ -246,8 +237,7 @@ def get_shard_assignments_for_tensor_parallel(
|
||||
):
|
||||
total_layers = model_card.n_layers
|
||||
world_size = len(cycle)
|
||||
runner_to_shard: dict[RunnerId, ShardMetadata] = {}
|
||||
node_to_runner: dict[NodeId, RunnerId] = {}
|
||||
shards: list[ShardWithId] = []
|
||||
|
||||
for i, node_id in enumerate(cycle):
|
||||
shard = TensorShardMetadata(
|
||||
@@ -260,14 +250,10 @@ def get_shard_assignments_for_tensor_parallel(
|
||||
)
|
||||
|
||||
runner_id = RunnerId()
|
||||
|
||||
runner_to_shard[runner_id] = shard
|
||||
node_to_runner[node_id] = runner_id
|
||||
shards.append(ShardWithId(node_id, runner_id, shard))
|
||||
|
||||
shard_assignments = ShardAssignments(
|
||||
model_id=model_card.model_id,
|
||||
runner_to_shard=runner_to_shard,
|
||||
node_to_runner=node_to_runner,
|
||||
model_id=model_card.model_id, shards=shards, primary_output_node=world_size - 1
|
||||
)
|
||||
|
||||
return shard_assignments
|
||||
|
||||
@@ -6,7 +6,6 @@ import pytest
|
||||
from loguru import logger
|
||||
|
||||
from exo.master.main import Master
|
||||
from exo.routing.router import get_node_zid
|
||||
from exo.shared.models.model_cards import ModelCard, ModelTask
|
||||
from exo.shared.types.backends import Backend
|
||||
from exo.shared.types.commands import (
|
||||
@@ -16,7 +15,7 @@ from exo.shared.types.commands import (
|
||||
PlaceInstance,
|
||||
TextGeneration,
|
||||
)
|
||||
from exo.shared.types.common import ModelId, SessionId, SystemId
|
||||
from exo.shared.types.common import ModelId, NodeId, SessionId, SystemId
|
||||
from exo.shared.types.events import (
|
||||
Event,
|
||||
GlobalForwarderEvent,
|
||||
@@ -42,14 +41,34 @@ from exo.shared.types.worker.instances import (
|
||||
MlxRingInstance,
|
||||
ShardAssignments,
|
||||
)
|
||||
from exo.shared.types.worker.runners import ShardWithId
|
||||
from exo.shared.types.worker.shards import PipelineShardMetadata, Sharding
|
||||
from exo.utils.channels import channel
|
||||
from exo.utils.info_gatherer.info_gatherer import NodeBackends
|
||||
|
||||
|
||||
class MockAggregator:
|
||||
def dump(self) -> dict[str, str]:
|
||||
return {}
|
||||
|
||||
|
||||
class MockStorage:
|
||||
async def get(self, _: str) -> None:
|
||||
return None
|
||||
|
||||
async def put(self, _1: str, _2: str) -> None:
|
||||
return None
|
||||
|
||||
async def delete(self, _: str) -> None:
|
||||
return None
|
||||
|
||||
async def dump(self, _: str) -> dict[str, str]:
|
||||
return {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_master():
|
||||
node_id = get_node_zid()
|
||||
node_id = NodeId("yoooo")
|
||||
session_id = SessionId(master_node_id=node_id, election_clock=0)
|
||||
|
||||
ge_sender, global_event_receiver = channel[GlobalForwarderEvent]()
|
||||
@@ -94,6 +113,8 @@ async def test_master():
|
||||
local_event_receiver=le_receiver,
|
||||
command_receiver=co_receiver,
|
||||
download_command_sender=fcds,
|
||||
aggregator=MockAggregator(), # pyright: ignore[reportArgumentType]
|
||||
storage=MockStorage(), # pyright: ignore[reportArgumentType]
|
||||
)
|
||||
logger.info("run the master")
|
||||
async with anyio.create_task_group() as tg:
|
||||
@@ -205,29 +226,33 @@ async def test_master():
|
||||
assert isinstance(events[2].event, InstanceCreated)
|
||||
created_instance = events[2].event.instance
|
||||
assert isinstance(created_instance, MlxRingInstance)
|
||||
runner_id = list(created_instance.shard_assignments.runner_to_shard.keys())[0]
|
||||
runner_id = created_instance.shard_assignments.shards[0].runner_id
|
||||
# Validate the shard assignments
|
||||
expected_shard_assignments = ShardAssignments(
|
||||
model_id=ModelId("llama-3.2-1b"),
|
||||
runner_to_shard={
|
||||
(runner_id): PipelineShardMetadata(
|
||||
start_layer=0,
|
||||
end_layer=16,
|
||||
n_layers=16,
|
||||
model_card=ModelCard(
|
||||
model_id=ModelId("llama-3.2-1b"),
|
||||
shards=[
|
||||
ShardWithId(
|
||||
node_id,
|
||||
runner_id,
|
||||
PipelineShardMetadata(
|
||||
start_layer=0,
|
||||
end_layer=16,
|
||||
n_layers=16,
|
||||
storage_size=Memory.from_bytes(678948),
|
||||
hidden_size=7168,
|
||||
supports_tensor=True,
|
||||
tasks=[ModelTask.TextGeneration],
|
||||
backends=[Backend.MlxMetal],
|
||||
model_card=ModelCard(
|
||||
model_id=ModelId("llama-3.2-1b"),
|
||||
n_layers=16,
|
||||
storage_size=Memory.from_bytes(678948),
|
||||
hidden_size=7168,
|
||||
supports_tensor=True,
|
||||
tasks=[ModelTask.TextGeneration],
|
||||
backends=[Backend.MlxMetal],
|
||||
),
|
||||
device_rank=0,
|
||||
world_size=1,
|
||||
),
|
||||
device_rank=0,
|
||||
world_size=1,
|
||||
)
|
||||
},
|
||||
node_to_runner={node_id: runner_id},
|
||||
],
|
||||
primary_output_node=0,
|
||||
)
|
||||
assert created_instance.shard_assignments == expected_shard_assignments
|
||||
# For single-node, hosts_by_node should have one entry with self-binding
|
||||
|
||||
@@ -49,16 +49,36 @@ from exo.shared.types.worker.instances import (
|
||||
MlxJacclInstance,
|
||||
MlxRingInstance,
|
||||
)
|
||||
from exo.shared.types.worker.runners import ShardAssignments
|
||||
from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardWithId
|
||||
from exo.shared.types.worker.shards import PipelineShardMetadata, Sharding
|
||||
|
||||
|
||||
class MockShard:
|
||||
def is_primary_output(self) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def instance() -> Instance:
|
||||
def instance(model_card: ModelCard) -> Instance:
|
||||
return MlxRingInstance(
|
||||
instance_id=InstanceId(),
|
||||
shard_assignments=ShardAssignments(
|
||||
model_id=ModelId("test-model"), runner_to_shard={}, node_to_runner={}
|
||||
model_id=ModelId("test-model"),
|
||||
shards=[
|
||||
ShardWithId(
|
||||
NodeId(),
|
||||
RunnerId(),
|
||||
PipelineShardMetadata(
|
||||
model_card=model_card,
|
||||
device_rank=0,
|
||||
world_size=1,
|
||||
start_layer=0,
|
||||
end_layer=model_card.n_layers,
|
||||
n_layers=model_card.n_layers,
|
||||
),
|
||||
)
|
||||
],
|
||||
primary_output_node=0,
|
||||
),
|
||||
hosts_by_node={},
|
||||
ephemeral_port=50000,
|
||||
@@ -123,6 +143,11 @@ def test_get_instance_placements_create_instance(
|
||||
node_id_a = NodeId()
|
||||
node_id_b = NodeId()
|
||||
node_id_c = NodeId()
|
||||
node_to_layers = {
|
||||
node_id_a: expected_layers[0],
|
||||
node_id_b: expected_layers[1],
|
||||
node_id_c: expected_layers[2],
|
||||
}
|
||||
|
||||
# fully connected (directed) between the 3 nodes
|
||||
conn_a_b = Connection(
|
||||
@@ -175,22 +200,11 @@ def test_get_instance_placements_create_instance(
|
||||
instance = placements[instance_id]
|
||||
assert instance.shard_assignments.model_id == model_card.model_id
|
||||
|
||||
runner_id_a = instance.shard_assignments.node_to_runner[node_id_a]
|
||||
runner_id_b = instance.shard_assignments.node_to_runner[node_id_b]
|
||||
runner_id_c = instance.shard_assignments.node_to_runner[node_id_c]
|
||||
for nid, _, shard in (shards := instance.shard_assignments.shards):
|
||||
assert shard.end_layer - shard.start_layer == node_to_layers[nid]
|
||||
|
||||
shard_a = instance.shard_assignments.runner_to_shard[runner_id_a]
|
||||
shard_b = instance.shard_assignments.runner_to_shard[runner_id_b]
|
||||
shard_c = instance.shard_assignments.runner_to_shard[runner_id_c]
|
||||
|
||||
assert shard_a.end_layer - shard_a.start_layer == expected_layers[0]
|
||||
assert shard_b.end_layer - shard_b.start_layer == expected_layers[1]
|
||||
assert shard_c.end_layer - shard_c.start_layer == expected_layers[2]
|
||||
|
||||
shards = [shard_a, shard_b, shard_c]
|
||||
shards_sorted = sorted(shards, key=lambda s: s.start_layer)
|
||||
assert shards_sorted[0].start_layer == 0
|
||||
assert shards_sorted[-1].end_layer == total_layers
|
||||
assert shards[0].shard.start_layer == 0
|
||||
assert shards[-1].shard.end_layer == total_layers
|
||||
|
||||
|
||||
def test_get_instance_placements_one_node_exact_fit() -> None:
|
||||
@@ -218,9 +232,7 @@ def test_get_instance_placements_one_node_exact_fit() -> None:
|
||||
instance_id = list(placements.keys())[0]
|
||||
instance = placements[instance_id]
|
||||
assert instance.shard_assignments.model_id == "test-model"
|
||||
assert len(instance.shard_assignments.node_to_runner) == 1
|
||||
assert len(instance.shard_assignments.runner_to_shard) == 1
|
||||
assert len(instance.shard_assignments.runner_to_shard) == 1
|
||||
assert len(instance.shard_assignments.shards) == 1
|
||||
|
||||
|
||||
def test_get_instance_placements_one_node_fits_with_extra_memory() -> None:
|
||||
@@ -248,9 +260,7 @@ def test_get_instance_placements_one_node_fits_with_extra_memory() -> None:
|
||||
instance_id = list(placements.keys())[0]
|
||||
instance = placements[instance_id]
|
||||
assert instance.shard_assignments.model_id == "test-model"
|
||||
assert len(instance.shard_assignments.node_to_runner) == 1
|
||||
assert len(instance.shard_assignments.runner_to_shard) == 1
|
||||
assert len(instance.shard_assignments.runner_to_shard) == 1
|
||||
assert len(instance.shard_assignments.shards) == 1
|
||||
|
||||
|
||||
def test_get_instance_placements_one_node_not_fit() -> None:
|
||||
@@ -381,7 +391,7 @@ def test_placement_selects_leaf_nodes(
|
||||
assert len(placements) == 1
|
||||
instance = list(placements.values())[0]
|
||||
|
||||
assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
|
||||
assigned_nodes = set(map(lambda it: it.node_id, instance.shard_assignments.shards))
|
||||
assert assigned_nodes == set((node_id_a, node_id_b)) or assigned_nodes == set(
|
||||
(
|
||||
node_id_c,
|
||||
@@ -498,8 +508,8 @@ def test_tensor_rdma_backend_connectivity_matrix(
|
||||
for i in range(3):
|
||||
assert matrix[i][i] is None
|
||||
|
||||
assigned_nodes = list(instance.shard_assignments.node_to_runner.keys())
|
||||
node_to_idx = {node_id: idx for idx, node_id in enumerate(assigned_nodes)}
|
||||
assigned_nodes = list(instance.shard_assignments.shards)
|
||||
node_to_idx = {node_id: idx for idx, (node_id, _, _) in enumerate(assigned_nodes)}
|
||||
|
||||
idx_a = node_to_idx[node_a]
|
||||
idx_b = node_to_idx[node_b]
|
||||
@@ -511,7 +521,7 @@ def test_tensor_rdma_backend_connectivity_matrix(
|
||||
|
||||
# Verify coordinators are set for all nodes
|
||||
assert len(instance.jaccl_coordinators) == 3
|
||||
for node_id in assigned_nodes:
|
||||
for node_id, _, _ in assigned_nodes:
|
||||
assert node_id in instance.jaccl_coordinators
|
||||
coordinator = instance.jaccl_coordinators[node_id]
|
||||
assert ":" in coordinator
|
||||
@@ -825,7 +835,7 @@ def test_placement_prefers_cycle_with_downloaded_model(
|
||||
|
||||
assert len(placements) == 1
|
||||
instance = list(placements.values())[0]
|
||||
assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
|
||||
assigned_nodes = set(map(lambda it: it.node_id, instance.shard_assignments.shards))
|
||||
assert assigned_nodes == {node_b}
|
||||
|
||||
|
||||
@@ -903,7 +913,7 @@ def test_placement_prefers_cycle_with_higher_download_progress(
|
||||
|
||||
assert len(placements) == 1
|
||||
instance = list(placements.values())[0]
|
||||
assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
|
||||
assigned_nodes = set(map(lambda it: it.node_id, instance.shard_assignments.shards))
|
||||
assert assigned_nodes == {node_b}
|
||||
|
||||
|
||||
@@ -957,7 +967,7 @@ def test_placement_does_not_prefer_cycle_with_failed_download(
|
||||
|
||||
assert len(placements) == 1
|
||||
instance = list(placements.values())[0]
|
||||
assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
|
||||
assigned_nodes = set(map(lambda it: it.node_id, instance.shard_assignments.shards))
|
||||
# node_a should win on RAM tiebreaker since failed download scores 0.0
|
||||
assert assigned_nodes == {node_a}
|
||||
|
||||
|
||||
@@ -204,6 +204,11 @@ def test_get_shard_assignments(
|
||||
node_a_id = NodeId()
|
||||
node_b_id = NodeId()
|
||||
node_c_id = NodeId()
|
||||
layers_by_node = {
|
||||
node_a_id: expected_layers[0],
|
||||
node_b_id: expected_layers[1],
|
||||
node_c_id: expected_layers[2],
|
||||
}
|
||||
|
||||
# create connections (A -> B -> C -> A forms a 3-cycle, plus B -> A also exists)
|
||||
connection1 = Connection(
|
||||
@@ -258,25 +263,8 @@ def test_get_shard_assignments(
|
||||
)
|
||||
|
||||
# assert
|
||||
runner_id_a = shard_assignments.node_to_runner[node_a_id]
|
||||
runner_id_b = shard_assignments.node_to_runner[node_b_id]
|
||||
runner_id_c = shard_assignments.node_to_runner[node_c_id]
|
||||
|
||||
assert (
|
||||
shard_assignments.runner_to_shard[runner_id_a].end_layer
|
||||
- shard_assignments.runner_to_shard[runner_id_a].start_layer
|
||||
== expected_layers[0]
|
||||
)
|
||||
assert (
|
||||
shard_assignments.runner_to_shard[runner_id_b].end_layer
|
||||
- shard_assignments.runner_to_shard[runner_id_b].start_layer
|
||||
== expected_layers[1]
|
||||
)
|
||||
assert (
|
||||
shard_assignments.runner_to_shard[runner_id_c].end_layer
|
||||
- shard_assignments.runner_to_shard[runner_id_c].start_layer
|
||||
== expected_layers[2]
|
||||
)
|
||||
for nid, _, shard in shard_assignments.shards:
|
||||
assert shard.end_layer - shard.start_layer == layers_by_node[nid]
|
||||
|
||||
|
||||
def test_get_mlx_jaccl_coordinators():
|
||||
@@ -543,11 +531,11 @@ class TestCfgParallelPlacement:
|
||||
model_card, cycle, node_memory
|
||||
)
|
||||
|
||||
shards = list(assignments.runner_to_shard.values())
|
||||
shards = list(assignments.shards)
|
||||
assert len(shards) == 2
|
||||
|
||||
# CFG models should get CfgShardMetadata
|
||||
for shard in shards:
|
||||
for _, _, shard in shards:
|
||||
assert isinstance(shard, CfgShardMetadata)
|
||||
# Both nodes should have all layers (no pipeline split)
|
||||
assert shard.start_layer == 0
|
||||
@@ -558,7 +546,7 @@ class TestCfgParallelPlacement:
|
||||
assert shard.pipeline_rank == 0
|
||||
|
||||
cfg_ranks = sorted(
|
||||
s.cfg_rank for s in shards if isinstance(s, CfgShardMetadata)
|
||||
s.shard.cfg_rank for s in shards if isinstance(s.shard, CfgShardMetadata)
|
||||
)
|
||||
assert cfg_ranks == [0, 1]
|
||||
|
||||
@@ -587,11 +575,11 @@ class TestCfgParallelPlacement:
|
||||
model_card, cycle, node_memory
|
||||
)
|
||||
|
||||
shards = list(assignments.runner_to_shard.values())
|
||||
shards = assignments.shards
|
||||
assert len(shards) == 4
|
||||
|
||||
# CFG models should get CfgShardMetadata
|
||||
for shard in shards:
|
||||
for _, _, shard in shards:
|
||||
assert isinstance(shard, CfgShardMetadata)
|
||||
assert shard.cfg_world_size == 2
|
||||
assert shard.pipeline_world_size == 2
|
||||
@@ -599,10 +587,14 @@ class TestCfgParallelPlacement:
|
||||
|
||||
# Check we have 2 nodes in each CFG group
|
||||
cfg_0_shards = [
|
||||
s for s in shards if isinstance(s, CfgShardMetadata) and s.cfg_rank == 0
|
||||
s.shard
|
||||
for s in shards
|
||||
if isinstance(s.shard, CfgShardMetadata) and s.shard.cfg_rank == 0
|
||||
]
|
||||
cfg_1_shards = [
|
||||
s for s in shards if isinstance(s, CfgShardMetadata) and s.cfg_rank == 1
|
||||
s.shard
|
||||
for s in shards
|
||||
if isinstance(s.shard, CfgShardMetadata) and s.shard.cfg_rank == 1
|
||||
]
|
||||
assert len(cfg_0_shards) == 2
|
||||
assert len(cfg_1_shards) == 2
|
||||
@@ -637,11 +629,11 @@ class TestCfgParallelPlacement:
|
||||
model_card, cycle, node_memory
|
||||
)
|
||||
|
||||
shards = list(assignments.runner_to_shard.values())
|
||||
shards = list(assignments.shards)
|
||||
assert len(shards) == 3
|
||||
|
||||
# Odd node count with CFG model falls back to PipelineShardMetadata (sequential CFG)
|
||||
for shard in shards:
|
||||
for _, _, shard in shards:
|
||||
assert isinstance(shard, PipelineShardMetadata)
|
||||
|
||||
def test_two_nodes_non_cfg_model_uses_pipeline(self):
|
||||
@@ -673,18 +665,18 @@ class TestCfgParallelPlacement:
|
||||
model_card, cycle, node_memory
|
||||
)
|
||||
|
||||
shards = list(assignments.runner_to_shard.values())
|
||||
shards = list(assignments.shards)
|
||||
assert len(shards) == 2
|
||||
|
||||
# Non-CFG models should get PipelineShardMetadata
|
||||
for shard in shards:
|
||||
for _, _, shard in shards:
|
||||
assert isinstance(shard, PipelineShardMetadata)
|
||||
|
||||
# Should have actual layer sharding (pipeline)
|
||||
layer_ranges = sorted(
|
||||
(s.start_layer, s.end_layer)
|
||||
(s.shard.start_layer, s.shard.end_layer)
|
||||
for s in shards
|
||||
if isinstance(s, PipelineShardMetadata)
|
||||
if isinstance(s.shard, PipelineShardMetadata)
|
||||
)
|
||||
# First shard starts at 0, last shard ends at 57
|
||||
assert layer_ranges[0][0] == 0
|
||||
|
||||
+46
-136
@@ -4,26 +4,19 @@ from datetime import datetime
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from exo.shared.models.model_cards import ModelCard
|
||||
from exo.shared.types.common import ModelId, NodeId
|
||||
from exo.shared.types.common import NodeId
|
||||
from exo.shared.types.events import (
|
||||
ChunkGenerated,
|
||||
CustomModelCardAdded,
|
||||
CustomModelCardDeleted,
|
||||
Event,
|
||||
IndexedEvent,
|
||||
InputChunkReceived,
|
||||
InstanceCreated,
|
||||
InstanceDeleted,
|
||||
InstanceLinkCreated,
|
||||
InstanceLinkDeleted,
|
||||
NodeDownloadProgress,
|
||||
NodeGatheredInfo,
|
||||
NodeTimedOut,
|
||||
RunnerStatusUpdated,
|
||||
TaskAcknowledged,
|
||||
TaskCreated,
|
||||
TaskDeleted,
|
||||
TaskFailed,
|
||||
TaskStatusUpdated,
|
||||
TestEvent,
|
||||
@@ -32,7 +25,6 @@ from exo.shared.types.events import (
|
||||
TracesCollected,
|
||||
TracesMerged,
|
||||
)
|
||||
from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
|
||||
from exo.shared.types.profiling import (
|
||||
NodeIdentity,
|
||||
NodeNetworkInfo,
|
||||
@@ -42,7 +34,6 @@ from exo.shared.types.profiling import (
|
||||
)
|
||||
from exo.shared.types.state import State
|
||||
from exo.shared.types.tasks import Task, TaskId, TaskStatus
|
||||
from exo.shared.types.topology import Connection, RDMAConnection
|
||||
from exo.shared.types.worker.downloads import DownloadProgress
|
||||
from exo.shared.types.worker.instances import Instance, InstanceId
|
||||
from exo.shared.types.worker.runners import (
|
||||
@@ -67,18 +58,6 @@ from exo.utils.info_gatherer.info_gatherer import (
|
||||
)
|
||||
|
||||
|
||||
def _is_rdma_ctl_enabled(
|
||||
node_id: NodeId, node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus]
|
||||
) -> bool:
|
||||
"""A node is RDMA-capable only if rdma_ctl status has been observed as enabled.
|
||||
|
||||
Missing entries default to ``False`` — if we have not yet observed (or the node
|
||||
cannot run) ``rdma_ctl``, it must not participate in an RDMA-backed instance.
|
||||
"""
|
||||
status = node_rdma_ctl.get(node_id)
|
||||
return status is not None and status.enabled
|
||||
|
||||
|
||||
def event_apply(event: Event, state: State) -> State:
|
||||
"""Apply an event to state."""
|
||||
match event:
|
||||
@@ -86,15 +65,10 @@ def event_apply(event: Event, state: State) -> State:
|
||||
TestEvent()
|
||||
| ChunkGenerated()
|
||||
| TaskAcknowledged()
|
||||
| InputChunkReceived()
|
||||
| TracesCollected()
|
||||
| TracesMerged()
|
||||
): # Pass-through events that don't modify state
|
||||
return state
|
||||
case CustomModelCardAdded():
|
||||
return apply_custom_model_card_added(event, state)
|
||||
case CustomModelCardDeleted():
|
||||
return apply_custom_model_card_deleted(event, state)
|
||||
case InstanceCreated():
|
||||
return apply_instance_created(event, state)
|
||||
case InstanceDeleted():
|
||||
@@ -109,8 +83,6 @@ def event_apply(event: Event, state: State) -> State:
|
||||
return apply_runner_status_updated(event, state)
|
||||
case TaskCreated():
|
||||
return apply_task_created(event, state)
|
||||
case TaskDeleted():
|
||||
return apply_task_deleted(event, state)
|
||||
case TaskFailed():
|
||||
return apply_task_failed(event, state)
|
||||
case TaskStatusUpdated():
|
||||
@@ -119,10 +91,6 @@ def event_apply(event: Event, state: State) -> State:
|
||||
return apply_topology_edge_created(event, state)
|
||||
case TopologyEdgeDeleted():
|
||||
return apply_topology_edge_deleted(event, state)
|
||||
case InstanceLinkCreated():
|
||||
return apply_instance_link_created(event, state)
|
||||
case InstanceLinkDeleted():
|
||||
return apply_instance_link_deleted(event, state)
|
||||
|
||||
|
||||
def apply(state: State, event: IndexedEvent) -> State:
|
||||
@@ -173,13 +141,6 @@ def apply_task_created(event: TaskCreated, state: State) -> State:
|
||||
return state.model_copy(update={"tasks": new_tasks})
|
||||
|
||||
|
||||
def apply_task_deleted(event: TaskDeleted, state: State) -> State:
|
||||
new_tasks: Mapping[TaskId, Task] = {
|
||||
tid: task for tid, task in state.tasks.items() if tid != event.task_id
|
||||
}
|
||||
return state.model_copy(update={"tasks": new_tasks})
|
||||
|
||||
|
||||
def apply_task_status_updated(event: TaskStatusUpdated, state: State) -> State:
|
||||
if event.task_id not in state.tasks:
|
||||
# maybe should raise
|
||||
@@ -222,38 +183,7 @@ def apply_instance_deleted(event: InstanceDeleted, state: State) -> State:
|
||||
new_instances: Mapping[InstanceId, Instance] = {
|
||||
iid: inst for iid, inst in state.instances.items() if iid != event.instance_id
|
||||
}
|
||||
new_links: dict[InstanceLinkId, InstanceLink] = {}
|
||||
for link_id, link in state.instance_links.items():
|
||||
prefill = [i for i in link.prefill_instances if i != event.instance_id]
|
||||
decode = [i for i in link.decode_instances if i != event.instance_id]
|
||||
if not prefill or not decode:
|
||||
continue
|
||||
if prefill == list(link.prefill_instances) and decode == list(
|
||||
link.decode_instances
|
||||
):
|
||||
new_links[link_id] = link
|
||||
else:
|
||||
new_links[link_id] = link.model_copy(
|
||||
update={"prefill_instances": prefill, "decode_instances": decode}
|
||||
)
|
||||
return state.model_copy(
|
||||
update={"instances": new_instances, "instance_links": new_links}
|
||||
)
|
||||
|
||||
|
||||
def apply_instance_link_created(event: InstanceLinkCreated, state: State) -> State:
|
||||
new_links: Mapping[InstanceLinkId, InstanceLink] = {
|
||||
**state.instance_links,
|
||||
event.link.link_id: event.link,
|
||||
}
|
||||
return state.model_copy(update={"instance_links": new_links})
|
||||
|
||||
|
||||
def apply_instance_link_deleted(event: InstanceLinkDeleted, state: State) -> State:
|
||||
new_links: Mapping[InstanceLinkId, InstanceLink] = {
|
||||
lid: link for lid, link in state.instance_links.items() if lid != event.link_id
|
||||
}
|
||||
return state.model_copy(update={"instance_links": new_links})
|
||||
return state.model_copy(update={"instances": new_instances})
|
||||
|
||||
|
||||
def apply_runner_status_updated(event: RunnerStatusUpdated, state: State) -> State:
|
||||
@@ -408,59 +338,26 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
|
||||
event.node_id: NodeThunderboltInfo(interfaces=info.idents),
|
||||
}
|
||||
case MacThunderboltConnections():
|
||||
conn_map = {
|
||||
tb_ident.domain_uuid: (nid, tb_ident.rdma_interface)
|
||||
for nid in state.node_thunderbolt
|
||||
for tb_ident in state.node_thunderbolt[nid].interfaces
|
||||
update["node_thunderbolt_connections"] = {
|
||||
**state.node_thunderbolt_connections,
|
||||
event.node_id: info,
|
||||
}
|
||||
source_is_rdma_enabled = _is_rdma_ctl_enabled(
|
||||
event.node_id, state.node_rdma_ctl
|
||||
)
|
||||
as_rdma_conns = [
|
||||
Connection(
|
||||
source=event.node_id,
|
||||
sink=conn_map[tb_conn.sink_uuid][0],
|
||||
edge=RDMAConnection(
|
||||
source_rdma_iface=conn_map[tb_conn.source_uuid][1],
|
||||
sink_rdma_iface=conn_map[tb_conn.sink_uuid][1],
|
||||
),
|
||||
)
|
||||
for tb_conn in info.conns
|
||||
if tb_conn.source_uuid in conn_map
|
||||
if tb_conn.sink_uuid in conn_map
|
||||
if source_is_rdma_enabled
|
||||
and _is_rdma_ctl_enabled(
|
||||
conn_map[tb_conn.sink_uuid][0], state.node_rdma_ctl
|
||||
)
|
||||
]
|
||||
topology.replace_all_out_rdma_connections(event.node_id, as_rdma_conns)
|
||||
case ThunderboltBridgeInfo():
|
||||
new_tb_bridge: dict[NodeId, ThunderboltBridgeStatus] = {
|
||||
**state.node_thunderbolt_bridge,
|
||||
event.node_id: info.status,
|
||||
}
|
||||
update["node_thunderbolt_bridge"] = new_tb_bridge
|
||||
# Only recompute cycles if the enabled status changed
|
||||
old_status = state.node_thunderbolt_bridge.get(event.node_id)
|
||||
old_enabled = old_status.enabled if old_status else False
|
||||
new_enabled = info.status.enabled
|
||||
if old_enabled != new_enabled:
|
||||
update["thunderbolt_bridge_cycles"] = (
|
||||
topology.get_thunderbolt_bridge_cycles(
|
||||
new_tb_bridge, state.node_network
|
||||
)
|
||||
update["thunderbolt_bridge_cycles"] = (
|
||||
topology.get_thunderbolt_bridge_cycles(
|
||||
new_tb_bridge, state.node_network
|
||||
)
|
||||
)
|
||||
case RdmaCtlStatus():
|
||||
update["node_rdma_ctl"] = {
|
||||
**state.node_rdma_ctl,
|
||||
event.node_id: NodeRdmaCtlStatus(enabled=info.enabled),
|
||||
}
|
||||
# If RDMA just got disabled on this node, drop any RDMA edges touching it
|
||||
# so placement / topology consumers cannot pick a disabled node for an
|
||||
# RDMA-backed instance. (Edges will repopulate on the next
|
||||
# MacThunderboltConnections poll once both endpoints are enabled again.)
|
||||
if not info.enabled:
|
||||
topology.remove_all_rdma_connections_touching(event.node_id)
|
||||
case NodeBackends():
|
||||
update["node_backends"] = {
|
||||
**state.node_backends,
|
||||
@@ -471,32 +368,45 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
|
||||
|
||||
|
||||
def apply_topology_edge_created(event: TopologyEdgeCreated, state: State) -> State:
|
||||
topology = copy.deepcopy(state.topology)
|
||||
topology.add_connection(event.conn)
|
||||
return state.model_copy(update={"topology": topology})
|
||||
source_connections = state.node_socket_connections.get(event.conn.source, {})
|
||||
sink_connections = source_connections.get(event.conn.sink, [])
|
||||
|
||||
update = {
|
||||
"node_socket_connections": {
|
||||
**state.node_socket_connections,
|
||||
event.conn.source: {
|
||||
**source_connections,
|
||||
event.conn.sink: sink_connections
|
||||
if event.conn.edge in sink_connections
|
||||
else [*sink_connections, event.conn.edge],
|
||||
},
|
||||
}
|
||||
}
|
||||
return state.model_copy(update=update)
|
||||
|
||||
|
||||
def apply_topology_edge_deleted(event: TopologyEdgeDeleted, state: State) -> State:
|
||||
topology = copy.deepcopy(state.topology)
|
||||
topology.remove_connection(event.conn)
|
||||
# TODO: Clean up removing the reverse connection
|
||||
return state.model_copy(update={"topology": topology})
|
||||
|
||||
|
||||
def apply_custom_model_card_added(event: CustomModelCardAdded, state: State) -> State:
|
||||
new_cards: Mapping[ModelId, ModelCard] = {
|
||||
**state.custom_model_cards,
|
||||
event.model_card.model_id: event.model_card,
|
||||
inner_update = {
|
||||
sink: final_edges
|
||||
for sink, edges in state.node_socket_connections.get(
|
||||
event.conn.source, {}
|
||||
).items()
|
||||
if (
|
||||
final_edges := [
|
||||
edge
|
||||
for edge in edges
|
||||
if (edge != event.conn.edge or sink != event.conn.sink)
|
||||
]
|
||||
)
|
||||
}
|
||||
return state.model_copy(update={"custom_model_cards": new_cards})
|
||||
|
||||
|
||||
def apply_custom_model_card_deleted(
|
||||
event: CustomModelCardDeleted, state: State
|
||||
) -> State:
|
||||
new_cards: Mapping[ModelId, ModelCard] = {
|
||||
model_id: card
|
||||
for model_id, card in state.custom_model_cards.items()
|
||||
if model_id != event.model_id
|
||||
update = {
|
||||
"node_socket_connections": {
|
||||
source: maps
|
||||
for source, maps in {
|
||||
**state.node_socket_connections,
|
||||
event.conn.source: inner_update,
|
||||
}.items()
|
||||
if maps
|
||||
}
|
||||
}
|
||||
return state.model_copy(update={"custom_model_cards": new_cards})
|
||||
return state.model_copy(update=update)
|
||||
@@ -48,6 +48,11 @@ class _CardCache:
|
||||
def get(self, model_id: ModelId) -> "ModelCard | None":
|
||||
return self.cc.get(model_id)
|
||||
|
||||
def list_cached(self) -> list["ModelCard"]:
|
||||
if EXO_ENABLE_IMAGE_MODELS:
|
||||
return list(self.cc.values())
|
||||
return [c for c in self.cc.values() if not _is_image_card(c)]
|
||||
|
||||
async def save(self, card: "ModelCard"):
|
||||
self.cc[card.model_id] = card
|
||||
try:
|
||||
@@ -55,7 +60,7 @@ class _CardCache:
|
||||
except OSError as e:
|
||||
logger.warning(f"failed to save custom model card ({e.strerror})")
|
||||
|
||||
async def pop(self, model_id: ModelId) -> "ModelCard | None":
|
||||
async def delete(self, model_id: ModelId) -> "ModelCard | None":
|
||||
"""Delete a user-added custom model card. Returns True if deleted."""
|
||||
card_path = _custom_cards_dir / (ModelId(model_id).normalize() + ".toml")
|
||||
try:
|
||||
@@ -68,9 +73,7 @@ class _CardCache:
|
||||
async def list_all(self) -> list["ModelCard"]:
|
||||
if len(self.cc) == 0:
|
||||
await self.refresh()
|
||||
if EXO_ENABLE_IMAGE_MODELS:
|
||||
return list(self.cc.values())
|
||||
return [c for c in self.cc.values() if not _is_image_card(c)]
|
||||
return self.list_cached()
|
||||
|
||||
async def _load_cards_from_dir(self, directory: Path, *, is_custom: bool) -> None:
|
||||
"""Load all TOML model cards from a directory into the cache."""
|
||||
|
||||
@@ -1,46 +0,0 @@
|
||||
from exo.shared.apply import apply
|
||||
from exo.shared.models.model_cards import ModelCard, ModelTask
|
||||
from exo.shared.types.backends import Backend
|
||||
from exo.shared.types.common import ModelId
|
||||
from exo.shared.types.events import (
|
||||
CustomModelCardAdded,
|
||||
CustomModelCardDeleted,
|
||||
IndexedEvent,
|
||||
)
|
||||
from exo.shared.types.memory import Memory
|
||||
from exo.shared.types.state import State
|
||||
|
||||
|
||||
def _model_card(model_id: ModelId) -> ModelCard:
|
||||
return ModelCard(
|
||||
model_id=model_id,
|
||||
n_layers=1,
|
||||
storage_size=Memory.from_bytes(1),
|
||||
hidden_size=1,
|
||||
supports_tensor=True,
|
||||
tasks=[ModelTask.TextGeneration],
|
||||
backends=[Backend.MlxMetal],
|
||||
)
|
||||
|
||||
|
||||
def test_custom_model_card_added_is_reduced_into_state() -> None:
|
||||
card = _model_card(ModelId("custom/model"))
|
||||
|
||||
state = apply(
|
||||
State(),
|
||||
IndexedEvent(idx=0, event=CustomModelCardAdded(model_card=card)),
|
||||
)
|
||||
|
||||
assert state.custom_model_cards == {card.model_id: card}
|
||||
|
||||
|
||||
def test_custom_model_card_deleted_removes_card_from_state() -> None:
|
||||
card = _model_card(ModelId("custom/model"))
|
||||
state = State(custom_model_cards={card.model_id: card}, last_event_applied_idx=0)
|
||||
|
||||
state = apply(
|
||||
state,
|
||||
IndexedEvent(idx=1, event=CustomModelCardDeleted(model_id=card.model_id)),
|
||||
)
|
||||
|
||||
assert state.custom_model_cards == {}
|
||||
@@ -1,72 +0,0 @@
|
||||
from exo.shared.apply import (
|
||||
apply_instance_deleted,
|
||||
apply_instance_link_created,
|
||||
apply_instance_link_deleted,
|
||||
)
|
||||
from exo.shared.types.events import (
|
||||
InstanceDeleted,
|
||||
InstanceLinkCreated,
|
||||
InstanceLinkDeleted,
|
||||
)
|
||||
from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
|
||||
from exo.shared.types.state import State
|
||||
from exo.shared.types.worker.instances import InstanceId
|
||||
|
||||
|
||||
def _link(
|
||||
prefill: list[InstanceId],
|
||||
decode: list[InstanceId],
|
||||
link_id: InstanceLinkId | None = None,
|
||||
) -> InstanceLink:
|
||||
return InstanceLink(
|
||||
link_id=link_id or InstanceLinkId(),
|
||||
prefill_instances=prefill,
|
||||
decode_instances=decode,
|
||||
)
|
||||
|
||||
|
||||
def test_create_link() -> None:
|
||||
state = State()
|
||||
link = _link([InstanceId("a")], [InstanceId("b")])
|
||||
new_state = apply_instance_link_created(InstanceLinkCreated(link=link), state)
|
||||
assert new_state.instance_links == {link.link_id: link}
|
||||
|
||||
|
||||
def test_update_replaces_existing_link() -> None:
|
||||
a, b, c = InstanceId("a"), InstanceId("b"), InstanceId("c")
|
||||
link = _link([a], [b])
|
||||
state = State(instance_links={link.link_id: link})
|
||||
|
||||
updated = link.model_copy(update={"decode_instances": [b, c]})
|
||||
new_state = apply_instance_link_created(InstanceLinkCreated(link=updated), state)
|
||||
assert set(new_state.instance_links[link.link_id].decode_instances) == {b, c}
|
||||
|
||||
|
||||
def test_delete_link() -> None:
|
||||
link = _link([InstanceId("a")], [InstanceId("b")])
|
||||
state = State(instance_links={link.link_id: link})
|
||||
|
||||
new_state = apply_instance_link_deleted(
|
||||
InstanceLinkDeleted(link_id=link.link_id), state
|
||||
)
|
||||
assert new_state.instance_links == {}
|
||||
|
||||
|
||||
def test_instance_deleted_strips_from_links() -> None:
|
||||
a, b, c = InstanceId("a"), InstanceId("b"), InstanceId("c")
|
||||
link = _link([a, c], [b])
|
||||
state = State(instance_links={link.link_id: link})
|
||||
|
||||
new_state = apply_instance_deleted(InstanceDeleted(instance_id=a), state)
|
||||
remaining = new_state.instance_links[link.link_id]
|
||||
assert remaining.prefill_instances == [c]
|
||||
assert remaining.decode_instances == [b]
|
||||
|
||||
|
||||
def test_instance_deleted_drops_link_when_role_empties() -> None:
|
||||
a, b = InstanceId("a"), InstanceId("b")
|
||||
link = _link([a], [b])
|
||||
state = State(instance_links={link.link_id: link})
|
||||
|
||||
new_state = apply_instance_deleted(InstanceDeleted(instance_id=a), state)
|
||||
assert link.link_id not in new_state.instance_links
|
||||
@@ -217,7 +217,7 @@ def test_topology_remove_all_rdma_connections_touching_keeps_socket_edges():
|
||||
)
|
||||
)
|
||||
socket_edge = SocketConnection(
|
||||
sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000")
|
||||
sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000"),
|
||||
)
|
||||
topology.add_connection(Connection(source=node_a, sink=node_b, edge=socket_edge))
|
||||
|
||||
|
||||
@@ -1,35 +0,0 @@
|
||||
from exo.shared.types.common import NodeId
|
||||
from exo.shared.types.multiaddr import Multiaddr
|
||||
from exo.shared.types.state import State
|
||||
from exo.shared.types.topology import Connection, SocketConnection
|
||||
|
||||
|
||||
def test_state_serialization_roundtrip() -> None:
|
||||
"""Verify that State → JSON → State round-trip preserves topology."""
|
||||
|
||||
# --- build a simple state ------------------------------------------------
|
||||
node_a = NodeId("node-a")
|
||||
node_b = NodeId("node-b")
|
||||
|
||||
connection = Connection(
|
||||
source=node_a,
|
||||
sink=node_b,
|
||||
edge=SocketConnection(
|
||||
sink_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/10001"),
|
||||
),
|
||||
)
|
||||
|
||||
state = State()
|
||||
state.topology.add_connection(connection)
|
||||
|
||||
json_repr = state.model_dump_json()
|
||||
restored_state = State.model_validate_json(json_repr)
|
||||
|
||||
assert (
|
||||
state.topology.to_snapshot().nodes
|
||||
== restored_state.topology.to_snapshot().nodes
|
||||
)
|
||||
assert set(state.topology.to_snapshot().connections) == set(
|
||||
restored_state.topology.to_snapshot().connections
|
||||
)
|
||||
assert restored_state.model_dump_json() == json_repr
|
||||
@@ -1,10 +1,8 @@
|
||||
import contextlib
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Iterable
|
||||
|
||||
import rustworkx as rx
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
from exo.shared.types.common import NodeId
|
||||
from exo.shared.types.profiling import (
|
||||
@@ -20,15 +18,6 @@ from exo.shared.types.topology import (
|
||||
)
|
||||
|
||||
|
||||
class TopologySnapshot(BaseModel):
|
||||
nodes: Sequence[NodeId]
|
||||
connections: Mapping[
|
||||
NodeId, Mapping[NodeId, Sequence[SocketConnection | RDMAConnection]]
|
||||
]
|
||||
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
|
||||
@dataclass
|
||||
class Topology:
|
||||
_graph: rx.PyDiGraph[NodeId, SocketConnection | RDMAConnection] = field(
|
||||
@@ -36,28 +25,6 @@ class Topology:
|
||||
)
|
||||
_vertex_indices: dict[NodeId, int] = field(init=False, default_factory=dict)
|
||||
|
||||
def to_snapshot(self) -> TopologySnapshot:
|
||||
return TopologySnapshot(
|
||||
nodes=list(self.list_nodes()), connections=self.map_connections()
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_snapshot(cls, snapshot: TopologySnapshot) -> "Topology":
|
||||
topology = cls()
|
||||
|
||||
for node_id in snapshot.nodes:
|
||||
with contextlib.suppress(ValueError):
|
||||
topology.add_node(node_id)
|
||||
|
||||
for source in snapshot.connections:
|
||||
for sink in snapshot.connections[source]:
|
||||
for edge in snapshot.connections[source][sink]:
|
||||
topology.add_connection(
|
||||
Connection(source=source, sink=sink, edge=edge)
|
||||
)
|
||||
|
||||
return topology
|
||||
|
||||
def add_node(self, node_id: NodeId) -> None:
|
||||
if node_id in self._vertex_indices:
|
||||
return
|
||||
|
||||
@@ -13,8 +13,6 @@ from exo.shared.models.model_cards import ModelId
|
||||
from exo.utils.pydantic_ext import TaggedModel
|
||||
from exo.worker.runner.diagnostics import KnownRunnerDiagnostic
|
||||
|
||||
from .common import CommandId
|
||||
|
||||
|
||||
class BaseChunk(TaggedModel):
|
||||
model: ModelId
|
||||
@@ -68,21 +66,6 @@ class ImageChunk(BaseChunk):
|
||||
yield name, value
|
||||
|
||||
|
||||
class InputImageChunk(BaseChunk):
|
||||
command_id: CommandId
|
||||
data: str
|
||||
chunk_index: int
|
||||
total_chunks: int
|
||||
image_index: int = 0
|
||||
|
||||
def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
|
||||
for name, value in super().__repr_args__(): # pyright: ignore[reportAny]
|
||||
if name == "data" and hasattr(value, "__len__"): # pyright: ignore[reportAny]
|
||||
yield name, f"<{len(self.data)} chars>"
|
||||
elif name is not None:
|
||||
yield name, value
|
||||
|
||||
|
||||
class PrefillProgressChunk(BaseChunk):
|
||||
"""Data class for prefill progress events during streaming."""
|
||||
|
||||
|
||||
@@ -5,9 +5,7 @@ from exo.api.types import (
|
||||
ImageGenerationTaskParams,
|
||||
)
|
||||
from exo.shared.models.model_cards import ModelCard, ModelId
|
||||
from exo.shared.types.chunks import InputImageChunk
|
||||
from exo.shared.types.common import CommandId, NodeId, SystemId
|
||||
from exo.shared.types.instance_link import InstanceLinkId
|
||||
from exo.shared.types.text_generation import TextGenerationTaskParams
|
||||
from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
|
||||
from exo.shared.types.worker.shards import Sharding, ShardMetadata
|
||||
@@ -46,7 +44,7 @@ class CreateInstance(BaseCommand):
|
||||
|
||||
|
||||
class DeleteInstance(BaseCommand):
|
||||
instance_id: InstanceId
|
||||
pass
|
||||
|
||||
|
||||
class TaskCancelled(BaseCommand):
|
||||
@@ -57,12 +55,6 @@ class TaskFinished(BaseCommand):
|
||||
finished_command_id: CommandId
|
||||
|
||||
|
||||
class SendInputChunk(BaseCommand):
|
||||
"""Command to send an input image chunk (converted to event by master)."""
|
||||
|
||||
chunk: InputImageChunk
|
||||
|
||||
|
||||
class RequestEventLog(BaseCommand):
|
||||
since_idx: int
|
||||
|
||||
@@ -82,24 +74,6 @@ class CancelDownload(BaseCommand):
|
||||
model_id: ModelId
|
||||
|
||||
|
||||
class AddCustomModelCard(BaseCommand):
|
||||
model_card: ModelCard
|
||||
|
||||
|
||||
class DeleteCustomModelCard(BaseCommand):
|
||||
model_id: ModelId
|
||||
|
||||
|
||||
class SetInstanceLink(BaseCommand):
|
||||
link_id: InstanceLinkId
|
||||
prefill_instances: list[InstanceId]
|
||||
decode_instances: list[InstanceId]
|
||||
|
||||
|
||||
class DeleteInstanceLink(BaseCommand):
|
||||
link_id: InstanceLinkId
|
||||
|
||||
|
||||
DownloadCommand = StartDownload | DeleteDownload | CancelDownload
|
||||
|
||||
|
||||
@@ -114,11 +88,6 @@ Command = (
|
||||
| DeleteInstance
|
||||
| TaskCancelled
|
||||
| TaskFinished
|
||||
| SendInputChunk
|
||||
| AddCustomModelCard
|
||||
| DeleteCustomModelCard
|
||||
| SetInstanceLink
|
||||
| DeleteInstanceLink
|
||||
)
|
||||
|
||||
|
||||
@@ -130,3 +99,15 @@ class ForwarderCommand(FrozenModel):
|
||||
class ForwarderDownloadCommand(FrozenModel):
|
||||
origin: SystemId
|
||||
command: DownloadCommand
|
||||
|
||||
|
||||
class JoinInstance(TaggedModel):
|
||||
# TODO: strip this down to less data
|
||||
instance: Instance
|
||||
|
||||
|
||||
class LeaveInstance(TaggedModel):
|
||||
instance_id: InstanceId
|
||||
|
||||
|
||||
Mail = JoinInstance | LeaveInstance
|
||||
@@ -3,11 +3,9 @@ from typing import final
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from exo.shared.models.model_cards import ModelCard
|
||||
from exo.shared.topology import Connection
|
||||
from exo.shared.types.chunks import Chunk, InputImageChunk
|
||||
from exo.shared.types.common import CommandId, Id, ModelId, NodeId, SessionId, SystemId
|
||||
from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
|
||||
from exo.shared.types.chunks import Chunk
|
||||
from exo.shared.types.common import CommandId, Id, NodeId, SessionId, SystemId
|
||||
from exo.shared.types.tasks import Task, TaskId, TaskStatus
|
||||
from exo.shared.types.worker.downloads import DownloadProgress
|
||||
from exo.shared.types.worker.instances import Instance, InstanceId
|
||||
@@ -41,10 +39,6 @@ class TaskAcknowledged(BaseEvent):
|
||||
task_id: TaskId
|
||||
|
||||
|
||||
class TaskDeleted(BaseEvent):
|
||||
task_id: TaskId
|
||||
|
||||
|
||||
class TaskStatusUpdated(BaseEvent):
|
||||
task_id: TaskId
|
||||
task_status: TaskStatus
|
||||
@@ -95,11 +89,6 @@ class ChunkGenerated(BaseEvent):
|
||||
chunk: Chunk
|
||||
|
||||
|
||||
class InputChunkReceived(BaseEvent):
|
||||
command_id: CommandId
|
||||
chunk: InputImageChunk
|
||||
|
||||
|
||||
class TopologyEdgeCreated(BaseEvent):
|
||||
conn: Connection
|
||||
|
||||
@@ -108,14 +97,6 @@ class TopologyEdgeDeleted(BaseEvent):
|
||||
conn: Connection
|
||||
|
||||
|
||||
class CustomModelCardAdded(BaseEvent):
|
||||
model_card: ModelCard
|
||||
|
||||
|
||||
class CustomModelCardDeleted(BaseEvent):
|
||||
model_id: ModelId
|
||||
|
||||
|
||||
@final
|
||||
class TraceEventData(FrozenModel):
|
||||
name: str
|
||||
@@ -138,20 +119,11 @@ class TracesMerged(BaseEvent):
|
||||
traces: list[TraceEventData]
|
||||
|
||||
|
||||
class InstanceLinkCreated(BaseEvent):
|
||||
link: InstanceLink
|
||||
|
||||
|
||||
class InstanceLinkDeleted(BaseEvent):
|
||||
link_id: InstanceLinkId
|
||||
|
||||
|
||||
Event = (
|
||||
TestEvent
|
||||
| TaskCreated
|
||||
| TaskStatusUpdated
|
||||
| TaskFailed
|
||||
| TaskDeleted
|
||||
| TaskAcknowledged
|
||||
| InstanceCreated
|
||||
| InstanceDeleted
|
||||
@@ -160,15 +132,10 @@ Event = (
|
||||
| NodeGatheredInfo
|
||||
| NodeDownloadProgress
|
||||
| ChunkGenerated
|
||||
| InputChunkReceived
|
||||
| TopologyEdgeCreated
|
||||
| TopologyEdgeDeleted
|
||||
| TracesCollected
|
||||
| TracesMerged
|
||||
| CustomModelCardAdded
|
||||
| CustomModelCardDeleted
|
||||
| InstanceLinkCreated
|
||||
| InstanceLinkDeleted
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -16,6 +16,10 @@ class MemoryUsage(FrozenModel):
|
||||
swap_total: Memory
|
||||
swap_available: Memory
|
||||
|
||||
@classmethod
|
||||
def tag(cls) -> str:
|
||||
return cls.__name__
|
||||
|
||||
@classmethod
|
||||
def from_bytes(
|
||||
cls, *, ram_total: int, ram_available: int, swap_total: int, swap_available: int
|
||||
|
||||
+160
-29
@@ -1,15 +1,22 @@
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from typing import Any, cast
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from pydantic import ConfigDict, Field, field_serializer, field_validator
|
||||
from exo_rs import LVAggregator
|
||||
from pydantic import ConfigDict, Field, TypeAdapter, model_serializer
|
||||
from pydantic.alias_generators import to_camel
|
||||
from pydantic_core.core_schema import SerializerFunctionWrapHandler
|
||||
|
||||
from exo.shared.models.model_cards import ModelCard
|
||||
from exo.shared.topology import Topology, TopologySnapshot
|
||||
from exo.shared.models import model_cards
|
||||
from exo.shared.topology import Topology
|
||||
from exo.shared.types.backends import Backend
|
||||
from exo.shared.types.common import ModelId, NodeId
|
||||
from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
|
||||
from exo.shared.types.common import NodeId
|
||||
from exo.shared.types.events import (
|
||||
InstanceCreated,
|
||||
NodeDownloadProgress,
|
||||
NodeGatheredInfo,
|
||||
RunnerStatusUpdated,
|
||||
)
|
||||
from exo.shared.types.profiling import (
|
||||
DiskUsage,
|
||||
MemoryUsage,
|
||||
@@ -21,11 +28,28 @@ from exo.shared.types.profiling import (
|
||||
ThunderboltBridgeStatus,
|
||||
)
|
||||
from exo.shared.types.tasks import Task, TaskId
|
||||
from exo.shared.types.worker.downloads import DownloadProgress
|
||||
from exo.shared.types.topology import (
|
||||
Connection,
|
||||
RDMAConnection,
|
||||
SocketConnection,
|
||||
SocketConnections,
|
||||
)
|
||||
from exo.shared.types.worker.downloads import DownloadPending, DownloadProgress
|
||||
from exo.shared.types.worker.instances import Instance, InstanceId
|
||||
from exo.shared.types.worker.runners import RunnerId, RunnerStatus
|
||||
from exo.shared.types.worker.shards import PipelineShardMetadata
|
||||
from exo.utils.info_gatherer.info_gatherer import (
|
||||
GatheredInfo,
|
||||
MacThunderboltConnections,
|
||||
)
|
||||
from exo.utils.pydantic_ext import FrozenModel
|
||||
|
||||
_DOWNLOAD_PROGRESS_ADAPTER = TypeAdapter[DownloadProgress](DownloadProgress)
|
||||
_INSTANCE_ADAPTER = TypeAdapter[Instance](Instance)
|
||||
_RUNNER_STATUS_ADAPTER = TypeAdapter[RunnerStatus](RunnerStatus)
|
||||
_SOCKET_CONNECTIONS_ADAPTER = TypeAdapter[SocketConnections](SocketConnections)
|
||||
_GATHERED_INFO_ADAPTER = TypeAdapter[GatheredInfo](GatheredInfo)
|
||||
|
||||
|
||||
class State(FrozenModel):
|
||||
"""Global system state.
|
||||
@@ -39,7 +63,6 @@ class State(FrozenModel):
|
||||
alias_generator=to_camel,
|
||||
validate_by_name=True,
|
||||
extra="forbid",
|
||||
# I want to reenable this ASAP, but it's causing an issue with TaskStatus
|
||||
strict=True,
|
||||
arbitrary_types_allowed=True,
|
||||
)
|
||||
@@ -48,7 +71,6 @@ class State(FrozenModel):
|
||||
downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {}
|
||||
tasks: Mapping[TaskId, Task] = {}
|
||||
last_seen: Mapping[NodeId, datetime] = {}
|
||||
topology: Topology = Field(default_factory=Topology)
|
||||
last_event_applied_idx: int = Field(default=-1, ge=-1)
|
||||
|
||||
# Granular node state mappings (update independently at different frequencies)
|
||||
@@ -61,34 +83,143 @@ class State(FrozenModel):
|
||||
node_thunderbolt_bridge: Mapping[NodeId, ThunderboltBridgeStatus] = {}
|
||||
node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus] = {}
|
||||
node_backends: Mapping[NodeId, list[Backend]] = {}
|
||||
node_socket_connections: Mapping[
|
||||
NodeId, Mapping[NodeId, Sequence[SocketConnection]]
|
||||
] = {}
|
||||
node_thunderbolt_connections: Mapping[NodeId, MacThunderboltConnections] = {}
|
||||
|
||||
# Detected cycles where all nodes have Thunderbolt bridge enabled (>2 nodes)
|
||||
thunderbolt_bridge_cycles: Sequence[Sequence[NodeId]] = []
|
||||
|
||||
instance_links: Mapping[InstanceLinkId, InstanceLink] = {}
|
||||
prefill_server_ports: Mapping[RunnerId, int] = {}
|
||||
|
||||
# User-added model cards. Workers can reconcile their on-disk custom card cache
|
||||
custom_model_cards: Mapping[ModelId, ModelCard] = {}
|
||||
@model_serializer(mode="wrap")
|
||||
def _serialize(self, handler: SerializerFunctionWrapHandler) -> dict[str, Any]:
|
||||
data = handler(self) # pyright: ignore[reportAny]
|
||||
data["topology"] = {
|
||||
"nodes": list(self.node_identities.keys()),
|
||||
"connections": self.topology.map_connections(),
|
||||
}
|
||||
return data # pyright: ignore[reportAny]
|
||||
|
||||
@field_serializer("topology", mode="plain")
|
||||
def _encode_topology(self, value: Topology) -> TopologySnapshot:
|
||||
return value.to_snapshot()
|
||||
@property
|
||||
def topology(self) -> Topology:
|
||||
topology = Topology()
|
||||
thunderbolt_by_uuid = {
|
||||
ident.domain_uuid: (node_id, ident.rdma_interface)
|
||||
for node_id, info in self.node_thunderbolt.items()
|
||||
for ident in info.interfaces
|
||||
}
|
||||
for node_id in self.node_identities:
|
||||
topology.add_node(node_id)
|
||||
|
||||
@field_validator("topology", mode="before")
|
||||
@classmethod
|
||||
def _deserialize_topology(cls, value: object) -> Topology: # noqa: D401 – Pydantic validator signature
|
||||
"""Convert an incoming *value* into a :class:`Topology` instance.
|
||||
for source, data in self.node_socket_connections.items():
|
||||
for sink, conns in data.items():
|
||||
for conn in conns:
|
||||
topology.add_connection(
|
||||
Connection(source=source, sink=sink, edge=conn)
|
||||
)
|
||||
|
||||
Accepts either an already constructed :class:`Topology` or a mapping
|
||||
representing :class:`~shared.topology.TopologySnapshot`.
|
||||
"""
|
||||
for source, connections in self.node_thunderbolt_connections.items():
|
||||
if not self.node_rdma_ctl.get(
|
||||
source, NodeRdmaCtlStatus(enabled=False)
|
||||
).enabled:
|
||||
continue
|
||||
for connection in connections.conns:
|
||||
if (
|
||||
source_iface := thunderbolt_by_uuid.get(connection.source_uuid)
|
||||
) is None or (
|
||||
sink_iface := thunderbolt_by_uuid.get(connection.sink_uuid)
|
||||
) is None:
|
||||
continue
|
||||
if not self.node_rdma_ctl.get(
|
||||
sink_iface[0], NodeRdmaCtlStatus(enabled=False)
|
||||
).enabled:
|
||||
continue
|
||||
assert source_iface[0] == source, "registered invalid source uuid"
|
||||
topology.add_connection(
|
||||
Connection(
|
||||
source=source_iface[0],
|
||||
sink=sink_iface[0],
|
||||
edge=RDMAConnection(
|
||||
source_rdma_iface=source_iface[1],
|
||||
sink_rdma_iface=sink_iface[1],
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
if isinstance(value, Topology):
|
||||
return value
|
||||
return topology
|
||||
|
||||
if isinstance(value, Mapping): # likely a snapshot-dict coming from JSON
|
||||
snapshot = TopologySnapshot(**cast(dict[str, Any], value)) # type: ignore[arg-type]
|
||||
return Topology.from_snapshot(snapshot)
|
||||
def with_aggregator(self, aggregator: LVAggregator) -> "State":
|
||||
from exo.shared.apply import event_apply
|
||||
|
||||
raise TypeError("Invalid representation for Topology field in State")
|
||||
state = self.model_copy()
|
||||
values = aggregator.dump()
|
||||
node_ids = {NodeId(key.split("/")[0]) for key in values}
|
||||
cached_cards = model_cards.card_cache.list_cached()
|
||||
if cached_cards and node_ids:
|
||||
state = state.model_copy(
|
||||
update={
|
||||
"downloads": {
|
||||
node_id: [
|
||||
DownloadPending(
|
||||
node_id=node_id,
|
||||
shard_metadata=PipelineShardMetadata(
|
||||
model_card=card,
|
||||
device_rank=0,
|
||||
world_size=1,
|
||||
start_layer=0,
|
||||
end_layer=card.n_layers,
|
||||
n_layers=card.n_layers,
|
||||
),
|
||||
total=card.storage_size,
|
||||
)
|
||||
for card in cached_cards
|
||||
]
|
||||
for node_id in node_ids
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
for key, value in values.items():
|
||||
try:
|
||||
parts = key.split("/")
|
||||
if len(parts) >= 3 and parts[1] == "downloads":
|
||||
progress = _DOWNLOAD_PROGRESS_ADAPTER.validate_json(value)
|
||||
event = NodeDownloadProgress(download_progress=progress)
|
||||
elif len(parts) == 3 and parts[1] == "desired_instances":
|
||||
instance = _INSTANCE_ADAPTER.validate_json(value)
|
||||
event = InstanceCreated(instance=instance)
|
||||
elif len(parts) == 4 and parts[1] == "runners":
|
||||
runner_status = _RUNNER_STATUS_ADAPTER.validate_json(value)
|
||||
event = RunnerStatusUpdated(
|
||||
runner_status=runner_status, runner_id=RunnerId(parts[2])
|
||||
)
|
||||
elif len(parts) == 2 and parts[1] == "socket_connections":
|
||||
socket_connections = _SOCKET_CONNECTIONS_ADAPTER.validate_json(
|
||||
value
|
||||
)
|
||||
state = state.model_copy(
|
||||
update={
|
||||
"node_socket_connections": {
|
||||
**state.node_socket_connections,
|
||||
NodeId(parts[0]): socket_connections.connections,
|
||||
}
|
||||
}
|
||||
)
|
||||
continue
|
||||
else:
|
||||
data = _GATHERED_INFO_ADAPTER.validate_json(value)
|
||||
node_id = NodeId(parts[0])
|
||||
event = NodeGatheredInfo(
|
||||
node_id=node_id,
|
||||
when=str(datetime.now(tz=timezone.utc)),
|
||||
info=data,
|
||||
)
|
||||
state = event_apply(event, state)
|
||||
except Exception as e:
|
||||
print(
|
||||
f"\n{'=' * 10}key: {key} with exception {str(e)}\nvalue: {value}{'=' * 10}\n"
|
||||
)
|
||||
|
||||
return state
|
||||
@@ -89,6 +89,10 @@ class Shutdown(BaseTask): # emitted by Worker
|
||||
runner_id: RunnerId
|
||||
|
||||
|
||||
class ForgetInstance(BaseTask): # emitted by Worker
|
||||
pass
|
||||
|
||||
|
||||
Task = (
|
||||
CreateRunner
|
||||
| DownloadModel
|
||||
@@ -100,6 +104,7 @@ Task = (
|
||||
| ImageGeneration
|
||||
| ImageEdits
|
||||
| Shutdown
|
||||
| ForgetInstance
|
||||
)
|
||||
TextTask = TextGeneration
|
||||
ImageTask = ImageGeneration | ImageEdits
|
||||
|
||||
@@ -69,10 +69,6 @@ class Base64Image(TruncatingString):
|
||||
truncate_length = 10
|
||||
|
||||
|
||||
class Base64ImageHash(TruncatingString):
|
||||
truncate_length = 10
|
||||
|
||||
|
||||
def _wrap_chat_value(x: Any) -> Any: # pyright: ignore[reportAny]
|
||||
if isinstance(x, (InputMessageContent, Base64Image)):
|
||||
return x
|
||||
@@ -130,7 +126,6 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
|
||||
presence_penalty: float | None = None
|
||||
frequency_penalty: float | None = None
|
||||
images: list[Base64Image] = Field(default_factory=list)
|
||||
image_hashes: dict[int, Base64ImageHash] = Field(default_factory=dict)
|
||||
|
||||
prefill_endpoint: str | None = None
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
from collections.abc import Iterator
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
|
||||
from exo.shared.types.common import NodeId
|
||||
@@ -29,6 +29,10 @@ class SocketConnection(FrozenModel):
|
||||
return hash(self.sink_multiaddr.ip_address)
|
||||
|
||||
|
||||
class SocketConnections(FrozenModel):
|
||||
connections: Mapping[NodeId, Sequence[SocketConnection]]
|
||||
|
||||
|
||||
class Connection(FrozenModel):
|
||||
source: NodeId
|
||||
sink: NodeId
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
from collections.abc import Iterable
|
||||
from enum import Enum
|
||||
|
||||
from pydantic import model_validator
|
||||
@@ -22,7 +23,19 @@ class BaseInstance(TaggedModel):
|
||||
shard_assignments: ShardAssignments
|
||||
|
||||
def shard(self, runner_id: RunnerId) -> ShardMetadata | None:
|
||||
return self.shard_assignments.runner_to_shard.get(runner_id, None)
|
||||
for _, rid, shard in self.shard_assignments.shards:
|
||||
if rid == runner_id:
|
||||
return shard
|
||||
|
||||
def runners_for(self, node_id: NodeId) -> Iterable[RunnerId]:
|
||||
for nid, rid, _ in self.shard_assignments.shards:
|
||||
if nid == node_id:
|
||||
yield rid
|
||||
|
||||
def primary_output_node(self) -> NodeId:
|
||||
return self.shard_assignments.shards[
|
||||
self.shard_assignments.primary_output_node
|
||||
].node_id
|
||||
|
||||
|
||||
class MlxRingInstance(BaseInstance):
|
||||
@@ -44,6 +57,12 @@ class BoundInstance(FrozenModel):
|
||||
bound_runner_id: RunnerId
|
||||
bound_node_id: NodeId
|
||||
|
||||
def is_primary_output_node(self) -> bool:
|
||||
return (
|
||||
self.instance.shard_assignments.primary_output_node
|
||||
== self.bound_shard.device_rank
|
||||
)
|
||||
|
||||
@property
|
||||
def bound_shard(self) -> ShardMetadata:
|
||||
shard = self.instance.shard(self.bound_runner_id)
|
||||
@@ -59,8 +78,9 @@ class BoundInstance(FrozenModel):
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_shard_exists(self) -> "BoundInstance":
|
||||
assert (
|
||||
self.bound_runner_id in self.instance.shard_assignments.runner_to_shard
|
||||
assert any(
|
||||
rid == self.bound_runner_id
|
||||
for (_, rid, _) in self.instance.shard_assignments.shards
|
||||
), (
|
||||
"Bound Instance must be constructed with a runner_id that is in the instances assigned shards"
|
||||
)
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Sequence
|
||||
from typing import NamedTuple
|
||||
|
||||
from pydantic import model_validator
|
||||
|
||||
@@ -18,8 +19,7 @@ class RunnerError(Exception):
|
||||
|
||||
|
||||
class BaseRunnerStatus(TaggedModel):
|
||||
def is_running(self):
|
||||
return isinstance(self, RunnerRunning)
|
||||
pass
|
||||
|
||||
|
||||
class RunnerIdle(BaseRunnerStatus):
|
||||
@@ -83,16 +83,32 @@ RunnerStatus = (
|
||||
)
|
||||
|
||||
|
||||
class ShardWithId(NamedTuple):
|
||||
node_id: NodeId
|
||||
runner_id: RunnerId
|
||||
shard: ShardMetadata
|
||||
|
||||
|
||||
class ShardAssignments(FrozenModel):
|
||||
model_id: ModelId
|
||||
runner_to_shard: Mapping[RunnerId, ShardMetadata]
|
||||
node_to_runner: Mapping[NodeId, RunnerId]
|
||||
shards: Sequence[ShardWithId]
|
||||
# this node needs to be connected to the API node for the stream to be considered ready
|
||||
# (this is a device rank)
|
||||
primary_output_node: int
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_runners_exist(self) -> "ShardAssignments":
|
||||
for runner_id in self.node_to_runner.values():
|
||||
if runner_id not in self.runner_to_shard:
|
||||
raise ValueError(
|
||||
f"Runner {runner_id} in node_to_runner does not exist in runner_to_shard"
|
||||
)
|
||||
for position, shard in enumerate(self.shards):
|
||||
if shard.shard.device_rank != position:
|
||||
raise ValueError("shard position does not correspond to device rank")
|
||||
|
||||
if not self.shards[self.primary_output_node].shard.is_primary_output():
|
||||
raise ValueError("primary output node does not correspond to primary shard")
|
||||
|
||||
return self
|
||||
|
||||
def runner_ids(self) -> list[RunnerId]:
|
||||
return [shard.runner_id for shard in self.shards]
|
||||
|
||||
def node_ids(self) -> list[NodeId]:
|
||||
return [shard.node_id for shard in self.shards]
|
||||
@@ -1,5 +1,5 @@
|
||||
from enum import Enum
|
||||
from typing import TypeAlias, final
|
||||
from typing import final
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
@@ -15,18 +15,14 @@ class Sharding(str, Enum):
|
||||
class BaseShardMetadata(TaggedModel):
|
||||
"""
|
||||
Defines a specific shard of the model that is ready to be run on a device.
|
||||
Replaces previous `Shard` object.
|
||||
Layers are represented as a half-open interval [start_layer, end_layer),
|
||||
where start_layer is inclusive and end_layer is exclusive.
|
||||
"""
|
||||
|
||||
model_card: ModelCard
|
||||
device_rank: int
|
||||
world_size: int
|
||||
|
||||
# Error handling; equivalent to monkey-patch, but we can't monkey-patch runner.py
|
||||
# This is kinda annoying because it allocates memory in the ShardMetadata object. Can be rethought after Shanghai.
|
||||
immediate_exception: bool = False
|
||||
should_timeout: float | None = None
|
||||
|
||||
start_layer: int = Field(ge=0)
|
||||
end_layer: int = Field(ge=0)
|
||||
n_layers: int = Field(ge=0)
|
||||
@@ -51,27 +47,59 @@ class BaseShardMetadata(TaggedModel):
|
||||
)
|
||||
)
|
||||
|
||||
def is_primary_output(self) -> bool:
|
||||
return self.device_rank == self.world_size - 1
|
||||
|
||||
def is_primary_output_node(self) -> bool:
|
||||
return self.is_primary_output()
|
||||
|
||||
|
||||
@final
|
||||
class PipelineShardMetadata(BaseShardMetadata):
|
||||
"""
|
||||
Pipeline parallelism shard meta.
|
||||
|
||||
Layers are represented as a half-open interval [start_layer, end_layer),
|
||||
where start_layer is inclusive and end_layer is exclusive.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
@final
|
||||
class CfgShardMetadata(BaseShardMetadata):
|
||||
"""Shard metadata for CFG-parallel image generation models."""
|
||||
# example
|
||||
# world_size 6
|
||||
# rank prank crank
|
||||
# 0 0 0
|
||||
# 1 1 0
|
||||
# 2 2 0
|
||||
# 3 2 1
|
||||
# 4 1 1
|
||||
# 5 0 1
|
||||
|
||||
cfg_rank: int # 0 = positive branch, 1 = negative branch
|
||||
cfg_world_size: int = 2
|
||||
@property
|
||||
def cfg_rank(self) -> int:
|
||||
# 0 = positive branch, 1 = negative branch
|
||||
return 0 if self.device_rank < self.world_size // 2 else 1
|
||||
|
||||
# Pipeline-relative coordinates (computed at placement time)
|
||||
pipeline_rank: int # rank within the pipeline group (0, 1, 2, ...)
|
||||
pipeline_world_size: int # number of nodes per pipeline group
|
||||
@property
|
||||
def cfg_world_size(self) -> int:
|
||||
return 2
|
||||
|
||||
@property
|
||||
def pipeline_rank(self) -> int:
|
||||
return (
|
||||
self.device_rank
|
||||
if self.cfg_rank == 0
|
||||
else (self.world_size - self.device_rank - 1)
|
||||
)
|
||||
|
||||
@property
|
||||
def pipeline_world_size(self) -> int:
|
||||
return self.world_size // 2
|
||||
|
||||
def is_primary_output(self) -> bool:
|
||||
"""
|
||||
For CFG models: the last pipeline stage in CFG group 0 (positive prompt).
|
||||
For non-CFG models: the last pipeline stage.
|
||||
"""
|
||||
assert self.pipeline_world_size == self.world_size // 2
|
||||
assert self.world_size % 2 == 0
|
||||
return self.device_rank == (self.world_size // 2) - 1
|
||||
|
||||
|
||||
@final
|
||||
@@ -79,6 +107,4 @@ class TensorShardMetadata(BaseShardMetadata):
|
||||
pass
|
||||
|
||||
|
||||
ShardMetadata: TypeAlias = (
|
||||
PipelineShardMetadata | CfgShardMetadata | TensorShardMetadata
|
||||
)
|
||||
ShardMetadata = PipelineShardMetadata | CfgShardMetadata | TensorShardMetadata
|
||||
@@ -1,4 +1,4 @@
|
||||
from typing import Any, Type
|
||||
from typing import Any, Callable, Iterable, Iterator, Type, TypeGuard
|
||||
|
||||
from .phantom import PhantomData
|
||||
|
||||
@@ -19,3 +19,11 @@ def todo[T](
|
||||
_phantom: PhantomData[T] = None,
|
||||
) -> T:
|
||||
raise NotImplementedError(msg)
|
||||
|
||||
|
||||
def not_none[T](t: T | None) -> TypeGuard[T]:
|
||||
return t is not None
|
||||
|
||||
|
||||
def fmap[T, U](f: Callable[[T], U | None], s: Iterable[T]) -> Iterator[U]:
|
||||
return filter(not_none, map(f, s))
|
||||
@@ -219,7 +219,7 @@ class AsyncProcess:
|
||||
await self.wait()
|
||||
|
||||
if self.exitcode is not None or not process.is_alive():
|
||||
logger.warning("Terminated nicely in the first attempt!")
|
||||
logger.info("Terminated nicely in the first attempt!")
|
||||
return
|
||||
|
||||
for attempt in range(2, _TERMINATE_ATTEMPTS + 1):
|
||||
|
||||
@@ -10,11 +10,13 @@ from typing import Self, cast
|
||||
import anyio
|
||||
from anyio import fail_after, open_process, to_thread
|
||||
from anyio.streams.buffered import BufferedByteReceiveStream
|
||||
from exo_rs import LVPublisher, SessionHandle
|
||||
from loguru import logger
|
||||
from pydantic import ValidationError
|
||||
|
||||
from exo.shared.constants import EXO_CONFIG_FILE, EXO_DEFAULT_MODELS_DIR
|
||||
from exo.shared.types.backends import Backend
|
||||
from exo.shared.types.common import NodeId
|
||||
from exo.shared.types.memory import Memory
|
||||
from exo.shared.types.profiling import (
|
||||
DiskUsage,
|
||||
@@ -27,7 +29,6 @@ from exo.shared.types.thunderbolt import (
|
||||
ThunderboltConnectivity,
|
||||
ThunderboltIdentifier,
|
||||
)
|
||||
from exo.utils.channels import Sender
|
||||
from exo.utils.pydantic_ext import TaggedModel
|
||||
from exo.utils.task_group import TaskGroup
|
||||
|
||||
@@ -401,10 +402,42 @@ GatheredInfo = (
|
||||
|
||||
@dataclass
|
||||
class InfoGatherer:
|
||||
info_sender: Sender[GatheredInfo]
|
||||
session_handle: SessionHandle
|
||||
node_id: NodeId
|
||||
info_senders: dict[str, LVPublisher] = field(init=False, default_factory=dict)
|
||||
_tg: TaskGroup = field(init=False, default_factory=TaskGroup)
|
||||
_psutil_enabled: bool = field(init=False, default=False)
|
||||
|
||||
async def send(self, info: GatheredInfo):
|
||||
if (tag := info.tag()) not in self.info_senders:
|
||||
self.info_senders[tag] = self.session_handle.last_value_publisher(
|
||||
f"node_metrics/{self.node_id}/{tag}"
|
||||
)
|
||||
await self.info_senders[tag].put(info.model_dump_json())
|
||||
|
||||
async def run(self):
|
||||
async with self._tg as tg:
|
||||
if IS_DARWIN:
|
||||
tg.start_soon(self._monitor_macmon, 1)
|
||||
tg.start_soon(self._monitor_system_profiler_thunderbolt_data, 5)
|
||||
tg.start_soon(self._monitor_thunderbolt_bridge_status, 10)
|
||||
tg.start_soon(self._monitor_rdma_ctl_status, 10)
|
||||
if not IS_DARWIN:
|
||||
tg.start_soon(self._monitor_memory_usage, 1)
|
||||
tg.start_soon(self._watch_system_info, 10)
|
||||
tg.start_soon(self._monitor_misc, 60)
|
||||
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.send(nc)
|
||||
|
||||
await self.send(await NodeBackends.gather())
|
||||
|
||||
def shutdown(self):
|
||||
self._tg.cancel_tasks()
|
||||
|
||||
async def _can_read_macmon_metrics(self, macmon_path: str) -> bool:
|
||||
try:
|
||||
with fail_after(5):
|
||||
@@ -441,34 +474,11 @@ class InfoGatherer:
|
||||
|
||||
return True
|
||||
|
||||
async def run(self):
|
||||
async with self._tg as tg:
|
||||
if IS_DARWIN:
|
||||
tg.start_soon(self._monitor_macmon, 1)
|
||||
tg.start_soon(self._monitor_system_profiler_thunderbolt_data, 5)
|
||||
tg.start_soon(self._monitor_thunderbolt_bridge_status, 10)
|
||||
tg.start_soon(self._monitor_rdma_ctl_status, 10)
|
||||
if not IS_DARWIN:
|
||||
tg.start_soon(self._monitor_memory_usage, 1)
|
||||
tg.start_soon(self._watch_system_info, 10)
|
||||
tg.start_soon(self._monitor_misc, 60)
|
||||
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):
|
||||
self._tg.cancel_tasks()
|
||||
|
||||
async def _monitor_static_info(self, static_info_poll_interval: float):
|
||||
while True:
|
||||
try:
|
||||
with fail_after(30):
|
||||
await self.info_sender.send(await StaticNodeInformation.gather())
|
||||
await self.send(await StaticNodeInformation.gather())
|
||||
except Exception as e:
|
||||
logger.opt(exception=e).warning("Error gathering static node info")
|
||||
await anyio.sleep(static_info_poll_interval)
|
||||
@@ -477,7 +487,7 @@ class InfoGatherer:
|
||||
while True:
|
||||
try:
|
||||
with fail_after(10):
|
||||
await self.info_sender.send(await MiscData.gather())
|
||||
await self.send(await MiscData.gather())
|
||||
except Exception as e:
|
||||
logger.opt(exception=e).warning("Error gathering misc data")
|
||||
await anyio.sleep(misc_poll_interval)
|
||||
@@ -498,12 +508,10 @@ class InfoGatherer:
|
||||
idents = [
|
||||
it for i in data if (it := i.ident(iface_map)) is not None
|
||||
]
|
||||
await self.info_sender.send(
|
||||
MacThunderboltIdentifiers(idents=idents)
|
||||
)
|
||||
await self.send(MacThunderboltIdentifiers(idents=idents))
|
||||
|
||||
conns = [it for i in data if (it := i.conn()) is not None]
|
||||
await self.info_sender.send(MacThunderboltConnections(conns=conns))
|
||||
await self.send(MacThunderboltConnections(conns=conns))
|
||||
except Exception as e:
|
||||
logger.opt(exception=e).warning("Error gathering Thunderbolt data")
|
||||
await anyio.sleep(system_profiler_interval)
|
||||
@@ -520,7 +528,7 @@ class InfoGatherer:
|
||||
)
|
||||
while True:
|
||||
try:
|
||||
await self.info_sender.send(
|
||||
await self.send(
|
||||
MemoryUsage.from_psutil(override_memory=override_memory)
|
||||
)
|
||||
except Exception as e:
|
||||
@@ -532,7 +540,7 @@ class InfoGatherer:
|
||||
try:
|
||||
with fail_after(10):
|
||||
nics = await get_network_interfaces()
|
||||
await self.info_sender.send(NodeNetworkInterfaces(ifaces=nics))
|
||||
await self.send(NodeNetworkInterfaces(ifaces=nics))
|
||||
except Exception as e:
|
||||
logger.opt(exception=e).warning("Error gathering network interfaces")
|
||||
await anyio.sleep(interface_watcher_interval)
|
||||
@@ -545,7 +553,7 @@ class InfoGatherer:
|
||||
with fail_after(30):
|
||||
curr = await ThunderboltBridgeInfo.gather()
|
||||
if curr is not None:
|
||||
await self.info_sender.send(curr)
|
||||
await self.send(curr)
|
||||
except Exception as e:
|
||||
logger.opt(exception=e).warning(
|
||||
"Error gathering Thunderbolt Bridge status"
|
||||
@@ -557,7 +565,7 @@ class InfoGatherer:
|
||||
try:
|
||||
curr = await RdmaCtlStatus.gather()
|
||||
if curr is not None:
|
||||
await self.info_sender.send(curr)
|
||||
await self.send(curr)
|
||||
except Exception as e:
|
||||
logger.opt(exception=e).warning("Error gathering RDMA ctl status")
|
||||
await anyio.sleep(rdma_ctl_poll_interval)
|
||||
@@ -566,7 +574,7 @@ class InfoGatherer:
|
||||
while True:
|
||||
try:
|
||||
with fail_after(5):
|
||||
await self.info_sender.send(await NodeDiskUsage.gather())
|
||||
await self.send(await NodeDiskUsage.gather())
|
||||
except Exception as e:
|
||||
logger.opt(exception=e).warning("Error gathering disk usage")
|
||||
await anyio.sleep(disk_poll_interval)
|
||||
@@ -611,7 +619,7 @@ class InfoGatherer:
|
||||
)
|
||||
text = data.decode("utf-8", errors="replace").strip()
|
||||
metrics = MacmonMetrics.from_raw_json(text)
|
||||
await self.info_sender.send(metrics)
|
||||
await self.send(metrics)
|
||||
except TimeoutError:
|
||||
logger.warning(
|
||||
f"MacMon produced no output for {read_timeout}s, restarting"
|
||||
|
||||
@@ -19,18 +19,22 @@ class FrozenModel(BaseModel):
|
||||
|
||||
|
||||
class TaggedModel(FrozenModel):
|
||||
@classmethod
|
||||
def tag(cls) -> str:
|
||||
return cls.__name__
|
||||
|
||||
@model_serializer(mode="wrap")
|
||||
def _serialize(self, handler: SerializerFunctionWrapHandler):
|
||||
inner = handler(self) # pyright: ignore[reportAny]
|
||||
return {self.__class__.__name__: inner}
|
||||
return {self.tag(): inner}
|
||||
|
||||
@model_validator(mode="wrap")
|
||||
@classmethod
|
||||
def _validate(cls, v: Any, handler: ValidatorFunctionWrapHandler) -> Self: # pyright: ignore[reportAny]
|
||||
if isinstance(v, dict) and len(v) == 1 and cls.__name__ in v: # pyright: ignore[reportUnknownArgumentType]
|
||||
return handler(v[cls.__name__]) # pyright: ignore[reportAny]
|
||||
if isinstance(v, dict) and len(v) == 1 and cls.tag() in v: # pyright: ignore[reportUnknownArgumentType]
|
||||
return handler(v[cls.tag()]) # pyright: ignore[reportAny]
|
||||
|
||||
return handler(v) # pyright: ignore[reportAny]
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.__class__.__name__}({super().__str__()})"
|
||||
return f"{self.tag()}({super().__str__()})"
|
||||
@@ -30,8 +30,6 @@ from exo.shared.types.worker.runner_response import (
|
||||
ModelLoadingResponse,
|
||||
)
|
||||
from exo.shared.types.worker.shards import (
|
||||
CfgShardMetadata,
|
||||
PipelineShardMetadata,
|
||||
ShardMetadata,
|
||||
)
|
||||
from exo.utils.channels import MpReceiver, MpSender
|
||||
@@ -49,22 +47,6 @@ from exo.worker.engines.mlx.utils_mlx import (
|
||||
)
|
||||
|
||||
|
||||
def _is_primary_output_node(shard_metadata: ShardMetadata) -> bool:
|
||||
"""Check if this node is the primary output node for image generation.
|
||||
|
||||
For CFG models: the last pipeline stage in CFG group 0 (positive prompt).
|
||||
For non-CFG models: the last pipeline stage.
|
||||
"""
|
||||
if isinstance(shard_metadata, CfgShardMetadata):
|
||||
is_pipeline_last = (
|
||||
shard_metadata.pipeline_rank == shard_metadata.pipeline_world_size - 1
|
||||
)
|
||||
return is_pipeline_last and shard_metadata.cfg_rank == 0
|
||||
elif isinstance(shard_metadata, PipelineShardMetadata):
|
||||
return shard_metadata.device_rank == shard_metadata.world_size - 1
|
||||
return False
|
||||
|
||||
|
||||
def _send_traces_if_enabled(
|
||||
event_sender: MpSender[Event],
|
||||
task_id: TaskId,
|
||||
@@ -171,7 +153,7 @@ class ImageEngine(Engine):
|
||||
resp = next(self.current_gen, None)
|
||||
return (
|
||||
(resp,)
|
||||
if resp is not None and _is_primary_output_node(self.shard_metadata)
|
||||
if resp is not None and self.shard_metadata.is_primary_output()
|
||||
else ()
|
||||
)
|
||||
|
||||
@@ -202,10 +184,10 @@ class ImageEngine(Engine):
|
||||
task=task_params,
|
||||
cancel_checker=cancel_checker,
|
||||
):
|
||||
if _is_primary_output_node(self.shard_metadata):
|
||||
if self.shard_metadata.is_primary_output():
|
||||
yield (task_id, response)
|
||||
except Exception as e:
|
||||
if _is_primary_output_node(self.shard_metadata):
|
||||
if self.shard_metadata.is_primary_output():
|
||||
yield (
|
||||
task_id,
|
||||
ErrorChunk(
|
||||
|
||||
@@ -38,11 +38,13 @@ class MlxBuilder(Builder):
|
||||
tokenizer: TokenizerWrapper | None = None
|
||||
group: mx.distributed.Group | None = None
|
||||
vision_processor: VisionProcessor | None = None
|
||||
is_primary_output_node: bool = False
|
||||
|
||||
def connect(self, bound_instance: BoundInstance) -> None:
|
||||
self.group = initialize_mlx(bound_instance)
|
||||
|
||||
def load(self, bound_instance: BoundInstance) -> Generator[ModelLoadingResponse]:
|
||||
self.is_primary_output_node = bound_instance.is_primary_output_node()
|
||||
(
|
||||
self.inference_model,
|
||||
self.tokenizer,
|
||||
@@ -82,7 +84,6 @@ 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(
|
||||
@@ -92,7 +93,7 @@ class MlxBuilder(Builder):
|
||||
tool_parser=tool_parser,
|
||||
kv_prefix_cache=kv_prefix_cache,
|
||||
model_id=self.model_id,
|
||||
device_rank=device_rank,
|
||||
is_primary_output_node=self.is_primary_output_node,
|
||||
cancel_receiver=self.cancel_receiver,
|
||||
event_sender=self.event_sender,
|
||||
vision_processor=vision_processor,
|
||||
@@ -106,7 +107,7 @@ class MlxBuilder(Builder):
|
||||
tool_parser=tool_parser,
|
||||
kv_prefix_cache=kv_prefix_cache,
|
||||
model_id=self.model_id,
|
||||
device_rank=device_rank,
|
||||
is_primary_output_node=self.is_primary_output_node,
|
||||
cancel_receiver=self.cancel_receiver,
|
||||
event_sender=self.event_sender,
|
||||
vision_processor=vision_processor,
|
||||
|
||||
@@ -154,7 +154,7 @@ def initialize_mlx(
|
||||
# TODO: pass in seed from params
|
||||
mx.random.seed(42)
|
||||
|
||||
assert len(bound_instance.instance.shard_assignments.node_to_runner) > 1, (
|
||||
assert len(bound_instance.instance.shard_assignments.shards) > 1, (
|
||||
"Tried to initialize mlx for a single node instance"
|
||||
)
|
||||
return mlx_distributed_init(bound_instance)
|
||||
|
||||
+158
-158
@@ -1,12 +1,12 @@
|
||||
import hashlib
|
||||
from collections import defaultdict
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import anyio
|
||||
from anyio import fail_after, to_thread
|
||||
from exo_rs import LVAggregator, LVPublisher, Mailbox, SessionHandle
|
||||
from loguru import logger
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from exo.api.types import ImageEditsTaskParams
|
||||
from exo.download.download_utils import is_read_only_model_dir, resolve_existing_model
|
||||
from exo.routing.event_router import (
|
||||
EventRouterBrokenResourceError,
|
||||
@@ -14,26 +14,26 @@ from exo.routing.event_router import (
|
||||
)
|
||||
from exo.shared.apply import apply
|
||||
from exo.shared.constants import EXO_MAX_INSTANCE_RETRIES
|
||||
from exo.shared.models.model_cards import ModelId, card_cache
|
||||
from exo.shared.types.chunks import InputImageChunk
|
||||
from exo.shared.models import model_cards
|
||||
from exo.shared.models.model_cards import ModelCard, ModelId
|
||||
from exo.shared.types.commands import (
|
||||
DeleteInstance,
|
||||
ForwarderCommand,
|
||||
ForwarderDownloadCommand,
|
||||
JoinInstance,
|
||||
LeaveInstance,
|
||||
Mail,
|
||||
StartDownload,
|
||||
)
|
||||
from exo.shared.types.common import CommandId, NodeId, SystemId
|
||||
from exo.shared.types.common import NodeId, SystemId
|
||||
from exo.shared.types.events import (
|
||||
Event,
|
||||
IndexedEvent,
|
||||
InputChunkReceived,
|
||||
InstanceDeleted,
|
||||
NodeDownloadProgress,
|
||||
NodeGatheredInfo,
|
||||
TaskCreated,
|
||||
TaskStatusUpdated,
|
||||
TopologyEdgeCreated,
|
||||
TopologyEdgeDeleted,
|
||||
)
|
||||
from exo.shared.types.multiaddr import Multiaddr
|
||||
from exo.shared.types.state import State
|
||||
@@ -41,19 +41,17 @@ from exo.shared.types.tasks import (
|
||||
CancelTask,
|
||||
CreateRunner,
|
||||
DownloadModel,
|
||||
ImageEdits,
|
||||
ForgetInstance,
|
||||
LoadModel,
|
||||
Shutdown,
|
||||
Task,
|
||||
TaskStatus,
|
||||
TextGeneration,
|
||||
)
|
||||
from exo.shared.types.text_generation import Base64Image, Base64ImageHash
|
||||
from exo.shared.types.topology import Connection, SocketConnection
|
||||
from exo.shared.types.topology import SocketConnection, SocketConnections
|
||||
from exo.shared.types.worker.downloads import DownloadCompleted
|
||||
from exo.shared.types.worker.instances import InstanceId
|
||||
from exo.shared.types.worker.instances import Instance, InstanceId
|
||||
from exo.shared.types.worker.runners import RunnerId
|
||||
from exo.utils.channels import Receiver, Sender, channel
|
||||
from exo.utils.channels import Receiver, Sender
|
||||
from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer
|
||||
from exo.utils.info_gatherer.net_profile import check_reachable
|
||||
from exo.utils.keyed_backoff import KeyedBackoff
|
||||
@@ -61,6 +59,8 @@ from exo.utils.task_group import TaskGroup
|
||||
from exo.worker.plan import plan
|
||||
from exo.worker.runner.supervisor import RunnerSupervisor
|
||||
|
||||
PRIMARY_RUNNER_MISSING_TIMEOUT_SECONDS = 10.0
|
||||
|
||||
|
||||
class Worker:
|
||||
def __init__(
|
||||
@@ -73,6 +73,7 @@ class Worker:
|
||||
# but I think it's the correct way to be thinking about commands
|
||||
command_sender: Sender[ForwarderCommand],
|
||||
download_command_sender: Sender[ForwarderDownloadCommand],
|
||||
session_handle: SessionHandle,
|
||||
api_port: int,
|
||||
):
|
||||
self.node_id: NodeId = node_id
|
||||
@@ -88,31 +89,38 @@ class Worker:
|
||||
|
||||
self._system_id = SystemId()
|
||||
|
||||
# Buffer for input image chunks (for image editing)
|
||||
self.input_chunk_buffer: dict[CommandId, dict[int, InputImageChunk]] = {}
|
||||
self.input_chunk_counts: dict[CommandId, int] = {}
|
||||
self.image_cache: dict[Base64ImageHash, Base64Image] = {}
|
||||
|
||||
self._download_backoff: KeyedBackoff[ModelId] = KeyedBackoff(base=0.5, cap=10.0)
|
||||
self._instance_backoff: KeyedBackoff[InstanceId] = KeyedBackoff(
|
||||
base=0.5, cap=10.0
|
||||
)
|
||||
self._stopped: anyio.Event = anyio.Event()
|
||||
self._sh: SessionHandle = session_handle
|
||||
self.aggregator: LVAggregator = session_handle.last_value_aggregator(
|
||||
"node_metrics"
|
||||
)
|
||||
self.mailbox: Mailbox = session_handle.mailbox(node_id)
|
||||
self.desired_instances: dict[InstanceId, Instance] = {}
|
||||
self._primary_runner_missing_since: dict[InstanceId, float] = {}
|
||||
self._primary_desired_instance_publishers: dict[InstanceId, LVPublisher] = {}
|
||||
self._socket_connections_publisher: LVPublisher = (
|
||||
session_handle.last_value_publisher(
|
||||
f"node_metrics/{self.node_id}/socket_connections"
|
||||
)
|
||||
)
|
||||
|
||||
async def run(self):
|
||||
logger.info("Starting Worker")
|
||||
|
||||
info_send, info_recv = channel[GatheredInfo]()
|
||||
info_gatherer: InfoGatherer = InfoGatherer(info_send)
|
||||
info_gatherer: InfoGatherer = InfoGatherer(self._sh, self.node_id)
|
||||
|
||||
try:
|
||||
async with self._tg as tg:
|
||||
tg.start_soon(info_gatherer.run)
|
||||
tg.start_soon(self._forward_info, info_recv)
|
||||
tg.start_soon(self.plan_step)
|
||||
tg.start_soon(self._event_applier)
|
||||
tg.start_soon(self._poll_connection_updates)
|
||||
tg.start_soon(self._reconcile_custom_cards)
|
||||
tg.start_soon(self._listen_to_mailbox)
|
||||
except* (EventRouterBrokenResourceError, EventRouterClosedResourceError):
|
||||
# Event router has been closed (try-star syntax handles error groups)
|
||||
pass
|
||||
@@ -124,8 +132,63 @@ class Worker:
|
||||
self.download_command_sender.close()
|
||||
for runner in self.runners.values():
|
||||
runner.shutdown()
|
||||
with anyio.CancelScope(shield=True):
|
||||
for publisher in self._primary_desired_instance_publishers.values():
|
||||
await publisher.delete()
|
||||
await self._socket_connections_publisher.delete()
|
||||
self._stopped.set()
|
||||
|
||||
async def _listen_to_mailbox(self):
|
||||
ta = TypeAdapter[Mail](Mail)
|
||||
while (mail := await self.mailbox.recv()) is not None:
|
||||
try:
|
||||
mail = ta.validate_json(mail)
|
||||
except ValidationError:
|
||||
logger.warning(f"discarding corrupt mail {mail}")
|
||||
continue
|
||||
|
||||
match mail:
|
||||
case JoinInstance(instance=instance):
|
||||
self.desired_instances[instance.instance_id] = instance
|
||||
if instance.primary_output_node() == self.node_id:
|
||||
await self._publish_primary_desired_instance(instance)
|
||||
case LeaveInstance(instance_id=instance_id):
|
||||
self.desired_instances.pop(instance_id, None)
|
||||
await self._delete_primary_desired_instance(instance_id)
|
||||
|
||||
async def _publish_primary_desired_instance(self, instance: Instance) -> None:
|
||||
publisher = self._primary_desired_instance_publishers.get(instance.instance_id)
|
||||
if publisher is None:
|
||||
publisher = self._sh.last_value_publisher(
|
||||
f"node_metrics/{self.node_id}/desired_instances/{instance.instance_id}"
|
||||
)
|
||||
self._primary_desired_instance_publishers[instance.instance_id] = publisher
|
||||
await publisher.put(instance.model_dump_json())
|
||||
|
||||
async def _delete_primary_desired_instance(self, instance_id: InstanceId) -> None:
|
||||
publisher = self._primary_desired_instance_publishers.pop(instance_id, None)
|
||||
if publisher is not None:
|
||||
await publisher.delete()
|
||||
|
||||
async def _forget_desired_instance_locally(self, instance_id: InstanceId) -> None:
|
||||
self.desired_instances.pop(instance_id, None)
|
||||
self._primary_runner_missing_since.pop(instance_id, None)
|
||||
await self._delete_primary_desired_instance(instance_id)
|
||||
|
||||
def _update_primary_runner_missing_since(
|
||||
self, live_runner_ids: set[RunnerId]
|
||||
) -> None:
|
||||
now = anyio.current_time()
|
||||
for instance_id, instance in list(self.desired_instances.items()):
|
||||
primary_runner_id = instance.shard_assignments.shards[
|
||||
instance.shard_assignments.primary_output_node
|
||||
].runner_id
|
||||
if primary_runner_id in live_runner_ids:
|
||||
self._primary_runner_missing_since.pop(instance_id, None)
|
||||
continue
|
||||
|
||||
self._primary_runner_missing_since.setdefault(instance_id, now)
|
||||
|
||||
async def _forward_info(self, recv: Receiver[GatheredInfo]):
|
||||
with recv as info_stream:
|
||||
async for info in info_stream:
|
||||
@@ -147,77 +210,68 @@ class Worker:
|
||||
if isinstance(event, InstanceDeleted):
|
||||
self._instance_backoff.reset(event.instance_id)
|
||||
|
||||
# Buffer input image chunks for image editing
|
||||
if isinstance(event, InputChunkReceived):
|
||||
cmd_id = event.command_id
|
||||
if cmd_id not in self.input_chunk_buffer:
|
||||
self.input_chunk_buffer[cmd_id] = {}
|
||||
self.input_chunk_counts[cmd_id] = event.chunk.total_chunks
|
||||
|
||||
self.input_chunk_buffer[cmd_id][event.chunk.chunk_index] = (
|
||||
event.chunk
|
||||
)
|
||||
if (
|
||||
len(self.input_chunk_buffer[cmd_id])
|
||||
== self.input_chunk_counts[cmd_id]
|
||||
):
|
||||
per_image: defaultdict[int, list[InputImageChunk]] = (
|
||||
defaultdict(list)
|
||||
)
|
||||
for chunk in self.input_chunk_buffer[cmd_id].values():
|
||||
per_image[chunk.image_index].append(chunk)
|
||||
for chunks_for_image in per_image.values():
|
||||
sorted_chunks = sorted(
|
||||
chunks_for_image, key=lambda c: c.chunk_index
|
||||
)
|
||||
img = Base64Image("".join(c.data for c in sorted_chunks))
|
||||
self.image_cache[
|
||||
Base64ImageHash(
|
||||
hashlib.sha256(img.encode("ascii")).hexdigest()
|
||||
)
|
||||
] = img
|
||||
|
||||
async def _reconcile_custom_cards(self) -> None:
|
||||
storage = self._sh.storage_interface()
|
||||
while True:
|
||||
await anyio.sleep(1)
|
||||
target = dict(self.state.custom_model_cards)
|
||||
for model_id, card in target.items():
|
||||
if card_cache.get(model_id) == card:
|
||||
await anyio.sleep(10)
|
||||
target: list[ModelId] = []
|
||||
for _, value in (await storage.dump("custom_model_cards/")).items():
|
||||
try:
|
||||
card = ModelCard.model_validate_json(value)
|
||||
except ValidationError:
|
||||
continue
|
||||
await card_cache.save(card)
|
||||
target.append(card.model_id)
|
||||
if model_cards.card_cache.get(card.model_id) == card:
|
||||
continue
|
||||
logger.info(f"Registered new custom model card for {card.model_id}")
|
||||
await model_cards.card_cache.save(card)
|
||||
|
||||
for card in await card_cache.list_all():
|
||||
for card in await model_cards.card_cache.list_all():
|
||||
if card.model_id not in target:
|
||||
await card_cache.pop(card.model_id)
|
||||
await model_cards.card_cache.delete(card.model_id)
|
||||
|
||||
async def plan_step(self):
|
||||
while True:
|
||||
await anyio.sleep(0.1)
|
||||
state = self.state.with_aggregator(self.aggregator)
|
||||
live_runner_ids = set(state.runners) | set(self.runners)
|
||||
now = anyio.current_time()
|
||||
self._update_primary_runner_missing_since(live_runner_ids)
|
||||
task: Task | None = plan(
|
||||
self.node_id,
|
||||
self.runners,
|
||||
self.state.downloads,
|
||||
self.state.instances,
|
||||
self.state.runners,
|
||||
self.state.tasks,
|
||||
self.input_chunk_buffer,
|
||||
self.image_cache,
|
||||
state.downloads, # comes from with_agg
|
||||
self.desired_instances, # comes from mailbox
|
||||
state.runners, # comes from with_agg
|
||||
self._instance_backoff,
|
||||
self._download_backoff,
|
||||
live_runner_ids=live_runner_ids,
|
||||
primary_runner_missing_since=self._primary_runner_missing_since,
|
||||
now=now,
|
||||
primary_runner_missing_timeout_seconds=PRIMARY_RUNNER_MISSING_TIMEOUT_SECONDS,
|
||||
)
|
||||
if task is None:
|
||||
continue
|
||||
|
||||
if isinstance(task, ForgetInstance):
|
||||
logger.warning(
|
||||
f"Instance {task.instance_id} primary runner missing for "
|
||||
f"{PRIMARY_RUNNER_MISSING_TIMEOUT_SECONDS:g}s; forgetting locally"
|
||||
)
|
||||
await self._forget_desired_instance_locally(task.instance_id)
|
||||
continue
|
||||
|
||||
if isinstance(task, CreateRunner):
|
||||
iid = task.instance_id
|
||||
if self._instance_backoff.attempts(iid) >= EXO_MAX_INSTANCE_RETRIES:
|
||||
logger.warning(
|
||||
f"Instance {iid} exceeded {EXO_MAX_INSTANCE_RETRIES} retries, requesting deletion"
|
||||
)
|
||||
await self._forget_desired_instance_locally(iid)
|
||||
await self.command_sender.send(
|
||||
ForwarderCommand(
|
||||
origin=self._system_id,
|
||||
command=DeleteInstance(instance_id=iid),
|
||||
command=DeleteInstance(),
|
||||
)
|
||||
)
|
||||
continue
|
||||
@@ -300,63 +354,6 @@ class Worker:
|
||||
task_id=task.task_id, task_status=TaskStatus.Complete
|
||||
)
|
||||
)
|
||||
case ImageEdits() if task.task_params.total_input_chunks > 0:
|
||||
# Assemble image from chunks and inject into task
|
||||
cmd_id = task.command_id
|
||||
chunks = self.input_chunk_buffer.get(cmd_id, {})
|
||||
assembled = "".join(chunks[i].data for i in range(len(chunks)))
|
||||
logger.info(
|
||||
f"Assembled input image from {len(chunks)} chunks, "
|
||||
f"total size: {len(assembled)} bytes"
|
||||
)
|
||||
# Create modified task with assembled image data
|
||||
modified_task = ImageEdits(
|
||||
task_id=task.task_id,
|
||||
command_id=task.command_id,
|
||||
instance_id=task.instance_id,
|
||||
task_status=task.task_status,
|
||||
task_params=ImageEditsTaskParams(
|
||||
image_data=assembled,
|
||||
total_input_chunks=task.task_params.total_input_chunks,
|
||||
prompt=task.task_params.prompt,
|
||||
model=task.task_params.model,
|
||||
n=task.task_params.n,
|
||||
quality=task.task_params.quality,
|
||||
output_format=task.task_params.output_format,
|
||||
response_format=task.task_params.response_format,
|
||||
size=task.task_params.size,
|
||||
image_strength=task.task_params.image_strength,
|
||||
bench=task.task_params.bench,
|
||||
stream=task.task_params.stream,
|
||||
partial_images=task.task_params.partial_images,
|
||||
advanced_params=task.task_params.advanced_params,
|
||||
),
|
||||
)
|
||||
# Cleanup buffers
|
||||
if cmd_id in self.input_chunk_buffer:
|
||||
del self.input_chunk_buffer[cmd_id]
|
||||
if cmd_id in self.input_chunk_counts:
|
||||
del self.input_chunk_counts[cmd_id]
|
||||
await self._start_runner_task(modified_task)
|
||||
|
||||
case TextGeneration() if task.task_params.image_hashes:
|
||||
cmd_id = task.command_id
|
||||
resolved_images = [
|
||||
self.image_cache[h]
|
||||
for _, h in sorted(task.task_params.image_hashes.items())
|
||||
]
|
||||
modified_task = task.model_copy(
|
||||
update={
|
||||
"task_params": task.task_params.model_copy(
|
||||
update={"images": resolved_images}
|
||||
)
|
||||
}
|
||||
)
|
||||
if cmd_id in self.input_chunk_buffer:
|
||||
del self.input_chunk_buffer[cmd_id]
|
||||
if cmd_id in self.input_chunk_counts:
|
||||
del self.input_chunk_counts[cmd_id]
|
||||
await self._start_runner_task(modified_task)
|
||||
case LoadModel(instance_id=instance_id):
|
||||
if (instance := self.state.instances.get(instance_id)) is not None:
|
||||
model_id = instance.shard_assignments.model_id
|
||||
@@ -371,16 +368,27 @@ class Worker:
|
||||
await self._stopped.wait()
|
||||
|
||||
async def _start_runner_task(self, task: Task):
|
||||
if (instance := self.state.instances.get(task.instance_id)) is not None:
|
||||
await self.runners[
|
||||
instance.shard_assignments.node_to_runner[self.node_id]
|
||||
].start_task(task)
|
||||
if (instance := self.desired_instances.get(task.instance_id)) is not None:
|
||||
for rid in instance.runners_for(self.node_id):
|
||||
await self.runners[rid].start_task(task)
|
||||
|
||||
async def _create_supervisor(self, task: CreateRunner) -> RunnerSupervisor:
|
||||
"""Creates and stores a new AssignedRunner with initial downloading status."""
|
||||
task_responder = (
|
||||
self._sh.task_responder(task.instance_id)
|
||||
if task.bound_instance.is_primary_output_node()
|
||||
else None
|
||||
)
|
||||
runner = await RunnerSupervisor.create(
|
||||
bound_instance=task.bound_instance,
|
||||
event_sender=self.event_sender.clone(),
|
||||
task_assignment_subscriber=self._sh.last_value_subscriber(
|
||||
f"task_assignments/{task.instance_id}/*"
|
||||
),
|
||||
runner_status_publisher=self._sh.last_value_publisher(
|
||||
f"node_metrics/{self.node_id}/runners/{task.bound_instance.bound_runner_id}/status"
|
||||
),
|
||||
task_responder=task_responder,
|
||||
)
|
||||
self.runners[task.bound_instance.bound_runner_id] = runner
|
||||
self._tg.start_soon(runner.run)
|
||||
@@ -388,45 +396,37 @@ class Worker:
|
||||
|
||||
async def _poll_connection_updates(self):
|
||||
while True:
|
||||
edges = set(
|
||||
conn.edge for conn in self.state.topology.out_edges(self.node_id)
|
||||
)
|
||||
state = self.state.with_aggregator(self.aggregator)
|
||||
conns: defaultdict[NodeId, set[str]] = defaultdict(set)
|
||||
async for ip, nid in check_reachable(
|
||||
self.state.topology,
|
||||
state.topology,
|
||||
self.node_id,
|
||||
self.state.node_network,
|
||||
state.node_network,
|
||||
api_port=self.api_port,
|
||||
):
|
||||
if ip in conns[nid]:
|
||||
continue
|
||||
conns[nid].add(ip)
|
||||
edge = SocketConnection(
|
||||
# nonsense multiaddr
|
||||
sink_multiaddr=Multiaddr(address=f"/ip4/{ip}/tcp/{self.api_port}")
|
||||
if "." in ip
|
||||
# nonsense multiaddr
|
||||
else Multiaddr(address=f"/ip6/{ip}/tcp/{self.api_port}"),
|
||||
)
|
||||
if edge not in edges:
|
||||
logger.debug(f"ping discovered {edge=}")
|
||||
await self.event_sender.send(
|
||||
TopologyEdgeCreated(
|
||||
conn=Connection(source=self.node_id, sink=nid, edge=edge)
|
||||
)
|
||||
)
|
||||
|
||||
for conn in self.state.topology.out_edges(self.node_id):
|
||||
if not isinstance(conn.edge, SocketConnection):
|
||||
continue
|
||||
# ignore mDNS discovered connections
|
||||
if conn.edge.sink_multiaddr.port != self.api_port:
|
||||
continue
|
||||
if (
|
||||
conn.sink not in conns
|
||||
or conn.edge.sink_multiaddr.ip_address not in conns[conn.sink]
|
||||
):
|
||||
logger.debug(f"ping failed to discover {conn=}")
|
||||
await self.event_sender.send(TopologyEdgeDeleted(conn=conn))
|
||||
socket_connections = SocketConnections(
|
||||
connections={
|
||||
nid: [
|
||||
SocketConnection(
|
||||
# nonsense multiaddr
|
||||
sink_multiaddr=Multiaddr(
|
||||
address=f"/ip4/{ip}/tcp/{self.api_port}"
|
||||
)
|
||||
if "." in ip
|
||||
# nonsense multiaddr
|
||||
else Multiaddr(address=f"/ip6/{ip}/tcp/{self.api_port}"),
|
||||
)
|
||||
for ip in sorted(ips)
|
||||
]
|
||||
for nid, ips in conns.items()
|
||||
}
|
||||
)
|
||||
await self._socket_connections_publisher.put(
|
||||
socket_connections.model_dump_json()
|
||||
)
|
||||
|
||||
await anyio.sleep(10)
|
||||
+92
-96
@@ -2,24 +2,17 @@
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
|
||||
from exo.shared.types.chunks import InputImageChunk
|
||||
from exo.shared.types.common import CommandId, ModelId, NodeId
|
||||
from exo.shared.types.common import ModelId, NodeId
|
||||
from exo.shared.types.tasks import (
|
||||
CancelTask,
|
||||
ConnectToGroup,
|
||||
CreateRunner,
|
||||
DownloadModel,
|
||||
ImageEdits,
|
||||
ImageGeneration,
|
||||
ForgetInstance,
|
||||
LoadModel,
|
||||
Shutdown,
|
||||
StartWarmup,
|
||||
Task,
|
||||
TaskId,
|
||||
TaskStatus,
|
||||
TextGeneration,
|
||||
)
|
||||
from exo.shared.types.text_generation import Base64Image, Base64ImageHash
|
||||
from exo.shared.types.worker.downloads import (
|
||||
DownloadCompleted,
|
||||
DownloadFailed,
|
||||
@@ -35,11 +28,10 @@ from exo.shared.types.worker.runners import (
|
||||
RunnerIdle,
|
||||
RunnerLoaded,
|
||||
RunnerLoading,
|
||||
RunnerReady,
|
||||
RunnerRunning,
|
||||
RunnerStatus,
|
||||
RunnerWarmingUp,
|
||||
)
|
||||
from exo.utils import fmap
|
||||
from exo.utils.keyed_backoff import KeyedBackoff
|
||||
from exo.worker.runner.supervisor import RunnerSupervisor
|
||||
|
||||
@@ -49,26 +41,44 @@ def plan(
|
||||
# Runners is expected to be FRESH and so should not come from state
|
||||
runners: Mapping[RunnerId, RunnerSupervisor],
|
||||
global_download_status: Mapping[NodeId, Sequence[DownloadProgress]],
|
||||
instances: Mapping[InstanceId, Instance],
|
||||
desired_instances: Mapping[InstanceId, Instance],
|
||||
all_runners: Mapping[RunnerId, RunnerStatus], # all global
|
||||
tasks: Mapping[TaskId, Task],
|
||||
input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]],
|
||||
image_cache: Mapping[Base64ImageHash, Base64Image],
|
||||
instance_backoff: KeyedBackoff[InstanceId],
|
||||
download_backoff: KeyedBackoff[ModelId],
|
||||
live_runner_ids: set[RunnerId] | None = None,
|
||||
primary_runner_missing_since: Mapping[InstanceId, float] | None = None,
|
||||
now: float = 0.0,
|
||||
primary_runner_missing_timeout_seconds: float = 10.0,
|
||||
) -> Task | None:
|
||||
live_runner_ids = set(all_runners) if live_runner_ids is None else live_runner_ids
|
||||
primary_runner_missing_since = (
|
||||
{} if primary_runner_missing_since is None else primary_runner_missing_since
|
||||
)
|
||||
# Python short circuiting OR logic should evaluate these sequentially.
|
||||
return (
|
||||
_cancel_tasks(runners, tasks)
|
||||
or _kill_runner(runners, all_runners, instances)
|
||||
or _create_runner(node_id, runners, all_runners, instances, instance_backoff)
|
||||
_kill_runner(runners, all_runners, desired_instances)
|
||||
or _forget_instance_without_primary_runner(
|
||||
node_id,
|
||||
desired_instances,
|
||||
live_runner_ids,
|
||||
primary_runner_missing_since,
|
||||
now,
|
||||
primary_runner_missing_timeout_seconds,
|
||||
)
|
||||
or _create_runner(
|
||||
node_id,
|
||||
runners,
|
||||
all_runners,
|
||||
desired_instances,
|
||||
live_runner_ids,
|
||||
instance_backoff,
|
||||
)
|
||||
or _model_needs_download(
|
||||
node_id, runners, global_download_status, download_backoff
|
||||
)
|
||||
or _init_distributed_backend(runners, all_runners)
|
||||
or _load_model(runners, all_runners, global_download_status)
|
||||
or _ready_to_warmup(runners, all_runners)
|
||||
or _pending_tasks(runners, tasks, all_runners, input_chunk_buffer, image_cache)
|
||||
)
|
||||
|
||||
|
||||
@@ -88,8 +98,10 @@ def _kill_runner(
|
||||
)
|
||||
|
||||
for (
|
||||
global_runner_id
|
||||
) in runner.bound_instance.instance.shard_assignments.node_to_runner.values():
|
||||
_,
|
||||
global_runner_id,
|
||||
_,
|
||||
) in runner.bound_instance.instance.shard_assignments.shards:
|
||||
if runner_id == global_runner_id:
|
||||
continue
|
||||
|
||||
@@ -100,25 +112,74 @@ def _kill_runner(
|
||||
)
|
||||
|
||||
|
||||
def _local_runner_id_for(instance: Instance, node_id: NodeId) -> RunnerId | None:
|
||||
return next(
|
||||
fmap(
|
||||
lambda it: it.runner_id if it.node_id == node_id else None,
|
||||
instance.shard_assignments.shards,
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def _primary_runner_id_for(instance: Instance) -> RunnerId:
|
||||
return instance.shard_assignments.shards[
|
||||
instance.shard_assignments.primary_output_node
|
||||
].runner_id
|
||||
|
||||
|
||||
def _forget_instance_without_primary_runner(
|
||||
node_id: NodeId,
|
||||
instances: Mapping[InstanceId, Instance],
|
||||
live_runner_ids: set[RunnerId],
|
||||
primary_runner_missing_since: Mapping[InstanceId, float],
|
||||
now: float,
|
||||
primary_runner_missing_timeout_seconds: float,
|
||||
) -> ForgetInstance | None:
|
||||
for instance_id, instance in instances.items():
|
||||
local_runner_id = _local_runner_id_for(instance, node_id)
|
||||
if local_runner_id is None:
|
||||
continue
|
||||
|
||||
primary_runner_id = _primary_runner_id_for(instance)
|
||||
if local_runner_id == primary_runner_id:
|
||||
continue
|
||||
if primary_runner_id in live_runner_ids:
|
||||
continue
|
||||
|
||||
missing_since = primary_runner_missing_since.get(instance_id)
|
||||
if missing_since is None:
|
||||
continue
|
||||
if now - missing_since < primary_runner_missing_timeout_seconds:
|
||||
continue
|
||||
|
||||
return ForgetInstance(instance_id=instance_id)
|
||||
|
||||
|
||||
def _create_runner(
|
||||
node_id: NodeId,
|
||||
runners: Mapping[RunnerId, RunnerSupervisor],
|
||||
all_runners: Mapping[RunnerId, RunnerStatus],
|
||||
instances: Mapping[InstanceId, Instance],
|
||||
live_runner_ids: set[RunnerId],
|
||||
instance_backoff: KeyedBackoff[InstanceId],
|
||||
) -> CreateRunner | None:
|
||||
for instance in instances.values():
|
||||
runner_id = instance.shard_assignments.node_to_runner.get(node_id, None)
|
||||
runner_id = _local_runner_id_for(instance, node_id)
|
||||
if runner_id is None:
|
||||
continue
|
||||
|
||||
if runner_id in runners:
|
||||
continue
|
||||
|
||||
primary_runner_id = _primary_runner_id_for(instance)
|
||||
if runner_id != primary_runner_id and primary_runner_id not in live_runner_ids:
|
||||
continue
|
||||
|
||||
# don't create runners if any other nodes have runners that have failed - wait for them to fix themselves first.
|
||||
instance_has_failed_runner = any(
|
||||
isinstance(all_runners.get(remote_runner_id), RunnerFailed)
|
||||
for remote_runner_id in instance.shard_assignments.node_to_runner.values()
|
||||
for (_, remote_runner_id, _) in instance.shard_assignments.shards
|
||||
if remote_runner_id != runner_id
|
||||
)
|
||||
we_have_failed_before = isinstance(all_runners.get(runner_id), RunnerFailed)
|
||||
@@ -175,7 +236,7 @@ def _init_distributed_backend(
|
||||
instance = runner.bound_instance.instance
|
||||
shard_assignments = instance.shard_assignments
|
||||
|
||||
is_single_node_instance = len(shard_assignments.runner_to_shard) == 1
|
||||
is_single_node_instance = len(shard_assignments.shards) == 1
|
||||
if is_single_node_instance:
|
||||
continue
|
||||
|
||||
@@ -185,7 +246,7 @@ def _init_distributed_backend(
|
||||
all_runners.get(global_runner_id),
|
||||
(RunnerConnecting, RunnerIdle),
|
||||
)
|
||||
for global_runner_id in shard_assignments.runner_to_shard
|
||||
for (_, global_runner_id, _) in shard_assignments.shards
|
||||
)
|
||||
|
||||
if not (runner_is_idle and all_runners_connecting):
|
||||
@@ -205,7 +266,7 @@ def _init_distributed_backend(
|
||||
# Rank = n-1
|
||||
connecting_rank_ready = device_rank == world_size - 1 and all(
|
||||
isinstance(all_runners.get(global_runner_id, None), RunnerConnecting)
|
||||
for global_runner_id in shard_assignments.runner_to_shard
|
||||
for (_, global_runner_id, _) in shard_assignments.shards
|
||||
if global_runner_id != runner_id
|
||||
)
|
||||
|
||||
@@ -233,12 +294,12 @@ def _load_model(
|
||||
and dp.shard_metadata.model_card.model_id == shard_assignments.model_id
|
||||
for dp in global_download_status[nid]
|
||||
)
|
||||
for nid in shard_assignments.node_to_runner
|
||||
for (nid, _, _) in shard_assignments.shards
|
||||
)
|
||||
if not all_local_downloads_complete:
|
||||
continue
|
||||
|
||||
is_single_node_instance = len(instance.shard_assignments.runner_to_shard) == 1
|
||||
is_single_node_instance = len(instance.shard_assignments.shards) == 1
|
||||
if is_single_node_instance and isinstance(runner.status, RunnerIdle):
|
||||
return LoadModel(instance_id=instance.instance_id)
|
||||
|
||||
@@ -249,7 +310,7 @@ def _load_model(
|
||||
all_runners.get(global_runner_id, None),
|
||||
(RunnerConnected, RunnerLoading, RunnerLoaded),
|
||||
)
|
||||
for global_runner_id in shard_assignments.runner_to_shard
|
||||
for (_, global_runner_id, _) in shard_assignments.shards
|
||||
)
|
||||
|
||||
if is_runner_waiting and all_ready_for_model:
|
||||
@@ -281,13 +342,13 @@ def _ready_to_warmup(
|
||||
all_runners.get(global_runner_id, None),
|
||||
(RunnerLoaded, RunnerWarmingUp),
|
||||
)
|
||||
for global_runner_id in shard_assignments.runner_to_shard
|
||||
for (_, global_runner_id, _) in shard_assignments.shards
|
||||
)
|
||||
|
||||
# Rank = 0
|
||||
connecting_rank_ready = device_rank == 0 and all(
|
||||
isinstance(all_runners.get(global_runner_id, None), RunnerWarmingUp)
|
||||
for global_runner_id in shard_assignments.runner_to_shard
|
||||
for (_, global_runner_id, _) in shard_assignments.shards
|
||||
if global_runner_id != runner_id
|
||||
)
|
||||
|
||||
@@ -295,68 +356,3 @@ def _ready_to_warmup(
|
||||
return StartWarmup(instance_id=instance.instance_id)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _pending_tasks(
|
||||
runners: Mapping[RunnerId, RunnerSupervisor],
|
||||
tasks: Mapping[TaskId, Task],
|
||||
all_runners: Mapping[RunnerId, RunnerStatus],
|
||||
input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]],
|
||||
image_cache: Mapping[Base64ImageHash, Base64Image],
|
||||
) -> Task | None:
|
||||
for task in tasks.values():
|
||||
# for now, just forward chat completions
|
||||
# TODO(ciaran): do this better!
|
||||
if not isinstance(task, (TextGeneration, ImageGeneration, ImageEdits)):
|
||||
continue
|
||||
if task.task_status not in (TaskStatus.Pending, TaskStatus.Running):
|
||||
continue
|
||||
|
||||
if isinstance(task, ImageEdits) and task.task_params.total_input_chunks > 0:
|
||||
received = len(input_chunk_buffer.get(task.command_id, {}))
|
||||
if received < task.task_params.total_input_chunks:
|
||||
continue # Wait for all chunks to arrive
|
||||
|
||||
if (
|
||||
isinstance(task, TextGeneration)
|
||||
and task.task_params.image_hashes
|
||||
and not all(
|
||||
h in image_cache for h in task.task_params.image_hashes.values()
|
||||
)
|
||||
):
|
||||
continue # Wait for all images to be assembled into the cache
|
||||
|
||||
for runner in runners.values():
|
||||
if task.instance_id != runner.bound_instance.instance.instance_id:
|
||||
continue
|
||||
|
||||
# the task status _should_ be set to completed by the LAST runner
|
||||
# it is currently set by the first
|
||||
# this is definitely a hack
|
||||
if task.task_id in runner.completed or task.task_id in runner.in_progress:
|
||||
continue
|
||||
|
||||
if isinstance(runner.status, (RunnerReady, RunnerRunning)) and all(
|
||||
isinstance(all_runners[global_runner_id], (RunnerReady, RunnerRunning))
|
||||
for global_runner_id in runner.bound_instance.instance.shard_assignments.runner_to_shard
|
||||
):
|
||||
return task
|
||||
|
||||
|
||||
def _cancel_tasks(
|
||||
runners: Mapping[RunnerId, RunnerSupervisor],
|
||||
tasks: Mapping[TaskId, Task],
|
||||
) -> Task | None:
|
||||
for task in tasks.values():
|
||||
if task.task_status != TaskStatus.Cancelled:
|
||||
continue
|
||||
for runner_id, runner in runners.items():
|
||||
if task.instance_id != runner.bound_instance.instance.instance_id:
|
||||
continue
|
||||
if task.task_id in runner.cancelled:
|
||||
continue
|
||||
return CancelTask(
|
||||
instance_id=task.instance_id,
|
||||
cancelled_task_id=task.task_id,
|
||||
runner_id=runner_id,
|
||||
)
|
||||
@@ -94,7 +94,7 @@ class SequentialGenerator(Engine):
|
||||
kv_prefix_cache: KVPrefixCache | None
|
||||
tool_parser: ToolParser | None
|
||||
model_id: ModelId
|
||||
device_rank: int
|
||||
is_primary_output_node: bool
|
||||
cancel_receiver: MpReceiver[TaskId]
|
||||
event_sender: MpSender[Event]
|
||||
vision_processor: VisionProcessor | None = None
|
||||
@@ -201,7 +201,8 @@ class SequentialGenerator(Engine):
|
||||
|
||||
return filter(
|
||||
lambda chunk: (
|
||||
not isinstance(chunk[1], GenerationChunk) or self.device_rank == 0
|
||||
self.is_primary_output_node
|
||||
or isinstance(chunk[1], (CancelledResponse, FinishedResponse))
|
||||
),
|
||||
itertools.chain(
|
||||
output,
|
||||
@@ -235,7 +236,7 @@ class SequentialGenerator(Engine):
|
||||
self._active = (task, gen, queue, output_generator)
|
||||
|
||||
def _send_error(self, task: TextGeneration, e: Exception) -> None:
|
||||
if self.device_rank == 0:
|
||||
if self.is_primary_output_node:
|
||||
self.event_sender.send(
|
||||
ChunkGenerated(
|
||||
command_id=task.command_id,
|
||||
@@ -252,7 +253,7 @@ class SequentialGenerator(Engine):
|
||||
prompt = apply_chat_template(self.tokenizer, task.task_params)
|
||||
|
||||
def on_prefill_progress(processed: int, total: int) -> None:
|
||||
if self.device_rank == 0:
|
||||
if self.is_primary_output_node:
|
||||
self.event_sender.send(
|
||||
ChunkGenerated(
|
||||
command_id=task.command_id,
|
||||
@@ -325,7 +326,7 @@ class BatchGenerator(Engine):
|
||||
kv_prefix_cache: KVPrefixCache | None
|
||||
tool_parser: ToolParser | None
|
||||
model_id: ModelId
|
||||
device_rank: int
|
||||
is_primary_output_node: bool
|
||||
cancel_receiver: MpReceiver[TaskId]
|
||||
event_sender: MpSender[Event]
|
||||
check_for_cancel_every: int = 50
|
||||
@@ -460,7 +461,8 @@ class BatchGenerator(Engine):
|
||||
|
||||
return filter(
|
||||
lambda chunk: (
|
||||
not isinstance(chunk[1], GenerationChunk) or self.device_rank == 0
|
||||
self.is_primary_output_node
|
||||
or isinstance(chunk[1], (CancelledResponse, FinishedResponse))
|
||||
),
|
||||
itertools.chain(output, self._apply_cancellations()),
|
||||
)
|
||||
@@ -494,7 +496,7 @@ class BatchGenerator(Engine):
|
||||
return iter(results)
|
||||
|
||||
def _send_error(self, task: TextGeneration, e: Exception) -> None:
|
||||
if self.device_rank == 0:
|
||||
if self.is_primary_output_node:
|
||||
self.event_sender.send(
|
||||
ChunkGenerated(
|
||||
command_id=task.command_id,
|
||||
@@ -511,7 +513,7 @@ class BatchGenerator(Engine):
|
||||
prompt = apply_chat_template(self.tokenizer, task.task_params)
|
||||
|
||||
def on_prefill_progress(processed: int, total: int) -> None:
|
||||
if self.device_rank == 0:
|
||||
if self.is_primary_output_node:
|
||||
self.event_sender.send(
|
||||
ChunkGenerated(
|
||||
command_id=task.command_id,
|
||||
|
||||
@@ -102,10 +102,6 @@ class Runner:
|
||||
self.device_rank = self.shard_metadata.device_rank
|
||||
|
||||
logger.info("hello from the runner")
|
||||
if getattr(self.shard_metadata, "immediate_exception", False):
|
||||
raise Exception("Fake exception - runner failed to spin up.")
|
||||
if timeout := getattr(self.shard_metadata, "should_timeout", 0):
|
||||
time.sleep(timeout)
|
||||
|
||||
self.setup_start_time = time.time()
|
||||
|
||||
|
||||
@@ -12,15 +12,39 @@ from anyio import (
|
||||
CancelScope,
|
||||
ClosedResourceError,
|
||||
)
|
||||
from exo_rs import (
|
||||
LVPublisher,
|
||||
LVSubscriber,
|
||||
TaskChunkSender,
|
||||
TaskRequest,
|
||||
TaskResponder,
|
||||
)
|
||||
from loguru import logger
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from exo.shared.constants import EXO_RUNNER_STDERR_LOG, EXO_RUNNER_STDOUT_LOG
|
||||
from exo.shared.types.chunks import ErrorChunk
|
||||
from exo.shared.types.chunks import ErrorChunk, PrefillProgressChunk
|
||||
from exo.shared.types.commands import (
|
||||
Command,
|
||||
TaskCancelled,
|
||||
TaskFinished,
|
||||
)
|
||||
from exo.shared.types.commands import (
|
||||
ImageEdits as ImageEditsCommand,
|
||||
)
|
||||
from exo.shared.types.commands import (
|
||||
ImageGeneration as ImageGenerationCommand,
|
||||
)
|
||||
from exo.shared.types.commands import (
|
||||
TextGeneration as TextGenerationCommand,
|
||||
)
|
||||
from exo.shared.types.common import CommandId
|
||||
from exo.shared.types.events import (
|
||||
ChunkGenerated,
|
||||
Event,
|
||||
RunnerStatusUpdated,
|
||||
TaskAcknowledged,
|
||||
TaskCreated,
|
||||
TaskStatusUpdated,
|
||||
)
|
||||
from exo.shared.types.tasks import (
|
||||
@@ -38,12 +62,12 @@ from exo.shared.types.worker.runners import (
|
||||
RunnerFailed,
|
||||
RunnerIdle,
|
||||
RunnerLoading,
|
||||
RunnerReady,
|
||||
RunnerRunning,
|
||||
RunnerShuttingDown,
|
||||
RunnerStatus,
|
||||
RunnerWarmingUp,
|
||||
)
|
||||
from exo.shared.types.worker.shards import ShardMetadata
|
||||
from exo.utils.async_process import AsyncProcess
|
||||
from exo.utils.channels import MpReceiver, MpSender, Receiver, Sender, mp_channel
|
||||
from exo.utils.fs import ensure_parent_directory_exists
|
||||
@@ -56,6 +80,22 @@ from exo.worker.runner.diagnostics import (
|
||||
|
||||
PREFILL_TIMEOUT_SECONDS = 60
|
||||
DECODE_TIMEOUT_SECONDS = 5
|
||||
type BridgeTask = TextGeneration | ImageGeneration | ImageEdits
|
||||
_BRIDGE_COMMAND_ADAPTER: TypeAdapter[Command] = TypeAdapter(Command)
|
||||
_BRIDGE_TASK_ADAPTER: TypeAdapter[BridgeTask] = TypeAdapter(BridgeTask)
|
||||
|
||||
|
||||
def _task_assignment_ids(key: str) -> tuple[str, TaskId] | None:
|
||||
prefix = "task_assignments/"
|
||||
if not key.startswith(prefix):
|
||||
return None
|
||||
|
||||
suffix = key.removeprefix(prefix)
|
||||
parts = suffix.split("/", maxsplit=1)
|
||||
if len(parts) != 2 or not parts[0] or not parts[1] or "/" in parts[1]:
|
||||
return None
|
||||
|
||||
return parts[0], TaskId(parts[1])
|
||||
|
||||
|
||||
@dataclass(eq=False)
|
||||
@@ -180,7 +220,6 @@ class RunnerStdioHandler:
|
||||
|
||||
@dataclass(eq=False)
|
||||
class RunnerSupervisor:
|
||||
shard_metadata: ShardMetadata
|
||||
bound_instance: BoundInstance
|
||||
runner_process: AsyncProcess
|
||||
_runner_stdio_handler: RunnerStdioHandler
|
||||
@@ -189,12 +228,23 @@ class RunnerSupervisor:
|
||||
_task_sender: MpSender[Task]
|
||||
_event_sender: Sender[Event]
|
||||
_cancel_sender: MpSender[TaskId]
|
||||
_task_responder: TaskResponder | None
|
||||
_task_assignment_subscriber: LVSubscriber
|
||||
runner_status_publisher: LVPublisher
|
||||
_assigned_tasks: dict[TaskId, BridgeTask] = field(default_factory=dict, init=False)
|
||||
_tg: TaskGroup = field(default_factory=TaskGroup, init=False)
|
||||
status: RunnerStatus = field(default_factory=RunnerIdle, init=False)
|
||||
pending: dict[TaskId, anyio.Event] = field(default_factory=dict, init=False)
|
||||
in_progress: dict[TaskId, Task] = field(default_factory=dict, init=False)
|
||||
completed: set[TaskId] = field(default_factory=set, init=False)
|
||||
cancelled: set[TaskId] = field(default_factory=set, init=False)
|
||||
bridge_tasks: dict[TaskId, BridgeTask] = field(default_factory=dict, init=False)
|
||||
bridge_command_tasks: dict[CommandId, TaskId] = field(
|
||||
default_factory=dict, init=False
|
||||
)
|
||||
bridge_chunk_senders: dict[CommandId, TaskChunkSender] = field(
|
||||
default_factory=dict, init=False
|
||||
)
|
||||
_cancel_watch_runner: anyio.CancelScope = field(
|
||||
default_factory=anyio.CancelScope, init=False
|
||||
)
|
||||
@@ -205,6 +255,9 @@ class RunnerSupervisor:
|
||||
*,
|
||||
bound_instance: BoundInstance,
|
||||
event_sender: Sender[Event],
|
||||
task_assignment_subscriber: LVSubscriber,
|
||||
runner_status_publisher: LVPublisher,
|
||||
task_responder: TaskResponder | None,
|
||||
initialize_timeout: float = 400,
|
||||
) -> Self:
|
||||
ev_send, ev_recv = mp_channel[Event | RunnerTerminationError]()
|
||||
@@ -226,11 +279,8 @@ class RunnerSupervisor:
|
||||
stdout_rx=runner_process.stdout, stderr_rx=runner_process.stderr
|
||||
)
|
||||
|
||||
shard_metadata = bound_instance.bound_shard
|
||||
|
||||
self = cls(
|
||||
bound_instance=bound_instance,
|
||||
shard_metadata=shard_metadata,
|
||||
runner_process=runner_process,
|
||||
_runner_stdio_handler=runner_stdio_handler,
|
||||
initialize_timeout=initialize_timeout,
|
||||
@@ -238,6 +288,9 @@ class RunnerSupervisor:
|
||||
_task_sender=task_sender,
|
||||
_cancel_sender=cancel_sender,
|
||||
_event_sender=event_sender,
|
||||
_task_responder=task_responder,
|
||||
_task_assignment_subscriber=task_assignment_subscriber,
|
||||
runner_status_publisher=runner_status_publisher,
|
||||
)
|
||||
|
||||
return self
|
||||
@@ -251,6 +304,12 @@ class RunnerSupervisor:
|
||||
|
||||
tg.start_soon(self._watch_runner)
|
||||
tg.start_soon(self._forward_events)
|
||||
if self._task_responder is not None:
|
||||
tg.start_soon(self._run_task_responder, self._task_responder)
|
||||
tg.start_soon(
|
||||
self._run_task_assignment_subscriber,
|
||||
self._task_assignment_subscriber,
|
||||
)
|
||||
finally:
|
||||
logger.info("Runner supervisor shutting down")
|
||||
if not self._cancel_watch_runner.cancel_called:
|
||||
@@ -267,6 +326,7 @@ class RunnerSupervisor:
|
||||
self._cancel_sender.close()
|
||||
|
||||
with anyio.CancelScope(shield=True):
|
||||
await self.runner_status_publisher.delete()
|
||||
await self.runner_process.stop()
|
||||
logger.info(
|
||||
f"Runner process successfully terminated: {self.runner_process.exitcode}"
|
||||
@@ -275,6 +335,44 @@ class RunnerSupervisor:
|
||||
def shutdown(self):
|
||||
self._tg.cancel_tasks()
|
||||
|
||||
async def _run_task_assignment_subscriber(self, subscriber: LVSubscriber) -> None:
|
||||
instance_id = self.bound_instance.instance.instance_id
|
||||
while (received := await subscriber.recv()) is not None:
|
||||
key, payload = received
|
||||
if payload is None:
|
||||
continue
|
||||
if (ids := _task_assignment_ids(key)) is None:
|
||||
continue
|
||||
assigned_instance_id, assigned_task_id = ids
|
||||
if assigned_instance_id != instance_id:
|
||||
continue
|
||||
|
||||
if payload == "":
|
||||
self._assigned_tasks.pop(assigned_task_id, None)
|
||||
continue
|
||||
|
||||
try:
|
||||
task = _BRIDGE_TASK_ADAPTER.validate_json(payload)
|
||||
except ValidationError:
|
||||
logger.warning(f"Ignoring invalid task assignment from {key}")
|
||||
continue
|
||||
|
||||
if task.instance_id != instance_id or task.task_id != assigned_task_id:
|
||||
logger.warning(f"Ignoring mismatched task assignment from {key}")
|
||||
continue
|
||||
|
||||
self._assigned_tasks[task.task_id] = task
|
||||
await self._reconcile_assigned_tasks()
|
||||
|
||||
async def _reconcile_assigned_tasks(self) -> None:
|
||||
if not isinstance(self.status, (RunnerReady, RunnerRunning)):
|
||||
return
|
||||
|
||||
for task in list(self._assigned_tasks.values()):
|
||||
if task.task_id in self.in_progress or task.task_id in self.completed:
|
||||
continue
|
||||
await self.start_task(task)
|
||||
|
||||
async def start_task(self, task: Task):
|
||||
if task.task_id in self.pending:
|
||||
logger.warning(
|
||||
@@ -320,33 +418,61 @@ class RunnerSupervisor:
|
||||
try:
|
||||
with self._ev_recv as events:
|
||||
async for event in events:
|
||||
if isinstance(event, RunnerTerminationError):
|
||||
# try to get exception if possible
|
||||
await self._check_runner(event)
|
||||
break
|
||||
if isinstance(event, RunnerStatusUpdated):
|
||||
self.status = event.runner_status
|
||||
if isinstance(event, TaskAcknowledged):
|
||||
self.pending.pop(event.task_id).set()
|
||||
continue
|
||||
if (
|
||||
isinstance(event, TaskStatusUpdated)
|
||||
and event.task_status == TaskStatus.Complete
|
||||
):
|
||||
# If a task has just been completed, we should be working on it.
|
||||
assert isinstance(
|
||||
self.status,
|
||||
(
|
||||
RunnerRunning,
|
||||
RunnerWarmingUp,
|
||||
RunnerLoading,
|
||||
RunnerConnecting,
|
||||
RunnerShuttingDown,
|
||||
),
|
||||
)
|
||||
self.in_progress.pop(event.task_id, None)
|
||||
self.completed.add(event.task_id)
|
||||
await self._event_sender.send(event)
|
||||
match event:
|
||||
case RunnerTerminationError():
|
||||
await self._check_runner(event)
|
||||
break
|
||||
|
||||
case RunnerStatusUpdated(runner_status=runner_status):
|
||||
self.status = runner_status
|
||||
await self.runner_status_publisher.put(
|
||||
self.status.model_dump_json()
|
||||
)
|
||||
await self._event_sender.send(event)
|
||||
await self._reconcile_assigned_tasks()
|
||||
|
||||
case TaskAcknowledged(task_id=task_id):
|
||||
self.pending.pop(task_id).set()
|
||||
|
||||
case TaskStatusUpdated(
|
||||
task_id=task_id, task_status=TaskStatus.Complete
|
||||
):
|
||||
# If a task has just been completed, we should be working on it.
|
||||
assert isinstance(
|
||||
self.status,
|
||||
(
|
||||
RunnerRunning,
|
||||
RunnerWarmingUp,
|
||||
RunnerLoading,
|
||||
RunnerConnecting,
|
||||
RunnerShuttingDown,
|
||||
),
|
||||
)
|
||||
self.in_progress.pop(task_id, None)
|
||||
self.completed.add(task_id)
|
||||
self._assigned_tasks.pop(task_id, None)
|
||||
await self._event_sender.send(event)
|
||||
if task_id in self.bridge_tasks:
|
||||
await self._finish_bridge_task_id(task_id)
|
||||
|
||||
case ChunkGenerated(command_id=command_id, chunk=chunk):
|
||||
task_id = self.bridge_command_tasks.get(command_id)
|
||||
chunk_sender = self.bridge_chunk_senders.get(command_id)
|
||||
if task_id is None or chunk_sender is None:
|
||||
logger.debug(
|
||||
f"Dropping bridge chunk for inactive command {command_id}"
|
||||
)
|
||||
continue
|
||||
|
||||
await chunk_sender.send(chunk.model_dump_json())
|
||||
if (
|
||||
not isinstance(chunk, PrefillProgressChunk)
|
||||
and chunk.finish_reason is not None
|
||||
):
|
||||
self.bridge_chunk_senders.pop(command_id, None)
|
||||
|
||||
case _:
|
||||
await self._event_sender.send(event)
|
||||
except (ClosedResourceError, BrokenResourceError):
|
||||
# this is the happy path shutdown - we don't need to spam log with it
|
||||
await self._check_runner()
|
||||
@@ -354,6 +480,108 @@ class RunnerSupervisor:
|
||||
for tid in self.pending:
|
||||
self.pending[tid].set()
|
||||
|
||||
async def _run_task_responder(self, responder: TaskResponder) -> None:
|
||||
while True:
|
||||
received = await responder.recv()
|
||||
if received is None:
|
||||
return
|
||||
request, chunk_sender, payload = received
|
||||
if payload is None:
|
||||
request.reply_err("Task command query did not include a payload")
|
||||
continue
|
||||
|
||||
try:
|
||||
command = _BRIDGE_COMMAND_ADAPTER.validate_json(payload)
|
||||
match command:
|
||||
case TextGenerationCommand():
|
||||
await self._submit_bridge_task(request, chunk_sender, command)
|
||||
case ImageGenerationCommand():
|
||||
await self._submit_bridge_task(request, chunk_sender, command)
|
||||
case ImageEditsCommand():
|
||||
await self._submit_bridge_task(request, chunk_sender, command)
|
||||
case TaskCancelled(cancelled_command_id=command_id):
|
||||
await self._cancel_bridge_task(command_id)
|
||||
request.reply(command_id)
|
||||
case TaskFinished(finished_command_id=command_id):
|
||||
await self._finish_bridge_task(command_id)
|
||||
request.reply(command_id)
|
||||
case _:
|
||||
request.reply_err(f"Unsupported bridge command: {command}")
|
||||
except Exception as exception:
|
||||
logger.opt(exception=exception).warning(
|
||||
"Failed to admit bridge command"
|
||||
)
|
||||
request.reply_err(str(exception))
|
||||
|
||||
async def _submit_bridge_task(
|
||||
self,
|
||||
request: TaskRequest,
|
||||
chunk_sender: TaskChunkSender,
|
||||
command: TextGenerationCommand | ImageGenerationCommand | ImageEditsCommand,
|
||||
) -> None:
|
||||
task_id = TaskId()
|
||||
match command:
|
||||
case TextGenerationCommand():
|
||||
task = TextGeneration(
|
||||
task_id=task_id,
|
||||
command_id=command.command_id,
|
||||
instance_id=self.bound_instance.instance.instance_id,
|
||||
task_status=TaskStatus.Pending,
|
||||
task_params=command.task_params,
|
||||
)
|
||||
case ImageGenerationCommand():
|
||||
task = ImageGeneration(
|
||||
task_id=task_id,
|
||||
command_id=command.command_id,
|
||||
instance_id=self.bound_instance.instance.instance_id,
|
||||
task_status=TaskStatus.Pending,
|
||||
task_params=command.task_params,
|
||||
)
|
||||
case ImageEditsCommand():
|
||||
task = ImageEdits(
|
||||
task_id=task_id,
|
||||
command_id=command.command_id,
|
||||
instance_id=self.bound_instance.instance.instance_id,
|
||||
task_status=TaskStatus.Pending,
|
||||
task_params=command.task_params,
|
||||
)
|
||||
|
||||
self.bridge_tasks[task.task_id] = task
|
||||
self.bridge_command_tasks[command.command_id] = task.task_id
|
||||
self.bridge_chunk_senders[command.command_id] = chunk_sender
|
||||
await self._event_sender.send(TaskCreated(task_id=task.task_id, task=task))
|
||||
assert self._task_responder is not None
|
||||
await self._task_responder.assign_task(task.task_id, task.model_dump_json())
|
||||
request.reply(command.command_id)
|
||||
|
||||
async def _cancel_bridge_task(self, command_id: CommandId) -> None:
|
||||
task_id = self.bridge_command_tasks.get(command_id)
|
||||
if task_id is None:
|
||||
logger.warning(f"Unable to cancel unknown bridge command {command_id}")
|
||||
return
|
||||
await self.cancel_task(task_id)
|
||||
await self._event_sender.send(
|
||||
TaskStatusUpdated(task_id=task_id, task_status=TaskStatus.Cancelled)
|
||||
)
|
||||
await self._finish_bridge_task_id(task_id)
|
||||
|
||||
async def _finish_bridge_task(self, command_id: CommandId) -> None:
|
||||
task_id = self.bridge_command_tasks.get(command_id)
|
||||
if task_id is None:
|
||||
logger.warning(f"Unable to finish unknown bridge command {command_id}")
|
||||
return
|
||||
await self._finish_bridge_task_id(task_id)
|
||||
|
||||
async def _finish_bridge_task_id(self, task_id: TaskId) -> None:
|
||||
task = self.bridge_tasks.pop(task_id, None)
|
||||
if task is None:
|
||||
logger.warning(f"Unable to finish unknown bridge task {task_id}")
|
||||
return
|
||||
self.bridge_command_tasks.pop(task.command_id, None)
|
||||
self.bridge_chunk_senders.pop(task.command_id, None)
|
||||
if self._task_responder is not None:
|
||||
await self._task_responder.unassign_task(task_id)
|
||||
|
||||
async def _watch_runner(self) -> None:
|
||||
with self._cancel_watch_runner:
|
||||
while True:
|
||||
@@ -412,7 +640,7 @@ class RunnerSupervisor:
|
||||
ChunkGenerated(
|
||||
command_id=task.command_id,
|
||||
chunk=ErrorChunk(
|
||||
model=self.shard_metadata.model_card.model_id,
|
||||
model=self.bound_instance.bound_shard.model_card.model_id,
|
||||
diagnostics=diagnostics,
|
||||
error_message=(
|
||||
"Runner shutdown before completing command "
|
||||
|
||||
@@ -11,7 +11,12 @@ from exo.shared.types.worker.instances import (
|
||||
InstanceId,
|
||||
MlxRingInstance,
|
||||
)
|
||||
from exo.shared.types.worker.runners import RunnerId, RunnerStatus, ShardAssignments
|
||||
from exo.shared.types.worker.runners import (
|
||||
RunnerId,
|
||||
RunnerStatus,
|
||||
ShardAssignments,
|
||||
ShardWithId,
|
||||
)
|
||||
from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
|
||||
|
||||
|
||||
@@ -52,16 +57,22 @@ def get_pipeline_shard_metadata(
|
||||
)
|
||||
|
||||
|
||||
# todo: clean up legacy formatted shards
|
||||
def get_shard_assignments(
|
||||
model_id: ModelId,
|
||||
node_to_runner: dict[NodeId, RunnerId],
|
||||
runner_to_shard: dict[RunnerId, ShardMetadata],
|
||||
) -> ShardAssignments:
|
||||
return ShardAssignments(
|
||||
model_id=model_id,
|
||||
node_to_runner=node_to_runner,
|
||||
runner_to_shard=runner_to_shard,
|
||||
)
|
||||
pon = 0
|
||||
shards = [
|
||||
ShardWithId(nid, rid := node_to_runner[nid], runner_to_shard[rid])
|
||||
for nid in node_to_runner
|
||||
]
|
||||
for i, (_, _, shard) in enumerate(shards):
|
||||
if shard.is_primary_output():
|
||||
pon = i
|
||||
|
||||
return ShardAssignments(model_id=model_id, shards=shards, primary_output_node=pon)
|
||||
|
||||
|
||||
def get_mlx_ring_instance(
|
||||
|
||||
@@ -53,8 +53,6 @@ def test_plan_requests_download_when_waiting_and_shard_not_downloaded():
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
@@ -109,8 +107,6 @@ def test_plan_loads_model_when_all_shards_downloaded_and_waiting():
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
@@ -155,8 +151,6 @@ def test_plan_does_not_request_download_when_shard_already_downloaded():
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
@@ -206,8 +200,6 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
@@ -230,8 +222,6 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
|
||||
@@ -53,8 +53,6 @@ def test_plan_kills_runner_when_instance_missing():
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
@@ -96,8 +94,6 @@ def test_plan_kills_runner_when_sibling_failed():
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
@@ -131,8 +127,6 @@ def test_plan_creates_runner_when_missing_for_node():
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
@@ -173,8 +167,6 @@ def test_plan_does_not_create_runner_when_supervisor_already_present():
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
@@ -206,8 +198,6 @@ def test_plan_does_not_create_runner_for_unassigned_node():
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
|
||||
@@ -1,303 +0,0 @@
|
||||
from typing import cast
|
||||
|
||||
import exo.worker.plan as plan_mod
|
||||
from exo.shared.types.tasks import Task, TaskId, TaskStatus, TextGeneration
|
||||
from exo.shared.types.text_generation import (
|
||||
InputMessage,
|
||||
InputMessageContent,
|
||||
TextGenerationTaskParams,
|
||||
)
|
||||
from exo.shared.types.worker.instances import BoundInstance, InstanceId
|
||||
from exo.shared.types.worker.runners import (
|
||||
RunnerIdle,
|
||||
RunnerReady,
|
||||
RunnerRunning,
|
||||
)
|
||||
from exo.utils.keyed_backoff import KeyedBackoff
|
||||
from exo.worker.tests.constants import (
|
||||
COMMAND_1_ID,
|
||||
INSTANCE_1_ID,
|
||||
MODEL_A_ID,
|
||||
NODE_A,
|
||||
NODE_B,
|
||||
RUNNER_1_ID,
|
||||
RUNNER_2_ID,
|
||||
TASK_1_ID,
|
||||
)
|
||||
from exo.worker.tests.unittests.conftest import (
|
||||
FakeRunnerSupervisor,
|
||||
OtherTask,
|
||||
get_mlx_ring_instance,
|
||||
get_pipeline_shard_metadata,
|
||||
)
|
||||
|
||||
|
||||
def test_plan_forwards_pending_chat_completion_when_runner_ready():
|
||||
"""
|
||||
When there is a pending TextGeneration for the local instance and all
|
||||
runners are Ready/Running, plan() should forward that task.
|
||||
"""
|
||||
shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2)
|
||||
shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2)
|
||||
instance = get_mlx_ring_instance(
|
||||
instance_id=INSTANCE_1_ID,
|
||||
model_id=MODEL_A_ID,
|
||||
node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID},
|
||||
runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1},
|
||||
)
|
||||
bound_instance = BoundInstance(
|
||||
instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A
|
||||
)
|
||||
local_runner = FakeRunnerSupervisor(
|
||||
bound_instance=bound_instance, status=RunnerReady()
|
||||
)
|
||||
|
||||
runners = {RUNNER_1_ID: local_runner}
|
||||
instances = {INSTANCE_1_ID: instance}
|
||||
all_runners = {
|
||||
RUNNER_1_ID: RunnerReady(),
|
||||
RUNNER_2_ID: RunnerReady(),
|
||||
}
|
||||
|
||||
task = TextGeneration(
|
||||
task_id=TASK_1_ID,
|
||||
instance_id=INSTANCE_1_ID,
|
||||
task_status=TaskStatus.Pending,
|
||||
command_id=COMMAND_1_ID,
|
||||
task_params=TextGenerationTaskParams(
|
||||
model=MODEL_A_ID,
|
||||
input=[InputMessage(role="user", content=InputMessageContent(""))],
|
||||
),
|
||||
)
|
||||
|
||||
result = plan_mod.plan(
|
||||
node_id=NODE_A,
|
||||
runners=runners, # type: ignore
|
||||
global_download_status={NODE_A: []},
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={TASK_1_ID: task},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
|
||||
assert result is task
|
||||
|
||||
|
||||
def test_plan_does_not_forward_chat_completion_if_any_runner_not_ready():
|
||||
"""
|
||||
Even with a pending TextGeneration, plan() should not forward it unless
|
||||
all runners for the instance are Ready/Running.
|
||||
"""
|
||||
shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2)
|
||||
shard2 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2)
|
||||
instance = get_mlx_ring_instance(
|
||||
instance_id=INSTANCE_1_ID,
|
||||
model_id=MODEL_A_ID,
|
||||
node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID},
|
||||
runner_to_shard={RUNNER_1_ID: shard1, RUNNER_2_ID: shard2},
|
||||
)
|
||||
bound_instance = BoundInstance(
|
||||
instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A
|
||||
)
|
||||
local_runner = FakeRunnerSupervisor(
|
||||
bound_instance=bound_instance, status=RunnerReady()
|
||||
)
|
||||
|
||||
runners = {RUNNER_1_ID: local_runner}
|
||||
instances = {INSTANCE_1_ID: instance}
|
||||
all_runners = {
|
||||
RUNNER_1_ID: RunnerReady(),
|
||||
RUNNER_2_ID: RunnerIdle(),
|
||||
}
|
||||
|
||||
task = TextGeneration(
|
||||
task_id=TASK_1_ID,
|
||||
instance_id=INSTANCE_1_ID,
|
||||
task_status=TaskStatus.Pending,
|
||||
command_id=COMMAND_1_ID,
|
||||
task_params=TextGenerationTaskParams(
|
||||
model=MODEL_A_ID,
|
||||
input=[InputMessage(role="user", content=InputMessageContent(""))],
|
||||
),
|
||||
)
|
||||
|
||||
result = plan_mod.plan(
|
||||
node_id=NODE_A,
|
||||
runners=runners, # type: ignore
|
||||
global_download_status={NODE_A: [], NODE_B: []},
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={TASK_1_ID: task},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_plan_does_not_forward_tasks_for_other_instances():
|
||||
"""
|
||||
plan() should ignore pending TextGeneration tasks whose instance_id does
|
||||
not match the local instance.
|
||||
"""
|
||||
shard = get_pipeline_shard_metadata(model_id=MODEL_A_ID, device_rank=0)
|
||||
local_instance = get_mlx_ring_instance(
|
||||
instance_id=INSTANCE_1_ID,
|
||||
model_id=MODEL_A_ID,
|
||||
node_to_runner={NODE_A: RUNNER_1_ID},
|
||||
runner_to_shard={RUNNER_1_ID: shard},
|
||||
)
|
||||
bound_instance = BoundInstance(
|
||||
instance=local_instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A
|
||||
)
|
||||
local_runner = FakeRunnerSupervisor(
|
||||
bound_instance=bound_instance, status=RunnerReady()
|
||||
)
|
||||
|
||||
runners = {RUNNER_1_ID: local_runner}
|
||||
instances = {INSTANCE_1_ID: local_instance}
|
||||
all_runners = {RUNNER_1_ID: RunnerReady()}
|
||||
|
||||
other_instance_id = InstanceId("instance-2")
|
||||
foreign_task = TextGeneration(
|
||||
task_id=TaskId("other-task"),
|
||||
instance_id=other_instance_id,
|
||||
task_status=TaskStatus.Pending,
|
||||
command_id=COMMAND_1_ID,
|
||||
task_params=TextGenerationTaskParams(
|
||||
model=MODEL_A_ID,
|
||||
input=[InputMessage(role="user", content=InputMessageContent(""))],
|
||||
),
|
||||
)
|
||||
|
||||
result = plan_mod.plan(
|
||||
node_id=NODE_A,
|
||||
runners=runners, # type: ignore
|
||||
global_download_status={NODE_A: []},
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={foreign_task.task_id: foreign_task},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_plan_ignores_non_pending_or_non_chat_tasks():
|
||||
"""
|
||||
_pending_tasks should not forward tasks that are either not TextGeneration
|
||||
or not in Pending/Running states.
|
||||
"""
|
||||
shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2)
|
||||
shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2)
|
||||
instance = get_mlx_ring_instance(
|
||||
instance_id=INSTANCE_1_ID,
|
||||
model_id=MODEL_A_ID,
|
||||
node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID},
|
||||
runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1},
|
||||
)
|
||||
bound_instance = BoundInstance(
|
||||
instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A
|
||||
)
|
||||
|
||||
local_runner = FakeRunnerSupervisor(
|
||||
bound_instance=bound_instance, status=RunnerReady()
|
||||
)
|
||||
|
||||
runners = {RUNNER_1_ID: local_runner}
|
||||
instances = {INSTANCE_1_ID: instance}
|
||||
all_runners = {
|
||||
RUNNER_1_ID: RunnerReady(),
|
||||
RUNNER_2_ID: RunnerReady(),
|
||||
}
|
||||
|
||||
completed_task = TextGeneration(
|
||||
task_id=TASK_1_ID,
|
||||
instance_id=INSTANCE_1_ID,
|
||||
task_status=TaskStatus.Complete,
|
||||
command_id=COMMAND_1_ID,
|
||||
task_params=TextGenerationTaskParams(
|
||||
model=MODEL_A_ID,
|
||||
input=[InputMessage(role="user", content=InputMessageContent(""))],
|
||||
),
|
||||
)
|
||||
|
||||
other_task_id = TaskId("other-task")
|
||||
|
||||
other_task = cast(
|
||||
Task,
|
||||
cast(
|
||||
object,
|
||||
OtherTask(
|
||||
task_id=other_task_id,
|
||||
instance_id=INSTANCE_1_ID,
|
||||
task_status=TaskStatus.Pending,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
result = plan_mod.plan(
|
||||
node_id=NODE_A,
|
||||
runners=runners, # type: ignore
|
||||
global_download_status={NODE_A: [], NODE_B: []},
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={TASK_1_ID: completed_task, other_task_id: other_task},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_plan_returns_none_when_nothing_to_do():
|
||||
"""
|
||||
If there are healthy runners, no downloads needed, and no pending tasks,
|
||||
plan() should return None (steady state).
|
||||
"""
|
||||
shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2)
|
||||
shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2)
|
||||
instance = get_mlx_ring_instance(
|
||||
instance_id=INSTANCE_1_ID,
|
||||
model_id=MODEL_A_ID,
|
||||
node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID},
|
||||
runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1},
|
||||
)
|
||||
bound_instance = BoundInstance(
|
||||
instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A
|
||||
)
|
||||
local_runner = FakeRunnerSupervisor(
|
||||
bound_instance=bound_instance, status=RunnerRunning()
|
||||
)
|
||||
|
||||
runners = {RUNNER_1_ID: local_runner}
|
||||
instances = {INSTANCE_1_ID: instance}
|
||||
all_runners = {
|
||||
RUNNER_1_ID: RunnerRunning(),
|
||||
RUNNER_2_ID: RunnerRunning(),
|
||||
}
|
||||
|
||||
result = plan_mod.plan(
|
||||
node_id=NODE_A,
|
||||
runners=runners, # type: ignore
|
||||
global_download_status={NODE_A: [], NODE_B: []},
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
|
||||
assert result is None
|
||||
@@ -62,8 +62,6 @@ def test_plan_starts_warmup_for_accepting_rank_when_all_loaded_or_warming():
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
@@ -107,8 +105,6 @@ def test_plan_starts_warmup_for_rank_zero_after_others_warming():
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
@@ -151,8 +147,6 @@ def test_plan_does_not_start_warmup_for_non_zero_rank_until_all_loaded_or_warmin
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
@@ -199,8 +193,6 @@ def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming():
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
@@ -219,8 +211,6 @@ def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming():
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
@@ -266,8 +256,6 @@ def test_plan_starts_warmup_for_connecting_rank_after_others_warming():
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
@@ -312,8 +300,6 @@ def test_plan_does_not_start_warmup_for_accepting_rank_until_all_loaded_or_warmi
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
@@ -357,8 +343,6 @@ def test_plan_does_not_start_warmup_for_connecting_rank_until_others_warming():
|
||||
instances=instances,
|
||||
all_runners=all_runners,
|
||||
tasks={},
|
||||
input_chunk_buffer={},
|
||||
image_cache={},
|
||||
instance_backoff=KeyedBackoff(),
|
||||
download_backoff=KeyedBackoff(),
|
||||
)
|
||||
|
||||
+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_ZENOH_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix run $remote_installable" 2>&1 |
|
||||
"ENABLE_DISAGGREGATION=true EXO_ZENOH_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix run $remote_installable" 2>&1 |
|
||||
awk -v p="${colour}[${host}]${reset}" '{ print p $0; fflush() }' &
|
||||
done
|
||||
|
||||
|
||||
@@ -52,18 +52,17 @@ def instance_id_from_instance(instance: dict[str, Any]) -> str:
|
||||
|
||||
def nodes_used_in_instance(instance: dict[str, Any]) -> int:
|
||||
inner = unwrap_instance(instance)
|
||||
return len(inner["shardAssignments"]["nodeToRunner"])
|
||||
return len(inner["shardAssignments"]["shards"])
|
||||
|
||||
|
||||
def runner_ids_from_instance(instance: dict[str, Any]) -> list[str]:
|
||||
inner = unwrap_instance(instance)
|
||||
runner_to_shard = inner["shardAssignments"]["runnerToShard"]
|
||||
return list(runner_to_shard.keys())
|
||||
return [r for (_, r, _) in inner["shardAssignments"]["shards"]]
|
||||
|
||||
|
||||
def node_ids_from_instance(instance: dict[str, Any]) -> list[str]:
|
||||
inner = unwrap_instance(instance)
|
||||
return list(inner["shardAssignments"]["nodeToRunner"].keys())
|
||||
return [n for (n, _, _) in inner["shardAssignments"]["shards"]]
|
||||
|
||||
|
||||
def runner_ready(runner: dict[str, Any]) -> bool:
|
||||
@@ -322,8 +321,8 @@ def run_planning_phase(
|
||||
|
||||
# Get nodes from preview
|
||||
inner = unwrap_instance(preview["instance"])
|
||||
node_ids = list(inner["shardAssignments"]["nodeToRunner"].keys())
|
||||
runner_to_shard = inner["shardAssignments"]["runnerToShard"]
|
||||
node_ids = [n for (n, _, _) in inner["shardAssignments"]["shards"]]
|
||||
shards = inner["shardAssignments"]["shards"]
|
||||
|
||||
needs_download = False
|
||||
|
||||
@@ -391,9 +390,7 @@ def run_planning_phase(
|
||||
|
||||
# Start downloads (idempotent)
|
||||
download_t0 = time.perf_counter() if needs_download else None
|
||||
for node_id in node_ids:
|
||||
runner_id = inner["shardAssignments"]["nodeToRunner"][node_id]
|
||||
shard = runner_to_shard[runner_id]
|
||||
for node_id, _, shard in shards:
|
||||
client.request_json(
|
||||
"POST",
|
||||
"/download/start",
|
||||
|
||||
@@ -450,7 +450,7 @@ build = [
|
||||
]
|
||||
mlx = [
|
||||
{ name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260506+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260512+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl" }, marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
@@ -461,7 +461,7 @@ mlx = [
|
||||
]
|
||||
mlx-cpu = [
|
||||
{ name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260506+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-mlx-cpu') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260512+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-mlx-cpu') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cpu') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra != 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra != 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra != 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl" }, marker = "(platform_machine != 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cpu') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra != 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra != 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra != 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx-cpu", marker = "sys_platform == 'linux'" },
|
||||
@@ -473,7 +473,7 @@ mlx-cpu = [
|
||||
]
|
||||
mlx-cuda12 = [
|
||||
{ name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260506+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260512+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra != 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl" }, marker = "(platform_machine != 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra != 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx-cuda-12", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx_cuda_12-0.32.0-py3-none-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra != 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
@@ -487,7 +487,7 @@ mlx-cuda12 = [
|
||||
]
|
||||
mlx-cuda13 = [
|
||||
{ name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260506+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260512+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl" }, marker = "(platform_machine != 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx-cuda-13", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx_cuda_13-0.32.0-py3-none-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none')" },
|
||||
@@ -1343,7 +1343,7 @@ dependencies = [
|
||||
{ name = "hf-transfer", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "matplotlib", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260506+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260512+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl" }, marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
@@ -1374,7 +1374,7 @@ sdist = { url = "https://files.pythonhosted.org/packages/55/fa/96d4cc7ada2833571
|
||||
|
||||
[[package]]
|
||||
name = "mlx"
|
||||
version = "0.32.0.dev20260506+cc3f3e60"
|
||||
version = "0.32.0.dev20260512+cc3f3e60"
|
||||
source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }
|
||||
resolution-markers = [
|
||||
"sys_platform == 'darwin'",
|
||||
@@ -1549,7 +1549,7 @@ version = "0.31.3"
|
||||
source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Fdeepseek-v4#6a3df6cd6b00a347ee40f12d97a182aaf86ea599" }
|
||||
dependencies = [
|
||||
{ name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260506+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260512+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
@@ -1565,7 +1565,7 @@ dependencies = [
|
||||
{ name = "datasets", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "fastapi", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "miniaudio", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260506+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260512+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine != 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx", version = "0.32.0", source = { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl" }, marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (platform_machine == 'aarch64' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (sys_platform != 'linux' and extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
{ name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
|
||||
|
||||
Reference in new issue
Block a user