Compare commits

...
13 Commits
Author SHA1 Message Date
Evan ded7840499 api streams 2026-05-08 16:04:25 +01:00
Evan 7ed3eaa617 json state proxy 2026-05-07 15:36:03 +01:00
Evan e113d42ba1 uncap 2026-05-07 13:28:51 +01:00
Evan c5245dd87e libp2p -> zenoh 2026-05-07 11:21:59 +01:00
Alex Cheema edef8004f8 Store custom model cards in State (#2024)
## Why

Workers currently update their custom model-card cache by reacting to
`CustomModelCardAdded` / `CustomModelCardDeleted` events directly. That
is another snapshot footgun: a worker restored from State may never see
the historical add/delete event, so the durable State must include the
desired custom-card set.

## How

- Add `State.custom_model_cards`, keyed by `ModelId`.
- Reduce `CustomModelCardAdded` into State.
- Reduce `CustomModelCardDeleted` into State.
- Add focused reducer tests for add and delete.

This PR only makes custom cards durable in State. A follow-up PR will
make workers reconcile their on-disk custom-card cache from this state
instead of relying on those events directly.

## Tests

- `uv run pytest
src/exo/shared/tests/test_apply/test_apply_custom_model_cards.py
src/exo/shared/tests/test_state_serialization.py`
- `uv run pytest`
- `uv run ruff check src/exo/shared/types/state.py
src/exo/shared/apply.py
src/exo/shared/tests/test_apply/test_apply_custom_model_cards.py`
- `uv run basedpyright`
- `nix fmt`
2026-05-07 09:06:39 +01:00
Alex CheemaandClaude Opus 4.7 a0c00f9dfd fix(placement): gate RDMA on nodeRdmaCtl.enabled at both endpoints (#2014)
## Summary

- Fixes a bug where `POST /place_instance` (and the dashboard UI) would
accept an MlxJaccl/RDMA instance spanning nodes whose
`nodeRdmaCtl.enabled` was `false`, because topology + placement
consulted Thunderbolt-derived RDMA edges without checking the per-node
`rdma_ctl` status.
- Three-layer fix: topology only emits `RDMAConnection` edges when both
endpoints have `nodeRdmaCtl.enabled = true`; flipping a node to disabled
immediately purges every RDMA edge touching it; `place_instance`
additionally rejects RDMA cycles containing any disabled or unobserved
node as a defense-in-depth check on the API/master path.

## Details

- `src/exo/shared/apply.py`
- `MacThunderboltConnections` case now filters out RDMA connections
whose source or sink lacks observed-and-enabled `rdma_ctl` status
(missing entry → treated as disabled).
- `RdmaCtlStatus` case now calls
`topology.remove_all_rdma_connections_touching(node_id)` when the node
reports disabled, so consumers don't have to wait for the next TB poll.
- `src/exo/shared/topology.py`
- New `Topology.remove_all_rdma_connections_touching(node_id)` removes
every RDMA edge incident to the node (incoming and outgoing) while
leaving socket edges intact.
- `src/exo/master/placement.py`
- `place_instance` accepts `node_rdma_ctl: Mapping[NodeId,
NodeRdmaCtlStatus] | None`. The `is_rdma_cycle` filter now also requires
`nodeRdmaCtl.enabled` for every node in the cycle. MlxJaccl placement
raises the existing "no RDMA-connected cycles available" error if no
qualifying cycle remains.
- `src/exo/api/main.py`, `src/exo/master/main.py`
  - Both placement entrypoints now pass `state.node_rdma_ctl` through.

## Tests

- `src/exo/shared/tests/test_apply/test_apply_rdma_gating.py` (new): six
unit tests covering enabled/disabled/missing combinations on apply, the
immediate-purge transition, and that purging RDMA edges leaves socket
edges untouched.
- `src/exo/master/tests/test_placement.py`: existing
`test_tensor_rdma_backend_connectivity_matrix` updated to pass
`node_rdma_ctl`. Two new tests assert MlxJaccl placement is rejected
when any cycle node is `enabled=false` or has no `rdma_ctl` entry.

## Test plan

- [x] `uv run basedpyright` — 0 errors
- [x] `uv run ruff check` — clean
- [x] `nix fmt`
- [x] `uv run pytest` — 429 passed, 1 skipped
- [ ] On a real mixed cluster (s15/s16 disabled, s17/s18 enabled),
confirm:
- [ ] `POST /place_instance` for an RDMA instance including s15 or s16
returns an error
  - [ ] An RDMA instance can still be placed across {s17, s18}
- [ ] `GET /state` shows no `sourceRdmaIface`/`sinkRdmaIface` on s15↔s16
connections
- [ ] Dashboard previews don't surface RDMA-spanning options that
include s15/s16

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-07 07:00:15 +00:00
Drifter4242 89d20c1888 fix(inference): prevent TP collective deadlock via agree_on_tasks order (#2048)
If you have two machines and make two requests at the same time, it can
crash. This is because the tasks can sometimes end up in different
orders on different machines. We need to sort the tasks and
mx_all_gather_tasks already sorts the tasks but the code ignores that
ordering. The fix is to make sure the sort order is preserved.

The rest is written by Sonnet (reviewed by me):

Tensor-parallel inference requires that every rank enqueues tasks in the
same order before running agree_on_tasks collectives. The old
implementation filtered from _maybe_queue:

self._queue.extend(task for task in self._maybe_queue if task in agreed)
self._maybe_queue = [task for task in self._maybe_queue if task in
different]

Because _maybe_queue is independently ordered per-rank (tasks arrive via
gRPC in whatever order the API server sends them), two concurrent
requests could produce different _maybe_queue orderings on rank 0 vs
rank 1. The filter then preserved those different orders into _queue, so
each rank started processing tasks in a different sequence. The next mlx
collective (all_reduce, all_gather, etc.) on rank 0 corresponded to a
different task than on rank 1 → permanent deadlock.

Fix: extend from agreed directly. mx_all_gather_tasks returns agreed as
a list sorted by task_id on all ranks, so every rank appends the same
sequence regardless of local arrival order.

Applies to both SequentialGenerator and BatchGenerator.

## Motivation

`agree_on_tasks` is called on every rank after accumulating new requests
in
`_maybe_queue`. Its job is to run an `all_gather` collective so all
ranks agree
on which tasks to promote to `_queue` before the next inference step.

The old implementation re-imposed **local arrival order** when extending
`_queue`:

```python
self._queue.extend(task for task in self._maybe_queue if task in agreed)
```

`mx_all_gather_tasks` already returns `agreed` sorted by `task_id` — the
same
deterministic order on every rank. But iterating `self._maybe_queue`
instead of
`agreed` discarded that sort and substituted the local gRPC arrival
order, which
differs per rank under concurrent load. Two concurrent requests arriving
in
`[A, B]` order on rank 0 and `[B, A]` on rank 1 caused the first MLX
collective
in the next step to hang permanently: each rank was executing a
different task's
collective and would never match.

## Changes

`SequentialGenerator.agree_on_tasks` and
`BatchGenerator.agree_on_tasks`:

```python
# Before
self._queue.extend(task for task in self._maybe_queue if task in agreed)
self._maybe_queue = [task for task in self._maybe_queue if task in different]

# After
self._queue.extend(agreed)          # preserves mx_all_gather_tasks sort order
self._maybe_queue = list(different) # already in local order; filter was redundant
```

## Why It Works

`mx_all_gather_tasks` (in `utils_mlx.py`) computes the agreed set then
sorts by
`task_id`:

```python
agreed = [local_tasks[tid] for tid in sorted(agreed_ids)]
```

Because `task_id` is a UUID and the sort is lexicographic, every rank
produces
the same `agreed` list regardless of local arrival order. Using `agreed`
directly
preserves this guarantee. The `different` list (tasks not yet seen on
all ranks)
is built by iterating `tasks` in local order, which is already correct.

## Test Plan

### Manual Testing

**Hardware:** 2× Mac Studio M3 Ultra 512 GB, Thunderbolt 5 direct
bridge,
`MlxJaccl` RDMA tensor-parallel (`moonshotai/Kimi-K2.6`, 595 GB INT4, 61
layers).

- Sent concurrent streaming requests; confirmed all complete without
deadlock.
- This hardware configuration (sub-millisecond inter-node latency) is
the most
likely to trigger the race, as requests from separate HTTP connections
can
reach rank 0 and rank 1 in opposite order before `agree_on_tasks` runs.

### Automated Testing

All existing tests pass: `pytest src -m "not slow"
--import-mode=importlib`
— 422/422 passed. The existing `test_event_ordering.py` covers the
`agree_on_tasks` call path with a mock that returns tasks in consistent
order;
the race requires real distributed hardware to reproduce
deterministically.
2026-05-06 12:24:58 +00:00
Evan Quineyandciaranbor dbcceaa50c Initialise _cancelled_tasks in ImageEngine (#2051)
we yielded nonsense chunks from engines; we didn't initialize the image
engine correctly. mostly rewrite of #2049

---------

Co-authored-by: ciaranbor <ciaranborourke-dev@proton.me>
2026-05-05 17:27:57 +01:00
Sam BradburyandSam Bradbury 9c6ff4ce95 feat: update rdma_ctl instructions (#1977)
## Motivation

The RDMA setup instructions were missing a step: after booting to
Recovery mode, users need to open Terminal from the Utilities menu
before they can run the `rdma_ctl` command. Without this step, users
following the instructions wouldn't know how to access a terminal in
Recovery mode. This step was already in the README just not in the UI
notifications.

## Changes

Added a missing instruction step — "Open Terminal from the Utilities
menu" — to three instances of the RDMA setup flow in
`dashboard/src/routes/+page.svelte`.

## Why It Works

N/A copy change only. 

## Test Plan

### Manual Testing
Hardware: MacBook Pro M4 Max 48GB

### Automated Testing
No automated tests affected; this is a UI copy change only.

Co-authored-by: Sam Bradbury <sam@consultbradbury.com>
2026-05-01 11:18:57 +00:00
ecohash-coandJordan Miller b26268dfaf fix(macos-app): disable URL response caching for cluster-state polling (#2005)
Fixes #2004.

`ClusterStateService` polls `/state` at 2 Hz via `URLSession.shared`,
which keeps an on-disk `URLCache` attached by default. Every polled
response body gets persisted under `~/Library/Caches/exolabs.EXO/`,
sustaining ~500–620 KB/sec of file-backed memory dirtied — far above
macOS's ~25 KB/sec per-process daily-average baseline. Six
microstackshot reports observed on a single Mac Studio M3 Ultra over
eight days, with one 15-hour run accumulating 34.36 GB of cache writes.

Heaviest stack on every diagnostic report (96–98% of samples):

```
_dispatch_workloop_worker_thread → _dispatch_block_async_invoke2 →
  __CFURLCache::CreateAndStoreCacheNode → write
```

Full diagnostic data and analysis in #2004.

## What changed

`ClusterStateService` now defaults to an ephemeral, non-caching
`URLSession` instead of `URLSession.shared`. Cluster-state responses are
time-sensitive and small; nothing benefits from being cached on disk.

```swift
private static func makeNonCachingSession() -> URLSession {
    let config = URLSessionConfiguration.ephemeral
    config.urlCache = nil
    config.requestCachePolicy = .reloadIgnoringLocalCacheData
    return URLSession(configuration: config)
}
```

The existing per-request `request.cachePolicy =
.reloadIgnoringLocalCacheData` calls are kept as defense in depth — they
only affect read behavior, but harmless to leave alongside the
session-level config.

## Scope

- **Behavioral**: none. Polled requests still go out at the same
cadence; responses still parse the same; no semantic change to any API
surface.
- **Test injection**: the `session:` parameter remains in `init`, so
tests can still inject a custom mock session unchanged.
- **`BugReportService` and other `URLSession.shared` callers**:
untouched. If maintainers prefer an app-wide URLCache disable instead,
happy to switch the approach (issue body has the alternative spelled
out).

## Verification

Verified locally that compiling EXO with this change produces a working
menubar app and `ClusterStateService` continues to fetch state
correctly. After ~30 min of idle polling, no new entries in
`/Library/Logs/DiagnosticReports/EXO_*.diag` and no growth in
`~/Library/Caches/exolabs.EXO/`.

## Test plan
- [ ] Build EXO from this branch on macOS 26.4
- [ ] Launch, let cluster state polling run for 30+ min
- [ ] Confirm no new microstackshot diagnostic reports
- [ ] Confirm `~/Library/Caches/exolabs.EXO/Cache.db*` does not grow

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-authored-by: Jordan Miller <jordan.d.miller@gmail.com>
2026-05-01 10:41:10 +00:00
ciaranbor 8dae3ecb9a A few targeted tweaks to address HF rate limits (#2009)
## Motivation

- exo bursts ~200 HF Hub-API requests on every cold start, blowing past
the anonymous 500-req/5-min budget.
- The existing retry loop catches 429 generically and gives up in ~3s —
well before HF's reset window.
- `file_meta` and `_download_file` had no 429 handling at all (became
`AssertionError`).
- Disk file-list cache was bypassed on every process restart.

## Changes

All in `src/exo/download/download_utils.py` + tests.

- Parse `t=` from HF's `RateLimit` header on 429; sleep `min(t, 300s) +
jitter`.
- Handle 429 at all three call sites (`_fetch_file_list`, `file_meta`,
`_download_file`).
- `n_attempts`: 3 → 5.
- Disk cache now primary across restarts (24h mtime TTL).
- `?recursive=true` instead of N+1 subdir walks.

## Why It Works

`t=<seconds>` is HF's "wait this long and you'll be unblocked" —
sleeping that long lets the window reset. Disk-cache-as-primary plus
recursive listing cuts cold-start Hub-API traffic by ~10×.

## Test Plan

### Manual Testing

MacBook Pro M1 Max. Tripped the real HF 429. Pre-fix: failed in 3.4s.
Post-fix: slept (HF returned `t=158`) and recovered.

### Automated Testing

- New `test_rate_limit_handling.py` (19 tests) — header parsing,
retry-loop behaviour, plus HTTP-level coverage that mocks aiohttp to
return a 429 and asserts each call site raises
`HuggingFaceRateLimitError(retry_after=52.0)`.
- New `TestFileListCacheTTL` in `test_offline_mode.py` — fresh cache
hits, stale cache refetches.
- 421 tests pass; basedpyright / ruff / nix fmt clean.
2026-04-30 18:06:15 +00:00
Alex CheemaandClaude Opus 4.7 fb12b403ea fix(app): tighten Share Bug Report prompt layout (#2008)
## Summary

Follow-ups to #2003 based on feedback that the Share Bug Report window
felt visually weighty: too much padding above and below, and a
description editor that invited an essay rather than a one-liner.

## Changes (one file)

`app/EXO/EXO/Views/BugReportWindowController.swift`:

- **Auto-size the window to its content.** Switched from `NSHostingView`
+ fixed `contentRect: 480x380` + SwiftUI `frame(minHeight: 320)` to
`NSHostingController` with `sizingOptions = [.preferredContentSize,
.minSize]`. The fixed-min combo was centering the form in dead vertical
space.
- **Smaller, lower-pressure editor.** Field is now labeled `Description
(optional)` with a placeholder hint (`What were you doing when it
broke?`) inside the editor. Editor height fixed at 72pt (was 120pt min).
Replaced the long lead-in paragraph and headline with a single one-line
caption between field and buttons: `Diagnostic logs will be uploaded
with your report.`
- **Tighter spacing.** Outer padding 20 -> 16, root spacing 16 -> 12,
prompting-section spacing 12 -> 8.
- **Remove em dash from copy.**

`BugReportService` and the menu wiring are unchanged.

## Test plan

- [ ] Click `Share Bug Report...` from the menu bar.
- [ ] The window opens centered and sized to its content (no big empty
bands top/bottom).
- [ ] Description editor is visibly compact, with the placeholder hint
showing when empty.
- [ ] The optional-ness is conveyed by the field label (no separate help
paragraph).
- [ ] Caption `Diagnostic logs will be uploaded with your report.`
appears in `.caption` style under the editor, above the buttons.
- [ ] Resize the window: persists across re-opens (frame autosave still
works).
- [ ] Send/Cancel/Try Again/Done flows behave the same as before.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-30 15:10:26 +01:00
Alex CheemaandClaude Opus 4.7 1606e63816 feat(app): open Share Bug Report in a dedicated window (#2003)
## Summary

- Adds a top-level **Share Bug Report…** menu item to the macOS popover
(between *Check for Updates* and *Quit*) with SF Symbol `ladybug`.
- Clicking it opens a dedicated resizable `NSWindow` ("Send a Bug
Report") that hosts the prompting / sending / success / failure flow.
- Removes the description-less duplicate from Settings → Debug Info, and
the dead `debugSection` it nominally lived behind.

## Why

PR #1959 added a user-description prompt to the bug-report flow, but its
trigger lived inside `ContentView.debugSection` — a view that's defined
but never rendered in the body. The path users actually hit was
`SettingsView.sendBugReportButton`, which called
`BugReportService.sendReport(isManual: true)` without ever passing
`userDescription`. So the description prompt was unreachable in the
built app.

## Approach

Per Apple HIG, an action that requires further input before completing
should open a dialog, not transform the menu inline. So:

- Add a top-level menu entry that ends in `…` (HIG: ellipsis indicates
"further input required").
- Move the prompting/sending/success/failure state machine into a
standalone `BugReportWindowController` modeled after the existing
`SettingsWindowController`.
- Single-instance window with frame-autosave name, sensible
`contentMinSize`, resizable, native button layout (`.cancelAction` /
`.defaultAction` keyboard shortcuts), light/dark-mode-correct
`.textBackgroundColor` and `.separatorColor`.
- Auto-focus the description field on open. `Try Again` from failure,
`Open GitHub Issue` + `Done` from success.

## Files

- `app/EXO/EXO/Views/BugReportWindowController.swift` (new) — controller
+ view.
- `app/EXO/EXO/EXOApp.swift` — wire `BugReportWindowController` as a
`@StateObject` and inject as environment object.
- `app/EXO/EXO/ContentView.swift` — replace inline state machine with
menu item that calls `bugReportWindowController.open()`. Remove
now-unused state, helpers, and dead `debugSection`.
- `app/EXO/EXO/Views/SettingsView.swift` — remove duplicate
`sendBugReportButton`, `sendBugReport()`, and related `@State`. Section
"Debug Info" keeps Thunderbolt / interface / RDMA info.

`BugReportService` is unchanged.

## Test plan

- [ ] Open the menu-bar popover → confirm **Share Bug Report…** appears
between *Check for Updates* and *Quit*, with a ladybug icon.
- [ ] Click it → a window titled "Send a Bug Report" appears, centered,
with the description editor focused.
- [ ] Resize the window → size persists across re-opens (frame
autosave).
- [ ] Type a description, press Return → upload succeeds, success card
with **Open GitHub Issue** + **Done** appears.
- [ ] Click **Open GitHub Issue** → browser opens with the description
pre-filled into the issue template.
- [ ] Send with empty description → upload still succeeds.
- [ ] Press Esc from the prompting state → window closes.
- [ ] On failure (e.g., offline) → error card with **Try Again** +
**Close** appears; Try Again returns to the editor with the description
preserved.
- [ ] Open the Settings window → Debug Info section is unchanged except
the Send Bug Report button is gone.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-30 13:05:07 +01:00
84 changed files with 5388 additions and 5096 deletions

No files matched your search

Generated
+2338 -2254
View File
File diff suppressed because it is too large. Load diff
+16 -15
View File
@@ -1,6 +1,6 @@
[workspace]
resolver = "3"
members = ["rust/networking", "rust/exo_pyo3_bindings", "rust/util"]
members = ["rust/exo_net", "rust/networking"]
[workspace.package]
version = "0.0.1"
@@ -20,30 +20,31 @@ opt-level = 3
[workspace.dependencies]
## Crate members as common dependencies
networking = { path = "rust/networking" }
util = { path = "rust/util" }
# Macro dependecies
extend = "1.2"
delegate = "0.13"
# Utility dependencies
keccak-const = "0.2"
# Async dependencies
async-stream = "0.3"
tokio = "1.46"
futures-lite = "2.6.1"
futures-timer = "3.0"
# Data structures
either = "1.15"
# Tracing/logging
log = "0.4"
env_logger = "0.11.10"
# networking
libp2p = "0.56"
libp2p-tcp = "0.44"
zenoh = "=1.9.0"
async-stream = "0.3.6"
netwatcher = "0.6.0"
parking_lot = "0.12.5"
pin-project = "1.1.10"
pyo3 = "0.27.2"
pyo3-async-runtimes = "0.27.0"
pyo3-log = "0.13.2"
pyo3-stub-gen = "0.22.2"
rand = "0.10.1"
serde_json = "1.0.149"
tracing = "0.1.44"
zenoh-plugin-storage-manager = { version = "=1.9.0", default-features = false }
zenoh-plugin-trait = "=1.9.0"
[workspace.lints.rust]
static_mut_refs = "warn" # Or use "warn" instead of deny
+8 -227
View File
@@ -16,22 +16,13 @@ struct ContentView: View {
@EnvironmentObject private var updater: SparkleUpdater
@EnvironmentObject private var thunderboltBridgeService: ThunderboltBridgeService
@EnvironmentObject private var settingsWindowController: SettingsWindowController
@EnvironmentObject private var bugReportWindowController: BugReportWindowController
@State private var focusedNode: NodeViewModel?
@State private var deletingInstanceIDs: Set<String> = []
@State private var showAllNodes = false
@State private var showAllInstances = false
@State private var baseURLCopied = false
@State private var showAdvanced = false
@State private var showDebugInfo = false
private enum BugReportPhase: Equatable {
case idle
case prompting
case sending(String)
case success(String)
case failure(String)
}
@State private var bugReportPhase: BugReportPhase = .idle
@State private var bugReportUserDescription: String = ""
@State private var uninstallInProgress = false
@State private var pendingNamespace: String = ""
@State private var pendingHFToken: String = ""
@@ -294,6 +285,13 @@ struct ContentView: View {
) {
updater.checkForUpdates()
}
HoverButton(
title: "Share Bug Report…",
tint: .primary,
trailingSystemImage: "ladybug"
) {
bugReportWindowController.open()
}
.padding(.bottom, 8)
HoverButton(title: "Quit", tint: .secondary) {
controller.stop()
@@ -477,40 +475,6 @@ struct ContentView: View {
}
}
private var debugSection: some View {
VStack(alignment: .leading, spacing: 4) {
HoverButton(
title: "Debug Info",
tint: .primary,
trailingSystemImage: showDebugInfo ? "chevron.up" : "chevron.down",
small: true
) {
showDebugInfo.toggle()
}
if showDebugInfo {
VStack(alignment: .leading, spacing: 4) {
Text("Version: \(buildTag)")
.font(.caption2)
.foregroundColor(.secondary)
Text("Commit: \(buildCommit)")
.font(.caption2)
.foregroundColor(.secondary)
Text(thunderboltStatusText)
.font(.caption2)
.foregroundColor(thunderboltStatusColor)
clusterThunderboltBridgeView
interfaceIpList
rdmaStatusView
sendBugReportButton
.padding(.top, 6)
}
.padding(.leading, 8)
.transition(.opacity)
}
}
.animation(.easeInOut(duration: 0.25), value: showDebugInfo)
}
private var rdmaStatusView: some View {
let rdmaStatuses = stateService.latestSnapshot?.nodeRdmaCtl ?? [:]
let localNodeId = stateService.localNodeId
@@ -559,127 +523,6 @@ struct ContentView: View {
}
}
private var sendBugReportButton: some View {
VStack(alignment: .leading, spacing: 6) {
switch bugReportPhase {
case .idle:
Button {
bugReportPhase = .prompting
bugReportUserDescription = ""
} label: {
HStack {
Text("Send Bug Report")
.font(.caption)
.fontWeight(.semibold)
Spacer()
}
.padding(.vertical, 6)
.padding(.horizontal, 8)
.background(
RoundedRectangle(cornerRadius: 6)
.fill(Color.accentColor.opacity(0.12))
)
}
.buttonStyle(.plain)
case .prompting:
VStack(alignment: .leading, spacing: 6) {
VStack(alignment: .leading, spacing: 2) {
Text("Tell us what went wrong (optional)")
.font(.caption2)
.foregroundColor(.secondary)
Text(
"A quick description of what you were doing and what happened helps us track down the bug for you."
)
.font(.caption2)
.foregroundColor(.secondary)
.opacity(0.8)
.fixedSize(horizontal: false, vertical: true)
}
TextEditor(text: $bugReportUserDescription)
.font(.caption2)
.frame(height: 60)
.overlay(
RoundedRectangle(cornerRadius: 4)
.stroke(Color.secondary.opacity(0.3), lineWidth: 1)
)
HStack(spacing: 8) {
Button("Send") {
Task {
await sendBugReport()
}
}
.font(.caption2)
.buttonStyle(.borderedProminent)
.controlSize(.small)
Button("Cancel") {
bugReportPhase = .idle
}
.font(.caption2)
.buttonStyle(.bordered)
.controlSize(.small)
}
}
.padding(8)
.background(
RoundedRectangle(cornerRadius: 6)
.fill(Color.accentColor.opacity(0.06))
)
case .sending(let message):
HStack(spacing: 6) {
ProgressView()
.scaleEffect(0.6)
Text(message)
.font(.caption2)
.foregroundColor(.secondary)
}
case .success(let message):
VStack(alignment: .leading, spacing: 6) {
Text(message)
.font(.caption2)
.foregroundColor(.secondary)
.fixedSize(horizontal: false, vertical: true)
Button {
openGitHubIssue()
} label: {
HStack(spacing: 4) {
Image(systemName: "arrow.up.right.square")
.imageScale(.small)
Text("Create GitHub Issue")
.font(.caption2)
}
}
.buttonStyle(.bordered)
.controlSize(.small)
Button("Done") {
bugReportPhase = .idle
bugReportUserDescription = ""
}
.font(.caption2)
.buttonStyle(.plain)
.foregroundColor(.secondary)
}
case .failure(let message):
VStack(alignment: .leading, spacing: 4) {
Text(message)
.font(.caption2)
.foregroundColor(.red)
.fixedSize(horizontal: false, vertical: true)
Button("Dismiss") {
bugReportPhase = .idle
}
.font(.caption2)
.buttonStyle(.plain)
.foregroundColor(.secondary)
}
}
}
.animation(.easeInOut(duration: 0.2), value: bugReportPhase)
}
private var processToggleBinding: Binding<Bool> {
Binding(
get: {
@@ -720,61 +563,6 @@ struct ContentView: View {
)
}
private func sendBugReport() async {
bugReportPhase = .sending("Collecting logs...")
let service = BugReportService()
let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines)
do {
let outcome = try await service.sendReport(
isManual: true,
userDescription: description.isEmpty ? nil : description
)
if outcome.success {
bugReportPhase = .success(outcome.message)
} else {
bugReportPhase = .failure(outcome.message)
}
} catch {
bugReportPhase = .failure(error.localizedDescription)
}
}
private func openGitHubIssue() {
let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines)
var bodyParts: [String] = []
bodyParts.append("## Describe the bug")
bodyParts.append("")
if !description.isEmpty {
bodyParts.append(description)
} else {
bodyParts.append("A clear and concise description of what the bug is.")
}
bodyParts.append("")
bodyParts.append("## Environment")
bodyParts.append("")
bodyParts.append("- macOS Version: \(ProcessInfo.processInfo.operatingSystemVersionString)")
bodyParts.append("- EXO Version: \(buildTag) (\(buildCommit))")
bodyParts.append("")
bodyParts.append("## Additional context")
bodyParts.append("")
bodyParts.append("A bug report with diagnostic logs was submitted via the app.")
let body = bodyParts.joined(separator: "\n")
var components = URLComponents(string: "https://github.com/exo-explore/exo/issues/new")!
components.queryItems = [
URLQueryItem(name: "template", value: "bug_report.md"),
URLQueryItem(name: "title", value: "[BUG] "),
URLQueryItem(name: "body", value: body),
URLQueryItem(name: "labels", value: "bug"),
]
if let url = components.url {
NSWorkspace.shared.open(url)
}
}
private func showUninstallConfirmationAlert() {
let alert = NSAlert()
alert.messageText = "Uninstall EXO"
@@ -857,13 +645,6 @@ struct ContentView: View {
}
}
private var buildTag: String {
Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown"
}
private var buildCommit: String {
Bundle.main.infoDictionary?["EXOBuildCommit"] as? String ?? "unknown"
}
}
private struct HoverButton: View {
+3
View File
@@ -22,6 +22,7 @@ struct EXOApp: App {
@StateObject private var updater: SparkleUpdater
@StateObject private var thunderboltBridgeService: ThunderboltBridgeService
@StateObject private var settingsWindowController: SettingsWindowController
@StateObject private var bugReportWindowController: BugReportWindowController
private let terminationObserver: TerminationObserver
private let firstLaunchPopout = FirstLaunchPopout()
private let ciContext = CIContext(options: nil)
@@ -46,6 +47,7 @@ struct EXOApp: App {
let thunderboltBridge = ThunderboltBridgeService(clusterStateService: service)
_thunderboltBridgeService = StateObject(wrappedValue: thunderboltBridge)
_settingsWindowController = StateObject(wrappedValue: SettingsWindowController())
_bugReportWindowController = StateObject(wrappedValue: BugReportWindowController())
enableLaunchAtLoginIfNeeded()
// Install LaunchDaemon to disable Thunderbolt Bridge on startup (prevents network loops)
NetworkSetupHelper.promptAndInstallIfNeeded()
@@ -66,6 +68,7 @@ struct EXOApp: App {
.environmentObject(updater)
.environmentObject(thunderboltBridgeService)
.environmentObject(settingsWindowController)
.environmentObject(bugReportWindowController)
} label: {
menuBarIcon
.onReceive(controller.$isFirstLaunchReady) { ready in
+18 -1
View File
@@ -17,7 +17,7 @@ final class ClusterStateService: ObservableObject {
init(
baseURL: URL = URL(string: "http://127.0.0.1:52415")!,
session: URLSession = .shared
session: URLSession = ClusterStateService.makeNonCachingSession()
) {
self.baseURL = baseURL
self.endpoint = baseURL.appendingPathComponent("state")
@@ -27,6 +27,23 @@ final class ClusterStateService: ObservableObject {
self.decoder = decoder
}
/// `URLSession.shared` carries an on-disk `URLCache` that persists every
/// response body under `~/Library/Caches/exolabs.EXO/`. We poll `/state`
/// at 2 Hz from `startPolling`, so leaving the shared cache attached
/// dirties ~500620 KB/sec of file-backed memory and trips macOS's
/// per-process `disk writes` resource limit (microstackshot reports
/// observed on M3 Ultra producing GBs of cached responses per hour).
/// Cluster-state polling responses are time-sensitive and small; they
/// gain nothing from being cached on disk. Use an ephemeral session
/// with `urlCache = nil` so neither response bodies nor metadata
/// touch disk.
private static func makeNonCachingSession() -> URLSession {
let config = URLSessionConfiguration.ephemeral
config.urlCache = nil
config.requestCachePolicy = .reloadIgnoringLocalCacheData
return URLSession(configuration: config)
}
func startPolling(interval: TimeInterval = 0.5) {
stopPolling()
Task {
@@ -0,0 +1,242 @@
import AppKit
import SwiftUI
/// Manages a standalone window for the bug-report flow.
/// Ensures only one instance exists and brings it to front on repeated opens.
@MainActor
final class BugReportWindowController: ObservableObject {
private var window: NSWindow?
func open() {
if let existing = window, existing.isVisible {
existing.makeKeyAndOrderFront(nil)
NSApp.activate()
return
}
let view = BugReportView(onDismiss: { [weak self] in
self?.window?.close()
})
let hostingController = NSHostingController(rootView: view)
hostingController.sizingOptions = [.preferredContentSize, .minSize]
let newWindow = NSWindow(contentViewController: hostingController)
newWindow.styleMask = [.titled, .closable, .resizable]
newWindow.title = "Send a Bug Report"
newWindow.center()
newWindow.setFrameAutosaveName("ExoBugReportWindow")
newWindow.isReleasedWhenClosed = false
newWindow.makeKeyAndOrderFront(nil)
NSApp.activate()
window = newWindow
}
}
private struct BugReportView: View {
fileprivate enum Phase: Equatable {
case prompting
case sending(String)
case success(String)
case failure(String)
}
let onDismiss: () -> Void
@State private var phase: Phase = .prompting
@State private var userDescription: String = ""
@FocusState private var descriptionFocused: Bool
var body: some View {
VStack(alignment: .leading, spacing: 12) {
switch phase {
case .prompting:
promptingView
case .sending(let message):
sendingView(message: message)
case .success(let message):
successView(message: message)
case .failure(let message):
failureView(message: message)
}
}
.padding(16)
.frame(minWidth: 380)
.animation(.easeInOut(duration: 0.2), value: phase)
.onAppear { descriptionFocused = true }
}
private var promptingView: some View {
VStack(alignment: .leading, spacing: 8) {
Text("Description (optional)")
.font(.subheadline)
.foregroundColor(.secondary)
ZStack(alignment: .topLeading) {
if userDescription.isEmpty {
Text("What were you doing when it broke?")
.font(.body)
.foregroundColor(Color(nsColor: .placeholderTextColor))
.padding(.horizontal, 10)
.padding(.vertical, 8)
.allowsHitTesting(false)
}
TextEditor(text: $userDescription)
.font(.body)
.scrollContentBackground(.hidden)
.padding(4)
.frame(height: 72)
.focused($descriptionFocused)
}
.background(
RoundedRectangle(cornerRadius: 6)
.fill(Color(nsColor: .textBackgroundColor))
)
.overlay(
RoundedRectangle(cornerRadius: 6)
.strokeBorder(Color(nsColor: .separatorColor), lineWidth: 1)
)
Text("Diagnostic logs will be uploaded with your report.")
.font(.caption)
.foregroundColor(.secondary)
HStack {
Spacer()
Button("Cancel") { onDismiss() }
.keyboardShortcut(.cancelAction)
Button("Send") {
Task { await send() }
}
.keyboardShortcut(.defaultAction)
}
.padding(.top, 4)
}
}
private func sendingView(message: String) -> some View {
VStack(alignment: .leading, spacing: 12) {
HStack(spacing: 10) {
ProgressView().controlSize(.small)
Text(message)
.foregroundColor(.secondary)
}
HStack {
Spacer()
Button("Cancel") { onDismiss() }
.keyboardShortcut(.cancelAction)
.disabled(true)
Button("Send") {}
.disabled(true)
}
}
}
private func successView(message: String) -> some View {
VStack(alignment: .leading, spacing: 12) {
HStack(alignment: .top, spacing: 10) {
Image(systemName: "checkmark.circle.fill")
.foregroundColor(.green)
.font(.title2)
Text(message)
.fixedSize(horizontal: false, vertical: true)
}
HStack {
Button {
openGitHubIssue()
} label: {
HStack(spacing: 4) {
Image(systemName: "arrow.up.right.square")
Text("Open GitHub Issue")
}
}
Spacer()
Button("Done") { onDismiss() }
.keyboardShortcut(.defaultAction)
}
}
}
private func failureView(message: String) -> some View {
VStack(alignment: .leading, spacing: 12) {
HStack(alignment: .top, spacing: 10) {
Image(systemName: "exclamationmark.triangle.fill")
.foregroundColor(.orange)
.font(.title2)
Text(message)
.fixedSize(horizontal: false, vertical: true)
}
HStack {
Spacer()
Button("Try Again") {
phase = .prompting
}
Button("Close") { onDismiss() }
.keyboardShortcut(.defaultAction)
}
}
}
private func send() async {
phase = .sending("Collecting logs and uploading…")
let service = BugReportService()
let description = userDescription.trimmingCharacters(in: .whitespacesAndNewlines)
do {
let outcome = try await service.sendReport(
isManual: true,
userDescription: description.isEmpty ? nil : description
)
if outcome.success {
phase = .success(outcome.message)
} else {
phase = .failure(outcome.message)
}
} catch {
phase = .failure(error.localizedDescription)
}
}
private func openGitHubIssue() {
let description = userDescription.trimmingCharacters(in: .whitespacesAndNewlines)
var bodyParts: [String] = []
bodyParts.append("## Describe the bug")
bodyParts.append("")
if !description.isEmpty {
bodyParts.append(description)
} else {
bodyParts.append("A clear and concise description of what the bug is.")
}
bodyParts.append("")
bodyParts.append("## Environment")
bodyParts.append("")
bodyParts.append("- macOS Version: \(ProcessInfo.processInfo.operatingSystemVersionString)")
bodyParts.append("- EXO Version: \(buildTag) (\(buildCommit))")
bodyParts.append("")
bodyParts.append("## Additional context")
bodyParts.append("")
bodyParts.append("A bug report with diagnostic logs was submitted via the app.")
let body = bodyParts.joined(separator: "\n")
var components = URLComponents(string: "https://github.com/exo-explore/exo/issues/new")!
components.queryItems = [
URLQueryItem(name: "template", value: "bug_report.md"),
URLQueryItem(name: "title", value: "[BUG] "),
URLQueryItem(name: "body", value: body),
URLQueryItem(name: "labels", value: "bug"),
]
if let url = components.url {
NSWorkspace.shared.open(url)
}
}
private var buildTag: String {
Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown"
}
private var buildCommit: String {
Bundle.main.infoDictionary?["EXOBuildCommit"] as? String ?? "unknown"
}
}
-46
View File
@@ -21,8 +21,6 @@ struct SettingsView: View {
@State private var pendingReadOnlyModelsDirs: String = ""
@State private var pendingCustomEnvironmentVariables: [CustomEnvironmentVariable] = []
@State private var needsRestart = false
@State private var bugReportInFlight = false
@State private var bugReportMessage: String?
@State private var uninstallInProgress = false
var body: some View {
@@ -202,8 +200,6 @@ struct SettingsView: View {
VStack(alignment: .leading, spacing: 2) {
rdmaStatusView
}
sendBugReportButton
}
Section("Danger Zone") {
@@ -504,50 +500,8 @@ struct SettingsView: View {
}
}
private var sendBugReportButton: some View {
VStack(alignment: .leading, spacing: 4) {
Button {
Task {
await sendBugReport()
}
} label: {
HStack {
if bugReportInFlight {
ProgressView()
.scaleEffect(0.6)
}
Text("Send Bug Report")
.font(.caption)
.fontWeight(.semibold)
Spacer()
}
}
.disabled(bugReportInFlight)
if let message = bugReportMessage {
Text(message)
.font(.caption2)
.foregroundColor(.secondary)
.fixedSize(horizontal: false, vertical: true)
}
}
}
// MARK: - Actions
private func sendBugReport() async {
bugReportInFlight = true
bugReportMessage = "Collecting logs..."
let service = BugReportService()
do {
let outcome = try await service.sendReport(isManual: true)
bugReportMessage = outcome.message
} catch {
bugReportMessage = error.localizedDescription
}
bugReportInFlight = false
}
private func showUninstallConfirmationAlert() {
let alert = NSAlert()
alert.messageText = "Uninstall EXO"
+3
View File
@@ -3435,6 +3435,7 @@
>
<li>Connect nodes with TB5 cables</li>
<li>Boot to Recovery (hold power 10s → Options)</li>
<li>Open Terminal from the Utilities menu</li>
<li>
Run
<code class="text-yellow-300 bg-yellow-400/10 px-1 rounded"
@@ -4822,6 +4823,7 @@
>
<li>Connect nodes with TB5 cables</li>
<li>Boot to Recovery (hold power 10s → Options)</li>
<li>Open Terminal from the Utilities menu</li>
<li>
Run
<code class="text-yellow-300 bg-yellow-400/10 px-1 rounded"
@@ -4968,6 +4970,7 @@
>
<li>Connect nodes with TB5 cables</li>
<li>Boot to Recovery (hold power 10s → Options)</li>
<li>Open Terminal from the Utilities menu</li>
<li>
Run
<code
Generated
+6 -6
View File
@@ -47,11 +47,11 @@
"rust-analyzer-src": "rust-analyzer-src"
},
"locked": {
"lastModified": 1775807984,
"narHash": "sha256-Redoe3D9zGN5I9QPHWL9vfMVQBehY1fKsMiRXQ83X3w=",
"lastModified": 1777708550,
"narHash": "sha256-Qif3UXT0l5OQq8H9pRWt4/ia4gF48MWK2oHKL8uVx8U=",
"owner": "nix-community",
"repo": "fenix",
"rev": "fcf90c0c4d368b2ca917a7afa6d08e98a397e5fd",
"rev": "74c1591efaff494756b8d35ebe357c6c2bbdca96",
"type": "github"
},
"original": {
@@ -218,11 +218,11 @@
"rust-analyzer-src": {
"flake": false,
"locked": {
"lastModified": 1775745684,
"narHash": "sha256-8MbfLwd60FNa8dRFkjE+G3TT/x21G3Rsplm1bMBQUtU=",
"lastModified": 1777639980,
"narHash": "sha256-6d7Hdurvbjc5uwJuc0YiK7rZBGj6Gs3uzfBFcTs+xCc=",
"owner": "rust-lang",
"repo": "rust-analyzer",
"rev": "64ddb549bc9a70d011328746fa46a8883f937b6b",
"rev": "64cdaeb06f69b6b769a492edd88b022ae88e8ca2",
"type": "github"
},
"original": {
+2 -2
View File
@@ -110,7 +110,7 @@
nixpkgs-fmt.enable = true;
ruff-format = {
enable = true;
excludes = [ "rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi" ];
excludes = [ "rust/exo_net/exo_net.pyi" ];
};
rustfmt = {
enable = true;
@@ -146,7 +146,7 @@
config.treefmt.build.wrapper
# PYTHON
self'.packages.editableVenv
#self'.packages.editableVenv
uv
# RUST
+2 -2
View File
@@ -22,8 +22,8 @@ sync-clean:
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
uv sync --reinstall-package exo_pyo3_bindings
cargo run --bin stub_gen
uv sync --reinstall-package exo_net
build-dashboard:
#!/usr/bin/env bash
+3 -3
View File
@@ -15,7 +15,7 @@ dependencies = [
"huggingface-hub>=1.8.0",
"psutil>=7.0.0",
"loguru>=0.7.3",
"exo-pyo3-bindings", # rust bindings
"exo-net", # rust bindings
"anyio==4.11.0",
"mlx==0.31.2; sys_platform == 'darwin'",
"mlx-lm; sys_platform=='darwin'",
@@ -75,10 +75,10 @@ cuda13 = [
###
[tool.uv.workspace]
members = ["rust/exo_pyo3_bindings", "bench"]
members = ["rust/exo_net", "bench"]
[tool.uv.sources]
exo-pyo3-bindings = { workspace = true }
exo-net = { workspace = true }
mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/deepseek-v4" }
torch = [
+6 -5
View File
@@ -35,17 +35,18 @@ let
# Replace workspace exo_pyo3_bindings with Nix-built wheel.
# Preserve passthru so mkVirtualEnv can resolve dependency groups.
# Copy .pyi stub + py.typed marker so basedpyright can find the types.
exo-pyo3-bindings = pkgs.stdenv.mkDerivation {
pname = "exo-pyo3-bindings";
exo-net = pkgs.stdenv.mkDerivation {
pname = "exo-net";
version = "0.1.0";
src = self'.packages.exo_pyo3_bindings;
src = self'.packages.exo-net;
# Install from pre-built wheel
nativeBuildInputs = [ final.pyprojectWheelHook ];
dontStrip = true;
passthru = prev.exo-pyo3-bindings.passthru or { };
postInstall = ''
local siteDir=$out/${final.python.sitePackages}/exo_pyo3_bindings
cp ${inputs.self}/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi $siteDir/
local siteDir=$out/${final.python.sitePackages}/exo_net
cp ${inputs.self}/rust/exo_net/exo_net.pyi $siteDir/
touch $siteDir/py.typed
'';
};
Symlink
+1
View File
@@ -0,0 +1 @@
/nix/store/3945qvxy6hla1aa6dkzlcv99d2d4bw56-exo
+52
View File
@@ -0,0 +1,52 @@
[package]
name = "exo_net"
version = { workspace = true }
edition = { workspace = true }
publish = false
[lib]
doctest = false
path = "src/lib.rs"
name = "exo_net"
# "cdylib" needed to produce shared library for Python to import
# "rlib" needed for stub-gen to run
crate-type = ["cdylib", "rlib"]
[[bin]]
path = "src/bin/stub_gen.rs"
name = "stub_gen"
doc = false
[lints]
workspace = true
[dependencies]
networking.workspace = true
extend.workspace = true
# interop
pyo3 = { workspace = true, features = ["experimental-async"] }
pyo3-stub-gen.workspace = true
pyo3-async-runtimes = { workspace = true, features = [
"attributes",
"tokio-runtime",
"testing",
] }
pyo3-log.workspace = true
# async runtime
tokio = { workspace = true, features = ["full", "tracing"] }
futures-lite.workspace = true
pin-project.workspace = true
# Tracing
log.workspace = true
env_logger.workspace = true
# Networking
zenoh.workspace = true
rand.workspace = true
serde_json.workspace = true
parking_lot.workspace = true
tracing.workspace = true
File renamed without changes.
+76
View File
@@ -0,0 +1,76 @@
# This file is automatically generated by pyo3_stub_gen
# ruff: noqa: E501, F401, F403, F405
import builtins
import collections.abc
import typing
__all__ = [
"NetReceiver",
"NetSender",
"NetworkingHandle",
"PyFromSwarm",
"PySession",
"StateProxy",
]
@typing.final
class NetReceiver:
def recv(self) -> collections.abc.Awaitable[bytes | None]: ...
@typing.final
class NetSender:
def send(self, data: bytes) -> collections.abc.Awaitable[bool]: ...
@typing.final
class NetworkingHandle:
@staticmethod
def new(identity: bytes, bootstrap_peers: typing.Sequence[builtins.str], listen_port: builtins.int) -> tuple[NetworkingHandle, PySession]: ...
async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
r"""
Subscribe to a `GossipSub` topic.
Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
"""
async def gossipsub_unsubscribe(self, topic: builtins.str) -> builtins.bool:
r"""
Unsubscribes from a `GossipSub` topic.
Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
"""
async def gossipsub_publish(self, topic: builtins.str, data: bytes) -> None:
r"""
Publishes a message with multiple topics to the `GossipSub` network.
If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
"""
async def recv(self) -> PyFromSwarm: ...
class PyFromSwarm:
@typing.final
class Connection(PyFromSwarm):
__match_args__ = ("connected",)
@property
def connected(self) -> builtins.bool: ...
def __new__(cls, connected: builtins.bool) -> PyFromSwarm.Connection: ...
@typing.final
class Message(PyFromSwarm):
__match_args__ = ("topic", "data",)
@property
def topic(self) -> builtins.str: ...
@property
def data(self) -> bytes: ...
def __new__(cls, topic: builtins.str, data: bytes) -> PyFromSwarm.Message: ...
@typing.final
class PySession:
def net_receiver(self, key: builtins.str) -> NetReceiver: ...
def net_sender(self, key: builtins.str) -> NetSender: ...
def state_proxy(self) -> StateProxy: ...
@typing.final
class StateProxy:
def snapshot(self) -> collections.abc.Awaitable[str]: ...
+24
View File
@@ -0,0 +1,24 @@
# This file is automatically generated by pyo3_stub_gen
# ruff: noqa: E501, F401
import builtins
import collections.abc
import typing
@typing.final
class NetReceiver: ...
@typing.final
class NetSender: ...
@typing.final
class PySession:
@staticmethod
def init() -> collections.abc.Awaitable[PySession]: ...
def net_receiver(self, key: builtins.str) -> NetReceiver: ...
def net_sender(self, key: builtins.str) -> NetSender: ...
def state_proxy(self) -> StateProxy: ...
@typing.final
class StateProxy:
def snapshot(self) -> collections.abc.Awaitable[str]: ...
@@ -3,24 +3,22 @@ requires = ["maturin>=1.0,<2.0"]
build-backend = "maturin"
[project]
name = "exo_pyo3_bindings"
version = "0.2.1"
name = "exo_net"
version = "0.3.0"
description = "Add your description here"
readme = "README.md"
authors = [
{ name = "Andrei Cravtov", email = "the.andrei.cravtov@gmail.com" },
{ name = "Evan Quiney", email = "evanev7@gmail.com" },
{ name = "Andrei Cravtov", email = "the.andrei.cravtov@gmail.com" },
]
requires-python = ">=3.13"
dependencies = []
[dependency-groups]
dev = ["exo_pyo3_bindings", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"]
dev = ["exo-net", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"]
[tool.maturin]
#purelib = true
#python-source = "python"
module-name = "exo_pyo3_bindings"
module-name = "exo_net"
features = ["pyo3/extension-module", "pyo3/experimental-async"]
[tool.pytest.ini_options]
@@ -2,7 +2,7 @@ use pyo3_stub_gen::Result;
fn main() -> Result<()> {
env_logger::Builder::from_env(env_logger::Env::default().filter_or("RUST_LOG", "info")).init();
let stub = exo_pyo3_bindings::stub_info()?;
let stub = exo_net::stub_info()?;
stub.generate()?;
Ok(())
}
@@ -5,21 +5,21 @@
//!
mod allow_threading;
mod ident;
// mod ident;
mod networking;
mod point_to_point;
mod session;
mod state;
use crate::ident::PyKeypair;
use crate::networking::networking_submodule;
use crate::point_to_point::{NetReceiver, NetSender};
use crate::session::PySession;
use crate::state::StateProxy;
use pyo3::prelude::PyModule;
use pyo3::types::PyModuleMethods;
use pyo3::{Bound, PyResult, pyclass, pymodule};
use pyo3::{Bound, PyResult, pymodule};
use pyo3_stub_gen::define_stub_info_gatherer;
/// Namespace for all the constants used by this crate.
pub(crate) mod r#const {
pub const MPSC_CHANNEL_SIZE: usize = 1024;
}
/// Namespace for crate-wide extension traits/methods
pub(crate) mod ext {
use crate::allow_threading::AllowThreads;
@@ -151,7 +151,7 @@ pub(crate) mod ext {
/// A Python module implemented in Rust. The name of this function must match
/// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to
/// import the module.
#[pymodule(name = "exo_pyo3_bindings")]
#[pymodule(name = "exo_net")]
fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
// install logger
pyo3_log::init();
@@ -162,7 +162,12 @@ fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
// TODO: for now this is all NOT a submodule, but figure out how to make the submodule system
// work with maturin, where the types generate correctly, in the right folder, without
// too many importing issues...
m.add_class::<PyKeypair>()?;
// m.add_class::<PyKeypair>()?;
// networking_submodule(m)?;
m.add_class::<StateProxy>()?;
m.add_class::<PySession>()?;
m.add_class::<NetReceiver>()?;
m.add_class::<NetSender>()?;
networking_submodule(m)?;
// top-level constructs
+192
View File
@@ -0,0 +1,192 @@
use std::pin::Pin;
use std::sync::Arc;
use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
use crate::ext::{ResultExt as _, TokioMpscSenderExt as _};
use crate::session::PySession;
use futures_lite::{Stream, StreamExt as _};
use networking::swarm::{FromSwarm, ToSwarm, create_swarm};
use pyo3::exceptions::{PyRuntimeError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::PyBytes;
use pyo3::{Bound, Py, PyAny, PyErr, PyResult, Python, pymethods};
use pyo3_stub_gen::derive::{
gen_methods_from_python, gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods,
};
use tokio::sync::{Mutex, mpsc, oneshot};
#[gen_stub_pyclass]
#[pyclass(name = "NetworkingHandle")]
struct PyNetworkingHandle {
// channels
pub to_swarm: mpsc::Sender<ToSwarm>,
pub swarm: Arc<Mutex<Pin<Box<dyn Stream<Item = FromSwarm> + Send>>>>,
}
#[gen_stub_pyclass_complex_enum]
#[pyclass]
enum PyFromSwarm {
Connection { connected: bool },
Message { topic: String, data: Py<PyBytes> },
}
impl From<FromSwarm> for PyFromSwarm {
fn from(value: FromSwarm) -> Self {
match value {
FromSwarm::Discovered {} => Self::Connection { connected: true },
FromSwarm::Expired {} => Self::Connection { connected: false },
FromSwarm::Message { topic, data } => Self::Message {
topic: topic,
data: data.pybytes(),
},
}
}
}
#[gen_stub_pymethods]
#[pymethods]
impl PyNetworkingHandle {
// NOTE: `async fn`s here that use `.await` will wrap the future in `.allow_threads_py()`
// immediately beforehand to release the interpreter.
// SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
// ---- Lifecycle management methods ----
#[staticmethod]
fn new<'py>(
identity: Bound<'py, PyBytes>,
bootstrap_peers: Vec<String>,
listen_port: u16,
) -> PyResult<(PyNetworkingHandle, PySession)> {
// create communication channels
let (to_swarm, from_client) = mpsc::channel(1024);
// get identity
let identity = u128::from_le_bytes(
identity
.extract::<'_, Vec<u8>>()?
.try_into()
.map_err(|_| PyValueError::new_err("invalid identity bytes"))?,
);
// create networking swarm (within tokio context!! or it crashes)
let swarm = pyo3_async_runtimes::tokio::get_runtime()
.block_on(create_swarm(
identity,
from_client,
bootstrap_peers,
listen_port,
))
.pyerr()?;
let session = swarm.session.z.clone();
Ok((
PyNetworkingHandle {
swarm: Arc::new(Mutex::new(swarm.into_stream())),
to_swarm,
},
PySession { session },
))
}
#[gen_stub(skip)]
fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let swarm = Arc::clone(&self.swarm);
pyo3_async_runtimes::tokio::future_into_py(py, async move {
swarm
.try_lock()
.map_err(|_| PyRuntimeError::new_err("called recv twice concurrently"))?
.next()
.await
.ok_or(PyErr::receiver_channel_closed())
.map(PyFromSwarm::from)
})
}
// ---- Gossipsub management methods ----
/// Subscribe to a `GossipSub` topic.
///
/// Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
async fn gossipsub_subscribe(&self, topic: String) -> PyResult<bool> {
let (tx, rx) = oneshot::channel();
// send off request to subscribe
self.to_swarm
.send_py(ToSwarm::Subscribe {
topic,
result_sender: tx,
})
.allow_threads_py() // allow-threads-aware async call
.await?;
// wait for response & return any errors
rx.allow_threads_py() // allow-threads-aware async call
.await
.map_err(|_| PyErr::receiver_channel_closed())?
.pyerr()
}
/// Unsubscribes from a `GossipSub` topic.
///
/// Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
async fn gossipsub_unsubscribe(&self, topic: String) -> PyResult<bool> {
let (tx, rx) = oneshot::channel();
// send off request to unsubscribe
self.to_swarm
.send_py(ToSwarm::Unsubscribe {
topic,
result_sender: tx,
})
.allow_threads_py() // allow-threads-aware async call
.await?;
// wait for response & convert any errors
rx.allow_threads_py() // allow-threads-aware async call
.await
.map_err(|_| PyErr::receiver_channel_closed())
}
/// Publishes a message with multiple topics to the `GossipSub` network.
///
/// If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
async fn gossipsub_publish(&self, topic: String, data: Py<PyBytes>) -> PyResult<()> {
let (tx, rx) = oneshot::channel();
// send off request to subscribe
let data = Python::attach(|py| Vec::from(data.as_bytes(py)));
self.to_swarm
.send_py(ToSwarm::Publish {
topic,
data,
result_sender: tx,
})
.allow_threads_py() // allow-threads-aware async call
.await?;
// wait for response & return any errors => ignore messageID for now!!!
let _ = rx
.allow_threads_py() // allow-threads-aware async call
.await
.map_err(|_| PyErr::receiver_channel_closed())?
.map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
Ok(())
}
}
pyo3_stub_gen::inventory::submit! {
gen_methods_from_python! {
r#"
class PyNetworkingHandle:
async def recv() -> PyFromSwarm: ...
"#
}
}
pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PyNetworkingHandle>()?;
m.add_class::<PyFromSwarm>()?;
Ok(())
}
+108
View File
@@ -0,0 +1,108 @@
use std::sync::Arc;
use pyo3::exceptions::PyConnectionError;
use pyo3::types::PyBytes;
use pyo3::types::PyNone;
use pyo3::{BoundObject, prelude::*};
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
use zenoh::Result;
use zenoh::{
handlers::FifoChannelHandler,
pubsub::{Publisher, Subscriber},
sample::Sample,
};
use crate::ext::ByteArrayExt;
#[gen_stub_pyclass]
#[pyclass]
pub struct NetReceiver {
pub subscriber: Subscriber<FifoChannelHandler<Sample>>,
}
#[gen_stub_pymethods]
#[pymethods]
impl NetReceiver {
#[gen_stub(override_return_type(
type_repr="collections.abc.Awaitable[bytes | 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, {
assert!(
self.subscriber.receiver_count() == 1,
"tried to receive twice on the same receiver"
);
let subscriber = self.subscriber.clone();
async move {
match subscriber.recv_async().await {
Err(_) => {
// stream closed;
Ok(Python::attach(|py| PyNone::get(py).unbind()).into_any())
}
Ok(sample) => Ok(sample.payload().to_bytes().to_vec().pybytes().into_any()),
}
}
})
}
}
#[gen_stub_pyclass]
#[pyclass]
pub struct NetSender {
pub publisher: Arc<Publisher<'static>>,
pub first: bool,
}
#[gen_stub_pymethods]
#[pymethods]
impl NetSender {
#[gen_stub(override_return_type(
type_repr="collections.abc.Awaitable[bool]",
imports=("collections.abc")
))]
pub fn send<'py>(
&'py mut self,
py: Python<'py>,
data: Bound<'py, PyBytes>,
) -> PyResult<Bound<'py, PyAny>> {
let is_first = self.first;
self.first = false;
pyo3_async_runtimes::tokio::future_into_py(py, {
let publisher = Arc::clone(&self.publisher);
// clone the data so py can have it back
let bytes = data.as_bytes().to_vec();
async move {
if is_first {
wait_for_listener(&*publisher)
.await
.map_err(|e| PyConnectionError::new_err(e.to_string()))?;
}
if !publisher
.matching_status()
.await
.map_err(|e| PyConnectionError::new_err(e.to_string()))?
.matching()
{
return Ok(false);
}
publisher
.put(&bytes)
.await
.map_err(|e| PyConnectionError::new_err(e.to_string()))?;
Ok(true)
}
})
}
}
async fn wait_for_listener<'a>(publisher: &Publisher<'a>) -> Result<()> {
let matcher = publisher.matching_listener().await?;
if publisher.matching_status().await?.matching() {
return Ok(());
}
while let Ok(status) = matcher.recv_async().await {
if status.matching() {
break;
}
}
Ok(())
}
+72
View File
@@ -0,0 +1,72 @@
use std::sync::Arc;
use pyo3::{exceptions::PyValueError, prelude::*};
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
use zenoh::Session;
use zenoh::Wait;
use zenoh::qos::CongestionControl;
use crate::{
point_to_point::{NetReceiver, NetSender},
state::StateProxy,
};
#[gen_stub_pyclass]
#[pyclass]
pub struct PySession {
pub session: Session,
}
#[gen_stub_pymethods]
#[pymethods]
impl PySession {
/* for now construct with NetworkingHandle
#[staticmethod]
#[gen_stub(override_return_type(
type_repr="collections.abc.Awaitable[PySession]",
imports=("collections.abc")
))]
pub fn init<'py>(py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
pyo3_async_runtimes::tokio::future_into_py(py, async move {
Ok(Self {
session: networking::open(
networking::cfg(rand::random(), 0).expect("default cfg is valid"),
)
.await
.map_err(|e| PyRuntimeError::new_err(e.to_string()))?,
})
})
}
*/
pub fn net_receiver<'py>(&self, key: String) -> PyResult<NetReceiver> {
Ok(NetReceiver {
subscriber: self
.session
.declare_subscriber(key)
.wait()
// C5: key format error
.map_err(|e| PyValueError::new_err(e.to_string()))?,
})
}
pub fn net_sender<'py>(&self, key: String) -> PyResult<NetSender> {
Ok(NetSender {
publisher: Arc::new(
self.session
.declare_publisher(key)
.congestion_control(CongestionControl::Block)
.wait()
// C5: key format error, could be declaration error
.map_err(|e| PyValueError::new_err(e.to_string()))?,
),
first: true,
})
}
pub fn state_proxy(&self) -> StateProxy {
StateProxy {
session: self.session.clone(),
}
}
}
+78
View File
@@ -0,0 +1,78 @@
use pyo3::{exceptions::PyValueError, prelude::*};
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
use serde_json::{Map, Value};
use zenoh::{Result, Session, sample::SampleFields};
#[gen_stub_pyclass]
#[pyclass]
pub struct StateProxy {
pub session: Session,
}
#[gen_stub_pymethods]
#[pymethods]
impl StateProxy {
#[gen_stub(override_return_type(
type_repr="collections.abc.Awaitable[str]",
imports=("collections.abc")
))]
pub fn snapshot<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
pyo3_async_runtimes::tokio::future_into_py(py, {
let session = self.session.clone();
async move {
Self::_snapshot(session)
.await
.map_err(|e| PyValueError::new_err(e.to_string()))
.map(|v| v.to_string())
}
})
}
}
impl StateProxy {
async fn _snapshot(session: Session) -> Result<Value> {
let q = session.get("storage/mem1/**").await?;
let mut v = Value::Object(Map::default());
while let Ok(sample) = q.recv_async().await {
let mut cur_v = &mut v;
let Ok(sample) = sample.into_result() else {
continue;
};
// skip storage/mem1
let SampleFields {
payload, key_expr, ..
} = sample.into();
let mut iter = key_expr.split('/').skip(2).peekable();
loop {
let Some(p) = iter.next() else {
break;
};
if iter.peek().is_none() {
// terminal; write value into json
let existing = cur_v
.as_object_mut()
.expect("path terminated unexpectedly - value stored at some/path and some/path/two")
.insert(p.to_owned(), Value::String(payload.try_to_string()?.to_string()));
if let Some(value) = existing {
assert!(value.is_string())
// could log, but string overwrites are fine
}
} else {
// non-terminal; ensure key exists in v, then replace cur with that object
cur_v = cur_v
.as_object_mut()
.expect("path terminated unexpectedly - value stored at some/path and some/path/two")
.entry(p)
.or_insert(Value::Object(Map::default()));
assert!(
cur_v.is_object(),
"path terminated unexpectedly - value stored at some/path and some/path/two"
)
}
}
}
Ok(v)
}
}
-66
View File
@@ -1,66 +0,0 @@
[package]
name = "exo_pyo3_bindings"
version = { workspace = true }
edition = { workspace = true }
publish = false
[lib]
doctest = false
path = "src/lib.rs"
name = "exo_pyo3_bindings"
# "cdylib" needed to produce shared library for Python to import
# "rlib" needed for stub-gen to run
crate-type = ["cdylib", "rlib"]
[[bin]]
path = "src/bin/stub_gen.rs"
name = "stub_gen"
doc = false
[lints]
workspace = true
[dependencies]
networking = { workspace = true }
# interop
pyo3 = { version = "0.27.2", features = [
# "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13
# "nightly", # enables better-supported GIL integration
"experimental-async", # async support in #[pyfunction] & #[pymethods]
#"experimental-inspect", # inspection of generated binary => easier to automate type-hint generation
#"py-clone", # adding Clone-ing of `Py<T>` without GIL (may cause panics - remove if panics happen)
# "multiple-pymethods", # allows multiple #[pymethods] sections per class
# integrations with other libraries
# "arc_lock", "bigdecimal", "either", "hashbrown", "indexmap", "num-bigint", "num-complex", "num-rational",
# "ordered-float", "rust_decimal", "smallvec",
# "anyhow", "chrono", "chrono-local", "chrono-tz", "eyre", "jiff-02", "lock_api", "parking-lot", "time", "serde",
] }
pyo3-stub-gen = { version = "0.17.2" }
pyo3-async-runtimes = { version = "0.27.0", features = [
"attributes",
"tokio-runtime",
"testing",
] }
pyo3-log = "0.13.2"
# macro dependencies
extend = { workspace = true }
delegate = { workspace = true }
# async runtime
tokio = { workspace = true, features = ["full", "tracing"] }
futures-lite = { workspace = true }
# utility dependencies
util = { workspace = true }
# Tracing
log = { workspace = true }
env_logger = "0.11"
# Networking
libp2p = { workspace = true, features = ["full"] }
pin-project = "1.1.10"
@@ -1,94 +0,0 @@
# This file is automatically generated by pyo3_stub_gen
# ruff: noqa: E501, F401
import builtins
import typing
@typing.final
class AllQueuesFullError(builtins.Exception):
def __new__(cls, *args: typing.Any) -> AllQueuesFullError: ...
def __repr__(self) -> builtins.str: ...
def __str__(self) -> builtins.str: ...
@typing.final
class Keypair:
r"""
Identity keypair of a node.
"""
@staticmethod
def generate() -> Keypair:
r"""
Generate a new Ed25519 keypair.
"""
@staticmethod
def from_bytes(bytes: bytes) -> Keypair:
r"""
Construct an Ed25519 keypair from secret key bytes
"""
def to_bytes(self) -> bytes:
r"""
Get the secret key bytes underlying the keypair
"""
def to_node_id(self) -> builtins.str:
r"""
Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
"""
@typing.final
class MessageTooLargeError(builtins.Exception):
def __new__(cls, *args: typing.Any) -> MessageTooLargeError: ...
def __repr__(self) -> builtins.str: ...
def __str__(self) -> builtins.str: ...
@typing.final
class NetworkingHandle:
def __new__(cls, identity: Keypair, bootstrap_peers: typing.Sequence[builtins.str], listen_port: builtins.int) -> NetworkingHandle: ...
async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
r"""
Subscribe to a `GossipSub` topic.
Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
"""
async def gossipsub_unsubscribe(self, topic: builtins.str) -> builtins.bool:
r"""
Unsubscribes from a `GossipSub` topic.
Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
"""
async def gossipsub_publish(self, topic: builtins.str, data: bytes) -> None:
r"""
Publishes a message with multiple topics to the `GossipSub` network.
If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
"""
async def recv(self) -> PyFromSwarm: ...
@typing.final
class NoPeersSubscribedToTopicError(builtins.Exception):
def __new__(cls, *args: typing.Any) -> NoPeersSubscribedToTopicError: ...
def __repr__(self) -> builtins.str: ...
def __str__(self) -> builtins.str: ...
class PyFromSwarm:
@typing.final
class Connection(PyFromSwarm):
__match_args__ = ("peer_id", "connected",)
@property
def peer_id(self) -> builtins.str: ...
@property
def connected(self) -> builtins.bool: ...
def __new__(cls, peer_id: builtins.str, connected: builtins.bool) -> PyFromSwarm.Connection: ...
@typing.final
class Message(PyFromSwarm):
__match_args__ = ("origin", "topic", "data",)
@property
def origin(self) -> builtins.str: ...
@property
def topic(self) -> builtins.str: ...
@property
def data(self) -> bytes: ...
def __new__(cls, origin: builtins.str, topic: builtins.str, data: bytes) -> PyFromSwarm.Message: ...
...
-47
View File
@@ -1,47 +0,0 @@
use crate::ext::ResultExt as _;
use libp2p::identity::Keypair;
use pyo3::types::{PyBytes, PyBytesMethods as _};
use pyo3::{Bound, PyResult, Python, pyclass, pymethods};
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
/// Identity keypair of a node.
#[gen_stub_pyclass]
#[pyclass(name = "Keypair", frozen)]
#[repr(transparent)]
pub struct PyKeypair(pub Keypair);
#[gen_stub_pymethods]
#[pymethods]
#[allow(clippy::needless_pass_by_value)]
impl PyKeypair {
/// Generate a new Ed25519 keypair.
#[staticmethod]
fn generate() -> Self {
Self(Keypair::generate_ed25519())
}
/// Construct an Ed25519 keypair from secret key bytes
#[staticmethod]
fn from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
let mut bytes = Vec::from(bytes.as_bytes());
Ok(Self(Keypair::ed25519_from_bytes(&mut bytes).pyerr()?))
}
/// Get the secret key bytes underlying the keypair
fn to_bytes<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
let bytes = self
.0
.clone()
.try_into_ed25519()
.pyerr()?
.secret()
.as_ref()
.to_vec();
Ok(PyBytes::new(py, &bytes))
}
/// Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
fn to_node_id(&self) -> String {
self.0.public().to_peer_id().to_base58()
}
}
-318
View File
@@ -1,318 +0,0 @@
use std::pin::Pin;
use std::sync::Arc;
use crate::r#const::MPSC_CHANNEL_SIZE;
use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
use crate::ext::{ResultExt as _, TokioMpscSenderExt as _};
use crate::ident::PyKeypair;
use crate::networking::exception::{
PyAllQueuesFullError, PyMessageTooLargeError, PyNoPeersSubscribedToTopicError,
};
use crate::pyclass;
use futures_lite::{Stream, StreamExt as _};
use libp2p::gossipsub::PublishError;
use networking::swarm::{FromSwarm, ToSwarm, create_swarm};
use pyo3::exceptions::PyRuntimeError;
use pyo3::prelude::{PyModule, PyModuleMethods as _};
use pyo3::types::PyBytes;
use pyo3::{Bound, Py, PyAny, PyErr, PyResult, Python, pymethods};
use pyo3_stub_gen::derive::{
gen_methods_from_python, gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods,
};
use tokio::sync::{Mutex, mpsc, oneshot};
mod exception {
use pyo3::types::PyTuple;
use pyo3::{exceptions::PyException, prelude::*};
use pyo3_stub_gen::derive::*;
#[gen_stub_pyclass]
#[pyclass(frozen, extends=PyException, name="NoPeersSubscribedToTopicError")]
pub struct PyNoPeersSubscribedToTopicError {}
impl PyNoPeersSubscribedToTopicError {
const MSG: &'static str = "\
No peers are currently subscribed to receive messages on this topic. \
Wait for peers to subscribe or check your network connectivity.";
/// Creates a new [ `PyErr` ] of this type.
///
/// [`PyErr`] : https://docs.rs/pyo3/latest/pyo3/struct.PyErr.html "PyErr in pyo3"
pub(crate) fn new_err() -> PyErr {
PyErr::new::<Self, _>(()) // TODO: check if this needs to be replaced???
}
}
#[gen_stub_pymethods]
#[pymethods]
impl PyNoPeersSubscribedToTopicError {
#[new]
#[pyo3(signature = (*args))]
#[allow(unused_variables)]
pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self {
Self {}
}
fn __repr__(&self) -> String {
format!("PeerId(\"{}\")", Self::MSG)
}
fn __str__(&self) -> String {
Self::MSG.to_string()
}
}
#[gen_stub_pyclass]
#[pyclass(frozen, extends=PyException, name="AllQueuesFullError")]
pub struct PyAllQueuesFullError {}
impl PyAllQueuesFullError {
const MSG: &'static str =
"All libp2p peers are unresponsive, resend the message or reconnect.";
/// Creates a new [ `PyErr` ] of this type.
///
/// [`PyErr`] : https://docs.rs/pyo3/latest/pyo3/struct.PyErr.html "PyErr in pyo3"
pub(crate) fn new_err() -> PyErr {
PyErr::new::<Self, _>(()) // TODO: check if this needs to be replaced???
}
}
#[gen_stub_pymethods]
#[pymethods]
impl PyAllQueuesFullError {
#[new]
#[pyo3(signature = (*args))]
#[allow(unused_variables)]
pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self {
Self {}
}
fn __repr__(&self) -> String {
format!("PeerId(\"{}\")", Self::MSG)
}
fn __str__(&self) -> String {
Self::MSG.to_string()
}
}
#[gen_stub_pyclass]
#[pyclass(frozen, extends=PyException, name="MessageTooLargeError")]
pub struct PyMessageTooLargeError {}
impl PyMessageTooLargeError {
const MSG: &'static str = "Gossipsub message exceeds max_transmit_size. Reduce prompt length or increase the limit.";
pub(crate) fn new_err() -> PyErr {
PyErr::new::<Self, _>(())
}
}
#[gen_stub_pymethods]
#[pymethods]
impl PyMessageTooLargeError {
#[new]
#[pyo3(signature = (*args))]
#[allow(unused_variables)]
pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self {
Self {}
}
fn __repr__(&self) -> String {
format!("MessageTooLargeError(\"{}\")", Self::MSG)
}
fn __str__(&self) -> String {
Self::MSG.to_string()
}
}
}
#[gen_stub_pyclass]
#[pyclass(name = "NetworkingHandle")]
struct PyNetworkingHandle {
// channels
pub to_swarm: mpsc::Sender<ToSwarm>,
pub swarm: Arc<Mutex<Pin<Box<dyn Stream<Item = FromSwarm> + Send>>>>,
}
#[gen_stub_pyclass_complex_enum]
#[pyclass]
enum PyFromSwarm {
Connection {
peer_id: String,
connected: bool,
},
Message {
origin: String,
topic: String,
data: Py<PyBytes>,
},
}
impl From<FromSwarm> for PyFromSwarm {
fn from(value: FromSwarm) -> Self {
match value {
FromSwarm::Discovered { peer_id } => Self::Connection {
peer_id: peer_id.to_base58(),
connected: true,
},
FromSwarm::Expired { peer_id } => Self::Connection {
peer_id: peer_id.to_base58(),
connected: false,
},
FromSwarm::Message { from, topic, data } => Self::Message {
origin: from.to_base58(),
topic: topic,
data: data.pybytes(),
},
}
}
}
#[gen_stub_pymethods]
#[pymethods]
impl PyNetworkingHandle {
// NOTE: `async fn`s here that use `.await` will wrap the future in `.allow_threads_py()`
// immediately beforehand to release the interpreter.
// SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
// ---- Lifecycle management methods ----
#[new]
#[pyo3(signature = (identity, bootstrap_peers, listen_port))]
fn py_new(
identity: Bound<'_, PyKeypair>,
bootstrap_peers: Vec<String>,
listen_port: u16,
) -> PyResult<Self> {
// create communication channels
let (to_swarm, from_client) = mpsc::channel(MPSC_CHANNEL_SIZE);
// get identity
let identity = identity.borrow().0.clone();
// create networking swarm (within tokio context!! or it crashes)
let _guard = pyo3_async_runtimes::tokio::get_runtime().enter();
let swarm = create_swarm(identity, from_client, bootstrap_peers, listen_port)
.pyerr()?
.into_stream();
Ok(Self {
swarm: Arc::new(Mutex::new(swarm)),
to_swarm,
})
}
#[gen_stub(skip)]
fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let swarm = Arc::clone(&self.swarm);
pyo3_async_runtimes::tokio::future_into_py(py, async move {
swarm
.try_lock()
.map_err(|_| PyRuntimeError::new_err("called recv twice concurrently"))?
.next()
.await
.ok_or(PyErr::receiver_channel_closed())
.map(PyFromSwarm::from)
})
}
// ---- Gossipsub management methods ----
/// Subscribe to a `GossipSub` topic.
///
/// Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
async fn gossipsub_subscribe(&self, topic: String) -> PyResult<bool> {
let (tx, rx) = oneshot::channel();
// send off request to subscribe
self.to_swarm
.send_py(ToSwarm::Subscribe {
topic,
result_sender: tx,
})
.allow_threads_py() // allow-threads-aware async call
.await?;
// wait for response & return any errors
rx.allow_threads_py() // allow-threads-aware async call
.await
.map_err(|_| PyErr::receiver_channel_closed())?
.pyerr()
}
/// Unsubscribes from a `GossipSub` topic.
///
/// Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
async fn gossipsub_unsubscribe(&self, topic: String) -> PyResult<bool> {
let (tx, rx) = oneshot::channel();
// send off request to unsubscribe
self.to_swarm
.send_py(ToSwarm::Unsubscribe {
topic,
result_sender: tx,
})
.allow_threads_py() // allow-threads-aware async call
.await?;
// wait for response & convert any errors
rx.allow_threads_py() // allow-threads-aware async call
.await
.map_err(|_| PyErr::receiver_channel_closed())
}
/// Publishes a message with multiple topics to the `GossipSub` network.
///
/// If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
async fn gossipsub_publish(&self, topic: String, data: Py<PyBytes>) -> PyResult<()> {
let (tx, rx) = oneshot::channel();
// send off request to subscribe
let data = Python::attach(|py| Vec::from(data.as_bytes(py)));
self.to_swarm
.send_py(ToSwarm::Publish {
topic,
data,
result_sender: tx,
})
.allow_threads_py() // allow-threads-aware async call
.await?;
// wait for response & return any errors => ignore messageID for now!!!
let _ = rx
.allow_threads_py() // allow-threads-aware async call
.await
.map_err(|_| PyErr::receiver_channel_closed())?
.map_err(|e| match e {
PublishError::AllQueuesFull(_) => PyAllQueuesFullError::new_err(),
PublishError::MessageTooLarge => PyMessageTooLargeError::new_err(),
PublishError::NoPeersSubscribedToTopic => {
PyNoPeersSubscribedToTopicError::new_err()
}
e => PyRuntimeError::new_err(e.to_string()),
})?;
Ok(())
}
}
pyo3_stub_gen::inventory::submit! {
gen_methods_from_python! {
r#"
class PyNetworkingHandle:
async def recv() -> PyFromSwarm: ...
"#
}
}
pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<exception::PyNoPeersSubscribedToTopicError>()?;
m.add_class::<exception::PyAllQueuesFullError>()?;
m.add_class::<exception::PyMessageTooLargeError>()?;
m.add_class::<PyNetworkingHandle>()?;
m.add_class::<PyFromSwarm>()?;
Ok(())
}
-54
View File
@@ -1,54 +0,0 @@
#[cfg(test)]
mod tests {
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;
let (tx, mut rx) = mpsc::channel::<Ping>(10);
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;
}
}
}
_ = tokio::time::sleep(Duration::from_secs_f32(0.1)) => {
println!("TASK: heartbeat");
}
}
}
println!("TASK: exited");
});
let tx2 = tx.clone();
tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
tx.send(Ping).await.expect("Should not fail");
drop(tx);
tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
tx2.send(Ping).await.expect("Should not fail");
drop(tx2);
tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
}
}
@@ -1,36 +0,0 @@
import asyncio
import pytest
from exo_pyo3_bindings import (
Keypair,
NetworkingHandle,
NoPeersSubscribedToTopicError,
PyFromSwarm,
)
@pytest.mark.asyncio
async def test_sleep_on_multiple_items() -> None:
print("PYTHON: starting handle")
h = NetworkingHandle(Keypair.generate(), [], 0)
rt = asyncio.create_task(_await_recv(h))
# sleep for 4 ticks
for i in range(4):
await asyncio.sleep(1)
try:
await h.gossipsub_publish("topic", b"somehting or other")
except NoPeersSubscribedToTopicError as e:
print("caught it", e)
async def _await_recv(h: NetworkingHandle):
while True:
event = await h.recv()
match event:
case PyFromSwarm.Connection() as c:
print(f"PYTHON: connection update: {c}")
case PyFromSwarm.Message() as m:
print(f"PYTHON: message: {m}")
+13 -36
View File
@@ -1,42 +1,19 @@
[package]
name = "networking"
version = { workspace = true }
edition = { workspace = true }
publish = false
version.workspace = true
edition.workspace = true
[lib]
doctest = false
name = "networking"
path = "src/lib.rs"
[dependencies]
async-stream.workspace = true
futures-lite.workspace = true
netwatcher = { workspace = true, features = ["tokio"] }
parking_lot.workspace = true
tokio = { workspace = true, features = ["full"] }
zenoh = { workspace = true, features = ["internal", "plugins", "unstable"] }
zenoh-plugin-storage-manager.workspace = true
zenoh-plugin-trait.workspace = true
rand.workspace = true
log.workspace = true
[lints]
workspace = true
[dependencies]
# datastructures
either = { workspace = true }
# macro dependencies
extend = { workspace = true }
delegate = { workspace = true }
# async
async-stream = { workspace = true }
futures-lite = { workspace = true }
futures-timer = { workspace = true }
tokio = { workspace = true, features = ["full"] }
# utility dependencies
util = { workspace = true }
tracing-subscriber = { version = "0.3.19", features = [
"default",
"env-filter",
] }
keccak-const = { workspace = true }
# tracing/logging
log = { workspace = true }
# networking
libp2p = { workspace = true, features = ["full"] }
pin-project = "1.1.10"
-86
View File
@@ -1,86 +0,0 @@
use futures_lite::StreamExt;
use libp2p::identity;
use networking::swarm;
use networking::swarm::{FromSwarm, ToSwarm};
use tokio::sync::{mpsc, oneshot};
use tokio::{io, io::AsyncBufReadExt as _};
use tracing_subscriber::EnvFilter;
use tracing_subscriber::filter::LevelFilter;
#[tokio::main]
async fn main() {
let _ = tracing_subscriber::fmt()
.with_env_filter(EnvFilter::from_default_env().add_directive(LevelFilter::INFO.into()))
.try_init();
let (to_swarm, from_client) = mpsc::channel(20);
// Configure swarm
let mut swarm = swarm::create_swarm(
identity::Keypair::generate_ed25519(),
from_client,
vec![],
0,
)
.expect("Swarm creation failed")
.into_stream();
// Create a Gossipsub topic & subscribe
let (tx, rx) = oneshot::channel();
_ = to_swarm
.send(ToSwarm::Subscribe {
topic: "test-net".to_string(),
result_sender: tx,
})
.await
.expect("should send");
// Read full lines from stdin
let mut stdin = io::BufReader::new(io::stdin()).lines();
println!("Enter messages via STDIN and they will be sent to connected peers using Gossipsub");
tokio::task::spawn(async move {
rx.await
.expect("tx not dropped")
.expect("subscribe shouldn't fail");
loop {
if let Ok(Some(line)) = stdin.next_line().await {
let (tx, rx) = oneshot::channel();
if let Err(e) = to_swarm
.send(swarm::ToSwarm::Publish {
topic: "test-net".to_string(),
data: line.as_bytes().to_vec(),
result_sender: tx,
})
.await
{
println!("Send error: {e:?}");
return;
};
match rx.await {
Ok(Err(e)) => println!("Publish error: {e:?}"),
Err(e) => println!("Publish error: {e:?}"),
Ok(_) => {}
}
}
}
});
// Kick it off
loop {
// on gossipsub outgoing
match swarm.next().await {
// on gossipsub incoming
Some(FromSwarm::Discovered { peer_id }) => {
println!("\n\nconnected to {peer_id}\n\n")
}
Some(FromSwarm::Expired { peer_id }) => {
println!("\n\ndisconnected from {peer_id}\n\n")
}
Some(FromSwarm::Message { from, topic, data }) => {
println!("{topic}/{from}:\n{}", String::from_utf8_lossy(&data))
}
None => {}
}
}
}
+23
View File
@@ -0,0 +1,23 @@
use log::info;
use networking;
use zenoh::Result;
#[tokio::main]
async fn main() -> Result<()> {
zenoh::init_log_from_env_or("info");
info!("Opening session...");
let cfg = networking::cfg(rand::random(), 0)?;
let session = networking::open(cfg).await?;
let _tok = session
.z
.liveliness()
.declare_token(format!("nodes/{}/live", session.z.zid()))
.await?;
let key_expr = "storage/mem1/name";
let payload = "me";
info!("Putting Data ('{key_expr}': '{payload}')...");
session.z.put(key_expr, payload).await?;
tokio::signal::ctrl_c().await?;
Ok(())
}
+25
View File
@@ -0,0 +1,25 @@
use log::info;
use networking;
use zenoh::Result;
#[tokio::main]
async fn main() -> Result<()> {
zenoh::init_log_from_env_or("info");
info!("Opening session...");
let cfg = networking::cfg(rand::random(), 0)?;
let session = networking::open(cfg).await?;
let _tok = session
.z
.liveliness()
.declare_token(format!("nodes/{}/live", session.z.zid()))
.await?;
let _sub = session
.z
.liveliness()
.declare_subscriber("nodes/*/live")
.history(true)
.callback(|tok| println!("{tok:?}"))
.await?;
tokio::signal::ctrl_c().await?;
Ok(())
}
-44
View File
@@ -1,44 +0,0 @@
https://github.com/ml-explore/mlx/commit/3fe98bacc7640d857acf3539f1d21b47a32e5609
^raw sockets distributed -> `<net/ndrv.h>` -> https://newosxbook.com/code/xnu-3247.1.106/bsd/net/ndrv.h.auto.html
--> header file for a networking component found in the macOS kernel (XNU) that defines structures for network device driver registration, specifically the ndrv_demux_desc and ndrv_protocol_desc structures used for demultiplexing protocol data at the network interface level. It specifies how to describe protocol data, such as an Ethernet type or a SNAP header, and how to associate these descriptions with a specific protocol family to receive matching packets.
--> Used to bind an NDRV socket so that packets that match given protocol demux descriptions can be received.
--> An NDRV socket is a special kind of socket in the Darwin/macOS operating system's XNU kernel, used for low-level network packet manipulation and binding to specific protocols for packet processing. It allows user-space applications or drivers to directly write Layer 2 (L2) network packets or interact with the network stack at a lower level, often by binding to protocol descriptors like the ndrv_protocol_desc. This type of socket is used for functions such as capturing and injecting packets, especially in network infrastructure software like routers or for kernel-level network monitoring and security tools.
--> also called PF_NDRV sockets --> https://newosxbook.com/bonus/vol1ch16.html
----> they are conceptually similar to https://scapy.disruptivelabs.in/networking/socket-interface PF_RAW or PF_PACKET
https://stackoverflow.com/questions/17169298/af-packet-on-osx
^AF_PACKET duplicates the packets as soon as it receives them from the physical layer (for incoming packets) or just before sending them out to the physical layer (for outgoing packets). -> this is on Linux only
^it doesn't exist on OS X so you can use /dev/bpfX (Berkeley Packet Filter) for sniffing
https://www.unix.com/man_page/mojave/4/ip/
^OS X manpages for IP
https://developer.apple.com/documentation/kernel/implementing_drivers_system_extensions_and_kexts
^driver kit, system extensions & kexts for macOS
----
To set up a Linux system to use a Thunderbolt connection as a network device, connect the two computers with a Thunderbolt cable, load the thunderbolt-net kernel module (usually automatic but modprobe is an option for manual loading), and then the operating system will create virtual Ethernet interfaces (e.g., thunderbolt0) for networking. You can then use standard tools like ifconfig or your desktop environment's network manager to configure these new interfaces for a link-local network.
--> https://gist.github.com/geosp/80fbd39e617b7d1d9421683df4ea224a
----> here is a guide on how to set up thunderbolt-ethernet on linux
----> I may be able to steal the thunderbolt-net code ideas to implement a kernel module for MacOS
https://chatgpt.com/s/t_68af8e41a8548191993281a014f846a7
^GPT discussion about making socket interface
https://chatgpt.com/s/t_68afb798a85c8191973c02a0fa7a48a3 --> link-local address,,??
https://chatgpt.com/s/t_68afb02987e08191b2b0044d3667ece2
^GPT discussion about accessing TB on MacOS low level interactions
--------------------------------
https://www.intel.com/content/www/us/en/support/articles/000098893/software.html
^Thunderbolt Share & Thunderbolt Networking Mode => intel's equivalent of thunderbolt bridge
---------------------------------
https://www.zerotier.com/blog/how-zerotier-eliminated-kernel-extensions-on-macos/
-->fake ethernet devices on MacOS -> omg??? we can detect thunderbolt bridge, then bind to it, then re-expose it as fake ethernet??
-->ps: https://chatgpt.com/s/t_68afb2b25fb881919526763fb5d7359c, AF/PF_NDRV are one and the same!!!
-->https://github.com/zerotier/ZeroTierOne/blob/dev/osdep/MacEthernetTapAgent.c
-390
View File
@@ -1,390 +0,0 @@
use crate::ext::MultiaddrExt;
use delegate::delegate;
use either::Either;
use futures_lite::FutureExt;
use futures_timer::Delay;
use libp2p::core::transport::PortUse;
use libp2p::core::{ConnectedPoint, Endpoint};
use libp2p::swarm::behaviour::ConnectionEstablished;
use libp2p::swarm::dial_opts::DialOpts;
use libp2p::swarm::{
CloseConnection, ConnectionClosed, ConnectionDenied, ConnectionHandler,
ConnectionHandlerSelect, ConnectionId, FromSwarm, NetworkBehaviour, THandler, THandlerInEvent,
THandlerOutEvent, ToSwarm, dummy,
};
use libp2p::{Multiaddr, PeerId, identity, mdns};
use std::collections::{BTreeSet, HashMap};
use std::convert::Infallible;
use std::io;
use std::net::IpAddr;
use std::task::{Context, Poll};
use std::time::Duration;
use util::wakerdeque::WakerDeque;
const RETRY_CONNECT_INTERVAL: Duration = Duration::from_secs(5);
mod managed {
use libp2p::swarm::NetworkBehaviour;
use libp2p::{identity, mdns, ping};
use std::io;
use std::time::Duration;
const MDNS_RECORD_TTL: Duration = Duration::from_secs(2_500);
const MDNS_QUERY_INTERVAL: Duration = Duration::from_secs(1_500);
const PING_TIMEOUT: Duration = Duration::from_millis(2_500);
const PING_INTERVAL: Duration = Duration::from_millis(2_500);
#[derive(NetworkBehaviour)]
pub struct Behaviour {
mdns: mdns::tokio::Behaviour,
ping: ping::Behaviour,
}
impl Behaviour {
pub fn new(keypair: &identity::Keypair) -> io::Result<Self> {
Ok(Self {
mdns: mdns_behaviour(keypair)?,
ping: ping_behaviour(),
})
}
}
fn mdns_behaviour(keypair: &identity::Keypair) -> io::Result<mdns::tokio::Behaviour> {
use mdns::{Config, tokio};
// mDNS config => enable IPv6
let mdns_config = Config {
ttl: MDNS_RECORD_TTL,
query_interval: MDNS_QUERY_INTERVAL,
// enable_ipv6: true, // TODO: for some reason, TCP+mDNS don't work well with ipv6?? figure out how to make work
..Default::default()
};
let mdns_behaviour = tokio::Behaviour::new(mdns_config, keypair.public().to_peer_id());
Ok(mdns_behaviour?)
}
fn ping_behaviour() -> ping::Behaviour {
ping::Behaviour::new(
ping::Config::new()
.with_timeout(PING_TIMEOUT)
.with_interval(PING_INTERVAL),
)
}
}
/// Events for when a listening connection is truly established and truly closed.
#[derive(Debug, Clone)]
pub enum Event {
ConnectionEstablished {
peer_id: PeerId,
connection_id: ConnectionId,
remote_ip: IpAddr,
remote_tcp_port: u16,
},
ConnectionClosed {
peer_id: PeerId,
connection_id: ConnectionId,
remote_ip: IpAddr,
remote_tcp_port: u16,
},
}
/// Discovery behavior that wraps mDNS to produce truly discovered durable peer-connections.
///
/// The behaviour operates as such:
/// 1) All true (listening) connections/disconnections are tracked, emitting corresponding events
/// to the swarm.
/// 1) mDNS discovered/expired peers are tracked; discovered but not connected peers are dialed
/// immediately, and expired but connected peers are disconnected from immediately.
/// 2) Every fixed interval: discovered but not connected peers are dialed, and expired but
/// connected peers are disconnected from.
pub struct Behaviour {
// state-tracking for managed behaviors & mDNS-discovered peers
managed: managed::Behaviour,
mdns_discovered: HashMap<PeerId, BTreeSet<Multiaddr>>,
bootstrap_peers: Vec<Multiaddr>,
retry_delay: Delay, // retry interval
// pending events to emmit => waker-backed Deque to control polling
pending_events: WakerDeque<ToSwarm<Event, Infallible>>,
}
impl Behaviour {
pub fn new(keypair: &identity::Keypair, bootstrap_peers: Vec<Multiaddr>) -> io::Result<Self> {
Ok(Self {
managed: managed::Behaviour::new(keypair)?,
mdns_discovered: HashMap::new(),
bootstrap_peers,
retry_delay: Delay::new(RETRY_CONNECT_INTERVAL),
pending_events: WakerDeque::new(),
})
}
fn dial(&mut self, peer_id: PeerId, addr: Multiaddr) {
self.pending_events.push_back(ToSwarm::Dial {
opts: DialOpts::peer_id(peer_id).addresses(vec![addr]).build(),
})
}
fn close_connection(&mut self, peer_id: PeerId, connection: ConnectionId) {
// push front to make this IMMEDIATE
self.pending_events.push_front(ToSwarm::CloseConnection {
peer_id,
connection: CloseConnection::One(connection),
})
}
fn handle_mdns_discovered(&mut self, peers: Vec<(PeerId, Multiaddr)>) {
for (p, ma) in peers {
self.dial(p, ma.clone()); // always connect
// get peer's multi-addresses or insert if missing
let Some(mas) = self.mdns_discovered.get_mut(&p) else {
self.mdns_discovered.insert(p, BTreeSet::from([ma]));
continue;
};
// multiaddress should never already be present - else something has gone wrong
let is_new_addr = mas.insert(ma);
assert!(is_new_addr, "cannot discover a discovered peer");
}
}
fn handle_mdns_expired(&mut self, peers: Vec<(PeerId, Multiaddr)>) {
for (p, ma) in peers {
// at this point, we *must* have the peer
let mas = self
.mdns_discovered
.get_mut(&p)
.expect("nonexistent peer cannot expire");
// at this point, we *must* have the multiaddress
let was_present = mas.remove(&ma);
assert!(was_present, "nonexistent multiaddress cannot expire");
// if empty, remove the peer-id entirely
if mas.is_empty() {
self.mdns_discovered.remove(&p);
}
}
}
fn on_connection_established(
&mut self,
peer_id: PeerId,
connection_id: ConnectionId,
remote_ip: IpAddr,
remote_tcp_port: u16,
) {
// send out connected event
self.pending_events
.push_back(ToSwarm::GenerateEvent(Event::ConnectionEstablished {
peer_id,
connection_id,
remote_ip,
remote_tcp_port,
}));
}
fn on_connection_closed(
&mut self,
peer_id: PeerId,
connection_id: ConnectionId,
remote_ip: IpAddr,
remote_tcp_port: u16,
) {
// send out disconnected event
self.pending_events
.push_back(ToSwarm::GenerateEvent(Event::ConnectionClosed {
peer_id,
connection_id,
remote_ip,
remote_tcp_port,
}));
}
}
impl NetworkBehaviour for Behaviour {
type ConnectionHandler =
ConnectionHandlerSelect<dummy::ConnectionHandler, THandler<managed::Behaviour>>;
type ToSwarm = Event;
// simply delegate to underlying mDNS behaviour
delegate! {
to self.managed {
fn handle_pending_inbound_connection(&mut self, connection_id: ConnectionId, local_addr: &Multiaddr, remote_addr: &Multiaddr) -> Result<(), ConnectionDenied>;
fn handle_pending_outbound_connection(&mut self, connection_id: ConnectionId, maybe_peer: Option<PeerId>, addresses: &[Multiaddr], effective_role: Endpoint) -> Result<Vec<Multiaddr>, ConnectionDenied>;
}
}
fn handle_established_inbound_connection(
&mut self,
connection_id: ConnectionId,
peer: PeerId,
local_addr: &Multiaddr,
remote_addr: &Multiaddr,
) -> Result<THandler<Self>, ConnectionDenied> {
Ok(ConnectionHandler::select(
dummy::ConnectionHandler,
self.managed.handle_established_inbound_connection(
connection_id,
peer,
local_addr,
remote_addr,
)?,
))
}
#[allow(clippy::needless_question_mark)]
fn handle_established_outbound_connection(
&mut self,
connection_id: ConnectionId,
peer: PeerId,
addr: &Multiaddr,
role_override: Endpoint,
port_use: PortUse,
) -> Result<THandler<Self>, ConnectionDenied> {
Ok(ConnectionHandler::select(
dummy::ConnectionHandler,
self.managed.handle_established_outbound_connection(
connection_id,
peer,
addr,
role_override,
port_use,
)?,
))
}
fn on_connection_handler_event(
&mut self,
peer_id: PeerId,
connection_id: ConnectionId,
event: THandlerOutEvent<Self>,
) {
match event {
Either::Left(ev) => libp2p::core::util::unreachable(ev),
Either::Right(ev) => {
self.managed
.on_connection_handler_event(peer_id, connection_id, ev)
}
}
}
// hook into these methods to drive behavior
fn on_swarm_event(&mut self, event: FromSwarm) {
self.managed.on_swarm_event(event); // let mDNS handle swarm events
// handle swarm events to update internal state:
match event {
FromSwarm::ConnectionEstablished(ConnectionEstablished {
peer_id,
connection_id,
endpoint,
..
}) => {
let remote_address = match endpoint {
ConnectedPoint::Dialer { address, .. } => address,
ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr,
};
if let Some((ip, port)) = remote_address.try_to_tcp_addr() {
// handle connection established event which is filtered correctly
self.on_connection_established(peer_id, connection_id, ip, port)
}
}
FromSwarm::ConnectionClosed(ConnectionClosed {
peer_id,
connection_id,
endpoint,
..
}) => {
let remote_address = match endpoint {
ConnectedPoint::Dialer { address, .. } => address,
ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr,
};
if let Some((ip, port)) = remote_address.try_to_tcp_addr() {
// handle connection closed event which is filtered correctly
self.on_connection_closed(peer_id, connection_id, ip, port)
}
}
// since we are running TCP/IP transport layer, we are assuming that
// no address changes can occur, hence encountering one is a fatal error
FromSwarm::AddressChange(a) => {
unreachable!("unhandlable: address change encountered: {:?}", a)
}
_ => {}
}
}
fn poll(&mut self, cx: &mut Context) -> Poll<ToSwarm<Self::ToSwarm, THandlerInEvent<Self>>> {
// delegate to managed behaviors for any behaviors they need to perform
match self.managed.poll(cx) {
Poll::Ready(ToSwarm::GenerateEvent(e)) => {
match e {
// handle discovered and expired events from mDNS
managed::BehaviourEvent::Mdns(e) => match e.clone() {
mdns::Event::Discovered(peers) => {
self.handle_mdns_discovered(peers);
}
mdns::Event::Expired(peers) => {
self.handle_mdns_expired(peers);
}
},
// handle ping events => if error then disconnect
managed::BehaviourEvent::Ping(e) => {
if let Err(_) = e.result {
self.close_connection(e.peer, e.connection.clone())
}
}
}
// since we just consumed an event, we should immediately wake just in case
// there are more events to come where that came from
cx.waker().wake_by_ref();
}
// forward any other mDNS event to the swarm or its connection handler(s)
Poll::Ready(e) => {
return Poll::Ready(
e.map_out(|_| unreachable!("events returning to swarm already handled"))
.map_in(Either::Right),
);
}
Poll::Pending => {}
}
// retry connecting to all mDNS peers periodically (fails safely if already connected)
if self.retry_delay.poll(cx).is_ready() {
for (p, mas) in self.mdns_discovered.clone() {
for ma in mas {
self.dial(p, ma)
}
}
// dial bootstrap peers (for environments where mDNS is unavailable)
for addr in &self.bootstrap_peers {
self.pending_events.push_back(ToSwarm::Dial {
opts: DialOpts::unknown_peer_id().address(addr.clone()).build(),
})
}
self.retry_delay.reset(RETRY_CONNECT_INTERVAL) // reset timeout
}
// send out any pending events from our own service
if let Some(e) = self.pending_events.pop_front(cx) {
return Poll::Ready(e.map_in(Either::Left));
}
// wait for pending events
Poll::Pending
}
}
+96 -36
View File
@@ -1,44 +1,104 @@
//! TODO: crate documentation
//!
//! this is here as a placeholder documentation
//!
//!
pub mod discovery;
use std::{env, panic, sync::Arc};
use netwatcher::WatchHandle;
use parking_lot::Mutex;
use tokio::{sync::mpsc, task::JoinHandle};
use zenoh::{Result, Session as ZSession, config::WhatAmI, internal::runtime::Runtime};
use zenoh_plugin_storage_manager::StoragesPlugin;
use zenoh_plugin_trait::PluginsManager;
pub use zenoh::{Config, config::ZenohId};
pub mod swarm;
/// Namespace for all the type/trait aliases used by this crate.
pub(crate) mod alias {
use std::error::Error;
pub type AnyError = Box<dyn Error + Send + Sync + 'static>;
pub type AnyResult<T> = Result<T, AnyError>;
pub fn cfg(identity: u128, listen_port: u16) -> Result<zenoh::Config> {
let namespace = env::var("EXO_ZENOH_NAMESPACE").unwrap_or_else(|_| "exo".to_string());
let mut cfg = zenoh::Config::default();
// todo: cleanup
cfg.insert_json5("id", &format!("\"{identity:x}\""))?;
cfg.insert_json5("mode", "\"peer\"")?;
cfg.insert_json5("listen/endpoints", &format!("[\"tcp/[::]:{listen_port}\"]"))?;
cfg.insert_json5("scouting/multicast/enabled", "true")?;
cfg.insert_json5("scouting/multicast/autoconnect", "[]")?;
cfg.insert_json5("scouting/gossip/multihop", "true")?;
cfg.insert_json5("namespace", &format!("{namespace:?}"))?;
cfg.insert_json5("transport/link/tx/batch_size", "9216")?;
cfg.insert_json5("timestamping/enabled", "true")?;
cfg.insert_json5("plugins/storage_manager/__required__", "true")?;
cfg.insert_json5(
"plugins/storage_manager/storages/mem1",
r#"{
key_expr: "storage/mem1/**",
strip_prefix: "storage/mem1",
volume: "memory",
replication: {
interval: 2,
}
}"#,
)?;
Ok(cfg)
}
/// Namespace for crate-wide extension traits/methods
pub(crate) mod ext {
use extend::ext;
use libp2p::Multiaddr;
use libp2p::multiaddr::Protocol;
use std::net::IpAddr;
#[ext(pub, name = MultiaddrExt)]
impl Multiaddr {
/// If the multiaddress corresponds to a TCP address, extracts it
fn try_to_tcp_addr(&self) -> Option<(IpAddr, u16)> {
let mut ps = self.into_iter();
let ip = if let Some(p) = ps.next() {
match p {
Protocol::Ip4(ip) => IpAddr::V4(ip),
Protocol::Ip6(ip) => IpAddr::V6(ip),
_ => return None,
pub async fn open(cfg: zenoh::Config) -> Result<Session> {
let mut plugins = PluginsManager::static_plugins_only();
plugins.declare_static_plugin::<StoragesPlugin, _>("storage_manager", true);
let mut runtime = zenoh::internal::runtime::RuntimeBuilder::new(cfg)
.plugins_manager(plugins)
.build()
.await?;
let z = zenoh::session::init(runtime.clone().into()).await?;
runtime.start().await?;
let _watch_all_handle = watch_all(runtime).await?;
Ok(Session {
z,
_watch_all_handle,
})
}
async fn watch_all(runtime: Runtime) -> Result<WatchAllHandle> {
log::info!("spawning scout");
let mut cfg = Config::default();
cfg.insert_json5("scouting/multicast/ttl", "3")?;
cfg.insert_json5("scouting/multicast/interface", "\"auto\"")?;
let mut scout = zenoh::scout(WhatAmI::Peer, cfg.clone()).await?;
let (send, mut recv) = mpsc::unbounded_channel();
let _sync = Arc::new(Mutex::new(netwatcher::watch_interfaces_with_callback(
move |u| _ = send.send(u),
)?));
let _async = tokio::task::spawn(async move {
loop {
tokio::select! {
u = recv.recv() => {
if u.is_none() {
return Ok(());
}
log::info!("reloading scout");
scout = zenoh::scout(WhatAmI::Peer, cfg.clone()).await?;
}
} else {
return None;
};
let Some(Protocol::Tcp(port)) = ps.next() else {
return None;
};
Some((ip, port))
hello = scout.recv_async() => {
if let Ok(hello) = hello {
// todo: auth
runtime
.connect_peer(&hello.zid().into(), hello.locators())
.await;
}
}
}
}
});
Ok(WatchAllHandle { _sync, _async })
}
pub struct Session {
pub z: ZSession,
_watch_all_handle: WatchAllHandle,
}
impl Drop for WatchAllHandle {
fn drop(&mut self) {
self._async.abort();
}
}
pub struct WatchAllHandle {
_sync: Arc<Mutex<WatchHandle>>,
_async: JoinHandle<Result<()>>,
}
+132 -230
View File
@@ -1,24 +1,20 @@
//! Compat shim for the old libp2p code
use std::collections::HashMap;
use std::pin::Pin;
use crate::swarm::transport::tcp_transport;
use crate::{alias, discovery};
pub use behaviour::{Behaviour, BehaviourEvent};
use futures_lite::{Stream, StreamExt};
use libp2p::{PeerId, SwarmBuilder, gossipsub, identity, swarm::SwarmEvent};
use tokio::sync::{mpsc, oneshot};
use futures_lite::Stream;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use zenoh::Result;
use zenoh::Session;
use zenoh::handlers::FifoChannelHandler;
use zenoh::liveliness::LivelinessToken;
use zenoh::pubsub::Subscriber;
use zenoh::sample::Sample;
use zenoh::sample::SampleKind;
/// The current version of the network: this prevents devices running different versions of the
/// software from interacting with each other.
///
/// TODO: right now this is a hardcoded constant; figure out what the versioning semantics should
/// even be, and how to inject the right version into this config/initialization. E.g. should
/// this be passed in as a parameter? What about rapidly changing versions in debug builds?
/// this is all VERY very hard to figure out and needs to be mulled over as a team.
pub const NETWORK_VERSION: &[u8] = b"v0.0.1";
pub const OVERRIDE_VERSION_ENV_VAR: &str = "EXO_LIBP2P_NAMESPACE";
// Uses oneshot senders to emulate function calling apis while avoiding requiring unique ownership
// of the Swarm.
#[derive(Debug)]
pub enum ToSwarm {
Unsubscribe {
topic: String,
@@ -26,52 +22,66 @@ pub enum ToSwarm {
},
Subscribe {
topic: String,
result_sender: oneshot::Sender<Result<bool, gossipsub::SubscriptionError>>,
result_sender: oneshot::Sender<Result<bool>>,
},
Publish {
topic: String,
data: Vec<u8>,
result_sender: oneshot::Sender<Result<gossipsub::MessageId, gossipsub::PublishError>>,
result_sender: oneshot::Sender<Result<()>>,
},
}
#[derive(Debug)]
pub enum FromSwarm {
Message {
from: PeerId,
topic: String,
data: Vec<u8>,
},
Discovered {
peer_id: PeerId,
},
Expired {
peer_id: PeerId,
},
Message { topic: String, data: Vec<u8> },
Discovered {},
Expired {},
}
pub type Topics = HashMap<String, Subscriber<()>>;
pub struct Swarm {
swarm: libp2p::Swarm<Behaviour>,
pub session: crate::Session,
from_client: mpsc::Receiver<ToSwarm>,
}
impl Swarm {
pub fn into_stream(self) -> Pin<Box<dyn Stream<Item = FromSwarm> + Send>> {
let Swarm {
mut swarm,
session,
mut from_client,
} = self;
let stream = async_stream::stream! {
let mut session = session;
let (mut to_topics, mut from_topics) = mpsc::channel(1024);
let mut topics = Topics::new();
let Ok((_token, discovery)) = register_liveness(&mut session.z).await else { return; };
loop {
tokio::select! {
msg = from_client.recv() => {
let Some(msg) = msg else { break };
on_message(&mut swarm, msg);
on_message(&mut session.z, &mut topics, &mut to_topics, msg).await;
}
event = swarm.next() => {
let Some(event) = event else { break };
if let Some(item) = filter_swarm_event(event) {
yield item;
event = from_topics.recv() => {
if let Some(event) = event {
yield event
}
}
token = discovery.recv_async() => {
if let Ok(token) = token {
let key_expr = token.key_expr().as_str().to_owned();
let nid = key_expr.strip_prefix("nodes/").and_then(|s| s.strip_suffix("/live"));
yield match token.kind() {
SampleKind::Put => {
log::info!("discovered: {nid:?}");
FromSwarm::Discovered {}
}
SampleKind::Delete => {
log::info!("expired: {nid:?}");
FromSwarm::Expired {}
}
}
}
}
}
}
};
@@ -79,208 +89,100 @@ impl Swarm {
}
}
fn on_message(swarm: &mut libp2p::Swarm<Behaviour>, message: ToSwarm) {
match message {
ToSwarm::Subscribe {
topic,
result_sender,
} => {
let result = swarm
.behaviour_mut()
.gossipsub
.subscribe(&gossipsub::IdentTopic::new(topic));
_ = result_sender.send(result);
}
ToSwarm::Unsubscribe {
topic,
result_sender,
} => {
let result = swarm
.behaviour_mut()
.gossipsub
.unsubscribe(&gossipsub::IdentTopic::new(topic));
_ = result_sender.send(result);
}
async fn register_liveness(
session: &mut Session,
) -> Result<(LivelinessToken, Subscriber<FifoChannelHandler<Sample>>)> {
let token = session
.liveliness()
.declare_token(format!("nodes/{}/live", session.zid()))
.await?;
let sub = session
.liveliness()
.declare_subscriber("nodes/*/live")
.history(true)
.await?;
Ok((token, sub))
}
async fn on_message(
session: &mut Session,
topics: &mut Topics,
to_topics: &mut mpsc::Sender<FromSwarm>,
msg: ToSwarm,
) {
match msg {
ToSwarm::Publish {
topic,
data,
result_sender,
} => {
let result = swarm
.behaviour_mut()
.gossipsub
.publish(gossipsub::IdentTopic::new(topic), data);
_ = result_sender.send(result);
let res = session.put(format!("topics/{topic}"), data).await;
_ = result_sender.send(res);
}
ToSwarm::Unsubscribe {
topic,
result_sender,
} => {
let Some((_, subscriber)) = topics.remove_entry(&topic) else {
_ = result_sender.send(false);
return;
};
_ = subscriber.undeclare().await;
_ = result_sender.send(true);
}
ToSwarm::Subscribe {
topic,
result_sender,
} => {
assert!(topic.is_ascii());
if topics.contains_key(&topic) {
_ = result_sender.send(Ok(false));
return;
}
let subscriber = match session
.declare_subscriber(format!("topics/{topic}"))
.allowed_origin(zenoh::sample::Locality::Remote)
.callback({
let sender = to_topics.clone();
let topic = topic.clone();
move |sample| {
if sample.kind() != SampleKind::Put {
return;
}
_ = sender.try_send(FromSwarm::Message {
topic: topic.clone(),
data: sample.payload().to_bytes().to_vec(),
});
}
})
.await
{
Ok(p) => p,
Err(e) => {
_ = result_sender.send(Err(e));
return;
}
};
assert!(topics.insert(topic, subscriber).is_none());
_ = result_sender.send(Ok(true));
}
}
}
fn filter_swarm_event(event: SwarmEvent<BehaviourEvent>) -> Option<FromSwarm> {
match event {
SwarmEvent::Behaviour(BehaviourEvent::Gossipsub(gossipsub::Event::Message {
message:
gossipsub::Message {
source: Some(peer_id),
topic,
data,
..
},
..
})) => Some(FromSwarm::Message {
from: peer_id,
topic: topic.into_string(),
data,
}),
SwarmEvent::Behaviour(BehaviourEvent::Discovery(
discovery::Event::ConnectionEstablished { peer_id, .. },
)) => Some(FromSwarm::Discovered { peer_id }),
SwarmEvent::Behaviour(BehaviourEvent::Discovery(discovery::Event::ConnectionClosed {
peer_id,
..
})) => Some(FromSwarm::Expired { peer_id }),
_ => None,
}
}
/// Create and configure a swarm.
///
/// - `listen_port`: TCP port to listen on. `0` lets the OS assign one.
/// - `bootstrap_peers`: multiaddrs to dial for environments without mDNS.
pub fn create_swarm(
keypair: identity::Keypair,
pub async fn create_swarm(
identity: u128,
from_client: mpsc::Receiver<ToSwarm>,
bootstrap_peers: Vec<String>,
listen_port: u16,
) -> alias::AnyResult<Swarm> {
let parsed_bootstrap_peers: Vec<libp2p::Multiaddr> = bootstrap_peers
.iter()
.filter(|s| !s.is_empty())
.filter_map(|s| s.parse().ok())
.collect();
let mut swarm = SwarmBuilder::with_existing_identity(keypair)
.with_tokio()
.with_other_transport(tcp_transport)?
.with_behaviour(|keypair| Behaviour::new(keypair, parsed_bootstrap_peers))?
.build();
swarm.listen_on(format!("/ip4/0.0.0.0/tcp/{listen_port}").parse()?)?;
Ok(Swarm { swarm, from_client })
}
mod transport {
use crate::alias;
use crate::swarm::{NETWORK_VERSION, OVERRIDE_VERSION_ENV_VAR};
use futures_lite::{AsyncRead, AsyncWrite};
use keccak_const::Sha3_256;
use libp2p::core::muxing;
use libp2p::core::transport::Boxed;
use libp2p::pnet::{PnetError, PnetOutput};
use libp2p::{PeerId, Transport, identity, noise, pnet, yamux};
use std::{env, sync::LazyLock};
/// Key used for networking's private network; parametrized on the [`NETWORK_VERSION`].
/// See [`pnet_upgrade`] for more.
static PNET_PRESHARED_KEY: LazyLock<[u8; 32]> = LazyLock::new(|| {
let builder = Sha3_256::new().update(b"exo_discovery_network");
if let Ok(var) = env::var(OVERRIDE_VERSION_ENV_VAR) {
let bytes = var.into_bytes();
builder.update(&bytes)
} else {
builder.update(NETWORK_VERSION)
}
.finalize()
});
/// Make the Swarm run on a private network, as to not clash with public libp2p nodes and
/// also different-versioned instances of this same network.
/// This is implemented as an additional "upgrade" ontop of existing [`libp2p::Transport`] layers.
async fn pnet_upgrade<TSocket>(
socket: TSocket,
_: impl Sized,
) -> Result<PnetOutput<TSocket>, PnetError>
where
TSocket: AsyncRead + AsyncWrite + Send + Unpin + 'static,
{
use pnet::{PnetConfig, PreSharedKey};
PnetConfig::new(PreSharedKey::new(*PNET_PRESHARED_KEY))
.handshake(socket)
.await
}
/// TCP/IP transport layer configuration.
pub fn tcp_transport(
keypair: &identity::Keypair,
) -> alias::AnyResult<Boxed<(PeerId, muxing::StreamMuxerBox)>> {
use libp2p::{
core::upgrade::Version,
tcp::{Config, tokio},
};
// `TCP_NODELAY` enabled => avoid latency
let tcp_config = Config::default().nodelay(true);
// V1 + lazy flushing => 0-RTT negotiation
let upgrade_version = Version::V1Lazy;
// Noise is faster than TLS + we don't care much for security
let noise_config = noise::Config::new(keypair)?;
// Use default Yamux config for multiplexing
let yamux_config = yamux::Config::default();
// Create new Tokio-driven TCP/IP transport layer
let base_transport = tokio::Transport::new(tcp_config)
.and_then(pnet_upgrade)
.upgrade(upgrade_version)
.authenticate(noise_config)
.multiplex(yamux_config);
// Return boxed transport (to flatten complex type)
Ok(base_transport.boxed())
}
}
mod behaviour {
use crate::{alias, discovery};
use libp2p::swarm::NetworkBehaviour;
use libp2p::{gossipsub, identity};
/// Behavior of the Swarm which composes all desired behaviors:
/// Right now its just [`discovery::Behaviour`] and [`gossipsub::Behaviour`].
#[derive(NetworkBehaviour)]
pub struct Behaviour {
pub discovery: discovery::Behaviour,
pub gossipsub: gossipsub::Behaviour,
}
impl Behaviour {
pub fn new(
keypair: &identity::Keypair,
bootstrap_peers: Vec<libp2p::Multiaddr>,
) -> alias::AnyResult<Self> {
Ok(Self {
discovery: discovery::Behaviour::new(keypair, bootstrap_peers)?,
gossipsub: gossipsub_behaviour(keypair),
})
}
}
fn gossipsub_behaviour(keypair: &identity::Keypair) -> gossipsub::Behaviour {
use gossipsub::{ConfigBuilder, MessageAuthenticity, ValidationMode};
// build a gossipsub network behaviour
// => signed message authenticity + strict validation mode means the message-ID is
// automatically provided by gossipsub w/out needing to provide custom message-ID function
gossipsub::Behaviour::new(
MessageAuthenticity::Signed(keypair.clone()),
ConfigBuilder::default()
.max_transmit_size(8 * 1024 * 1024)
.validation_mode(ValidationMode::Strict)
.build()
.expect("the configuration should always be valid"),
)
.expect("creating gossipsub behavior should always work")
}
) -> Result<Swarm> {
// todo: bootstrap
if !bootstrap_peers.is_empty() || listen_port != 0 {
todo!();
}
let cfg = crate::cfg(identity, listen_port)?;
let session = crate::open(cfg).await?;
Ok(Swarm {
session,
from_client,
})
}
-107
View File
@@ -1,107 +0,0 @@
use futures_lite::StreamExt;
use networking::swarm::{FromSwarm, create_swarm};
use std::time::Duration;
use tokio::sync::mpsc;
use tokio::time::timeout;
/// Helper: find a free TCP port.
fn free_port() -> u16 {
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
listener.local_addr().unwrap().port()
}
/// Two nodes connect via bootstrap peers — no mDNS needed.
///
/// Node A listens on a fixed port. Node B bootstraps to A's address.
/// We verify that B emits `FromSwarm::Discovered` for A's peer ID.
#[tokio::test]
async fn two_nodes_connect_via_bootstrap_peers() {
let port_a = free_port();
// Node A: listens on a known port, no bootstrap peers
let keypair_a = libp2p::identity::Keypair::generate_ed25519();
let peer_id_a = keypair_a.public().to_peer_id();
let (_tx_a, rx_a) = mpsc::channel(16);
let swarm_a = create_swarm(keypair_a, rx_a, vec![], port_a).expect("create swarm A");
let mut stream_a = swarm_a.into_stream();
// Node B: bootstraps to A's address
let keypair_b = libp2p::identity::Keypair::generate_ed25519();
let (_tx_b, rx_b) = mpsc::channel(16);
let swarm_b = create_swarm(
keypair_b,
rx_b,
vec![format!("/ip4/127.0.0.1/tcp/{port_a}")],
0,
)
.expect("create swarm B");
let mut stream_b = swarm_b.into_stream();
// Wait for B to discover A (connection established)
let connected = timeout(Duration::from_secs(10), async {
loop {
tokio::select! {
Some(event) = stream_a.next() => {
// A will also see B connect, but we check from B's perspective
let _ = event;
}
Some(event) = stream_b.next() => {
if let FromSwarm::Discovered { peer_id } = event {
if peer_id == peer_id_a {
return true;
}
}
}
}
}
})
.await;
assert!(
connected.is_ok() && connected.unwrap(),
"Node B should discover Node A via bootstrap peer"
);
}
/// Empty bootstrap peers should work (backward compatible).
#[tokio::test]
async fn create_swarm_with_empty_bootstrap_peers() {
let keypair = libp2p::identity::Keypair::generate_ed25519();
let (_tx, rx) = mpsc::channel(16);
let swarm = create_swarm(keypair, rx, vec![], 0);
assert!(
swarm.is_ok(),
"create_swarm with no bootstrap peers should succeed"
);
}
/// Invalid multiaddr strings are silently filtered out.
#[tokio::test]
async fn create_swarm_ignores_invalid_bootstrap_addrs() {
let keypair = libp2p::identity::Keypair::generate_ed25519();
let (_tx, rx) = mpsc::channel(16);
let swarm = create_swarm(
keypair,
rx,
vec![
"not-a-valid-multiaddr".to_string(),
"".to_string(),
"/ip4/10.0.0.1/tcp/30000".to_string(), // valid
],
0,
);
assert!(
swarm.is_ok(),
"create_swarm should succeed even with invalid bootstrap addrs"
);
}
/// Fixed listen port works correctly.
#[tokio::test]
async fn create_swarm_with_fixed_port() {
let port = free_port();
let keypair = libp2p::identity::Keypair::generate_ed25519();
let (_tx, rx) = mpsc::channel(16);
let swarm = create_swarm(keypair, rx, vec![], port);
assert!(swarm.is_ok(), "create_swarm with fixed port should succeed");
}
-7
View File
@@ -1,7 +0,0 @@
// maybe this will hold test in the future...??
#[cfg(test)]
mod tests {
#[test]
fn does_nothing() {}
}
+4 -3
View File
@@ -55,6 +55,7 @@
];
OPENSSL_NO_VENDOR = "1";
MATURIN_NO_INSTALL_RUST = "1";
# Required for pyo3 tests to find libpython
LD_LIBRARY_PATH = lib.makeLibraryPath [ pkgs.python313 ];
@@ -81,11 +82,11 @@
config = {
packages = {
# Python bindings wheel via maturin
exo_pyo3_bindings = craneLib.buildPackage (
exo-net = craneLib.buildPackage (
commonArgs
// {
inherit cargoArtifacts;
pname = "exo_pyo3_bindings";
pname = "exo-net";
nativeBuildInputs = commonArgs.nativeBuildInputs ++ [
pkgs.maturin
@@ -95,7 +96,7 @@
maturin build \
--release \
--manylinux off \
--manifest-path rust/exo_pyo3_bindings/Cargo.toml \
--manifest-path rust/exo_net/Cargo.toml \
--features "pyo3/extension-module,pyo3/experimental-async" \
--interpreter ${pkgs.python313}/bin/python \
--out dist
-15
View File
@@ -1,15 +0,0 @@
[package]
name = "util"
version = { workspace = true }
edition = { workspace = true }
publish = false
[lib]
doctest = false
name = "util"
path = "src/lib.rs"
[lints]
workspace = true
[dependencies]
-1
View File
@@ -1 +0,0 @@
pub mod wakerdeque;
-55
View File
@@ -1,55 +0,0 @@
use std::collections::VecDeque;
use std::fmt::{Debug, Formatter};
use std::task::{Context, Waker};
/// A wrapper around [`VecDeque`] which wakes (if it can) on any `push_*` methods,
/// and updates the internally stored waker by consuming [`Context`] on any `pop_*` methods.
pub struct WakerDeque<T> {
waker: Option<Waker>,
deque: VecDeque<T>,
}
impl<T: Debug> Debug for WakerDeque<T> {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
self.deque.fmt(f)
}
}
impl<T> WakerDeque<T> {
pub fn new() -> Self {
Self {
waker: None,
deque: VecDeque::new(),
}
}
fn update(&mut self, cx: &mut Context<'_>) {
self.waker = Some(cx.waker().clone());
}
fn wake(&mut self) {
let Some(ref mut w) = self.waker else { return };
w.wake_by_ref();
self.waker = None;
}
pub fn pop_front(&mut self, cx: &mut Context<'_>) -> Option<T> {
self.update(cx);
self.deque.pop_front()
}
pub fn pop_back(&mut self, cx: &mut Context<'_>) -> Option<T> {
self.update(cx);
self.deque.pop_back()
}
pub fn push_front(&mut self, value: T) {
self.wake();
self.deque.push_front(value);
}
pub fn push_back(&mut self, value: T) {
self.wake();
self.deque.push_back(value);
}
}
-5
View File
@@ -20,7 +20,6 @@ from exo.shared.types.chunks import (
TokenChunk,
ToolCallChunk,
)
from exo.shared.types.common import CommandId
from exo.shared.types.text_generation import (
Base64Image,
InputMessage,
@@ -181,7 +180,6 @@ def ollama_request_to_text_generation(
async def generate_ollama_chat_stream(
_command_id: CommandId,
chunk_stream: AsyncGenerator[
ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
],
@@ -264,7 +262,6 @@ async def generate_ollama_chat_stream(
async def collect_ollama_chat_response(
_command_id: CommandId,
chunk_stream: AsyncGenerator[
ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
],
@@ -369,7 +366,6 @@ def ollama_generate_request_to_text_generation(
async def generate_ollama_generate_stream(
_command_id: CommandId,
chunk_stream: AsyncGenerator[
ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
],
@@ -442,7 +438,6 @@ async def generate_ollama_generate_stream(
async def collect_ollama_generate_response(
_command_id: CommandId,
chunk_stream: AsyncGenerator[
ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
],
+275 -312
View File
@@ -5,6 +5,7 @@ import json
import random
import time
from collections.abc import AsyncGenerator, Awaitable, Callable, Iterable
from dataclasses import dataclass, field
from datetime import datetime, timezone
from http import HTTPStatus
from pathlib import Path
@@ -12,7 +13,7 @@ from typing import Annotated, Any, Literal, cast
from uuid import uuid4
import anyio
from anyio import BrokenResourceError, ClosedResourceError
from exo_net import NetSender, PySession
from fastapi import FastAPI, File, Form, HTTPException, Query, Request, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
@@ -21,6 +22,7 @@ from hypercorn.asyncio import serve # pyright: ignore[reportUnknownVariableType
from hypercorn.config import Config
from hypercorn.typing import ASGIFramework
from loguru import logger
from pydantic import TypeAdapter
from exo.api.adapters.chat_completions import (
chat_request_to_text_generation,
@@ -133,19 +135,19 @@ from exo.shared.constants import (
)
from exo.shared.election import ElectionMessage
from exo.shared.logging import InterceptLogger
from exo.shared.models import model_cards
from exo.shared.models.model_cards import (
ModelCard,
ModelId,
add_to_card_cache,
get_card,
get_model_cards,
)
from exo.shared.tracing import TraceEvent, compute_stats, export_trace, load_trace_file
from exo.shared.types.chunks import (
ErrorChunk,
ImageChunk,
Chunk,
ImageGenerationChunk,
InputImageChunk,
PrefillProgressChunk,
StatusChunk,
TextGenerationChunk,
TokenChunk,
ToolCallChunk,
)
@@ -159,7 +161,6 @@ from exo.shared.types.commands import (
DeleteInstance,
DeleteInstanceLink,
DownloadCommand,
ForwarderCommand,
ForwarderDownloadCommand,
ImageEdits,
ImageGeneration,
@@ -173,7 +174,6 @@ from exo.shared.types.commands import (
)
from exo.shared.types.common import CommandId, Id, NodeId, SystemId
from exo.shared.types.events import (
ChunkGenerated,
Event,
IndexedEvent,
InstanceDeleted,
@@ -199,7 +199,7 @@ from exo.shared.types.worker.downloads import DownloadCompleted
from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
from exo.shared.types.worker.shards import Sharding
from exo.utils.banner import print_startup_banner
from exo.utils.channels import Receiver, Sender, channel
from exo.utils.channels import Receiver, Sender
from exo.utils.disk_event_log import DiskEventLog
from exo.utils.power_sampler import PowerSampler
from exo.utils.task_group import TaskGroup
@@ -232,6 +232,95 @@ def _require_disaggregation_enabled() -> None:
)
@dataclass
class Transport:
session: PySession
cancel_scopes: dict[CommandId, anyio.CancelScope] = field(
init=False, default_factory=dict
)
command_sender: NetSender = field(init=False)
paused: bool = field(init=False, default=False)
paused_ev: anyio.Event = field(init=False, default_factory=anyio.Event)
def __post_init__(self):
# TODO: retire root keyspace
self.command_sender = self.session.net_sender("orchestrator")
async def send_command(self, command: Command) -> bool:
while self.paused:
await self.paused_ev.wait()
return await self.command_sender.send(command.model_dump_json().encode("utf-8"))
async def stream_text(
self,
command_id: CommandId,
) -> AsyncGenerator[TextGenerationChunk | StatusChunk]:
async for chunk in self.stream(command_id):
if isinstance(chunk, (TextGenerationChunk | StatusChunk)):
yield chunk
async def stream_images(
self,
command_id: CommandId,
) -> AsyncGenerator[ImageGenerationChunk]:
async for chunk in self.stream(command_id):
if isinstance(chunk, (ImageGenerationChunk)):
yield chunk
async def stream(
self,
command_id: CommandId,
) -> AsyncGenerator[Chunk,]:
try:
with anyio.CancelScope() as cs:
self.cancel_scopes[command_id] = cs
# recv from any node
receiver = self.session.net_receiver(
f"runners/*/active_tasks/{command_id}/chunks"
)
while True:
data = await receiver.recv()
if data is None:
logger.warning(
"stream terminated early without finish reason EOF"
)
break
yield (
chunk := cast(
Chunk,
TypeAdapter(Chunk).validate_json(
data, strict=True, extra="forbid"
),
)
)
if (
not isinstance(chunk, StatusChunk)
and chunk.finish_reason is not None
):
break
except anyio.get_cancelled_exc_class():
with anyio.CancelScope(shield=True):
await self.command_sender.send(
TaskCancelled(cancelled_command_id=command_id)
.model_dump_json()
.encode("utf-8")
)
finally:
self.cancel_scopes.pop(command_id, None)
with anyio.CancelScope(shield=True):
await self.command_sender.send(
TaskFinished(finished_command_id=command_id)
.model_dump_json()
.encode("utf-8")
)
def cancel(self, command_id: CommandId) -> bool:
if (cs := self.cancel_scopes.get(command_id, None)) is not None:
cs.cancel()
return True
return False
class API:
def __init__(
self,
@@ -239,15 +328,14 @@ class API:
*,
port: int,
event_receiver: Receiver[IndexedEvent],
command_sender: Sender[ForwarderCommand],
download_command_sender: Sender[ForwarderDownloadCommand],
# This lets us pause the API if an election is running
election_receiver: Receiver[ElectionMessage],
session: PySession,
) -> None:
self.state = State()
self._event_log = DiskEventLog(_API_EVENT_LOG_DIR)
self._system_id = SystemId()
self.command_sender = command_sender
self.download_command_sender = download_command_sender
self.event_receiver = event_receiver
self.election_receiver = election_receiver
@@ -256,9 +344,6 @@ class API:
self.port = port
self._sent_image_hashes: set[str] = set()
self.paused: bool = False
self.paused_ev: anyio.Event = anyio.Event()
self.app = FastAPI()
@self.app.middleware("http")
@@ -282,13 +367,7 @@ class API:
name="dashboard",
)
self._text_generation_queues: dict[
CommandId,
Sender[TokenChunk | ErrorChunk | ToolCallChunk | PrefillProgressChunk],
] = {}
self._image_generation_queues: dict[
CommandId, Sender[ImageChunk | ErrorChunk]
] = {}
self.transport = Transport(session)
self._image_store = ImageStore(EXO_IMAGE_CACHE_DIR)
self._tg: TaskGroup = TaskGroup()
@@ -309,9 +388,9 @@ class API:
def unpause(self, result_clock: int):
logger.info("Unpausing API")
self.last_completed_election = result_clock
self.paused = False
self.paused_ev.set()
self.paused_ev = anyio.Event()
self.transport.paused = False
self.transport.paused_ev.set()
self.transport.paused_ev = anyio.Event()
def _setup_exception_handlers(self) -> None:
self.app.exception_handler(HTTPException)(self.http_exception_handler)
@@ -426,7 +505,7 @@ class API:
instance_meta=payload.instance_meta,
min_nodes=payload.min_nodes,
)
await self._send(command)
await self.transport.send_command(command)
return CreateInstanceResponse(
message="Command received.",
@@ -451,7 +530,7 @@ class API:
command = CreateInstance(
instance=instance,
)
await self._send(command)
await self.transport.send_command(command)
return CreateInstanceResponse(
message="Command received.",
@@ -481,6 +560,7 @@ class API:
topology=self.state.topology,
current_instances=self.state.instances,
download_status=self.state.downloads,
node_rdma_ctl=self.state.node_rdma_ctl,
)
except ValueError as exc:
raise HTTPException(status_code=400, detail=str(exc)) from exc
@@ -544,6 +624,7 @@ class API:
current_instances=self.state.instances,
required_nodes=required_nodes,
download_status=self.state.downloads,
node_rdma_ctl=self.state.node_rdma_ctl,
)
except ValueError as exc:
if (model_card.model_id, sharding, instance_meta, 0) not in seen:
@@ -632,7 +713,7 @@ class API:
command = DeleteInstance(
instance_id=instance_id,
)
await self._send(command)
await self.transport.send_command(command)
return DeleteInstanceResponse(
message="Command received.",
command_id=command.command_id,
@@ -667,7 +748,7 @@ class API:
prefill_instances=list(body.prefill_instances),
decode_instances=list(body.decode_instances),
)
await self._send(command)
await self.transport.send_command(command)
return InstanceLinkResponse(
message="Command received.", command_id=command.command_id
)
@@ -677,64 +758,24 @@ class API:
) -> InstanceLinkResponse:
_require_disaggregation_enabled()
command = DeleteInstanceLink(link_id=link_id)
await self._send(command)
await self.transport.send_command(command)
return InstanceLinkResponse(
message="Command received.", command_id=command.command_id
)
async def cancel_command(self, command_id: CommandId) -> CancelCommandResponse:
"""Cancel an active command by closing its stream and notifying workers."""
sender = self._text_generation_queues.get(
command_id
) or self._image_generation_queues.get(command_id)
if sender is None:
if self.transport.cancel(command_id):
return CancelCommandResponse(
message="Command cancelled.",
command_id=command_id,
)
else:
raise HTTPException(
status_code=404,
detail="Command not found or already completed",
)
await self._send(TaskCancelled(cancelled_command_id=command_id))
sender.close()
return CancelCommandResponse(
message="Command cancelled.",
command_id=command_id,
)
async def _token_chunk_stream(
self, command_id: CommandId
) -> AsyncGenerator[
TokenChunk | ErrorChunk | ToolCallChunk | PrefillProgressChunk, None
]:
"""Yield chunks for a given command until completion.
This is the internal low-level stream used by all API adapters.
"""
try:
self._text_generation_queues[command_id], recv = channel[
TokenChunk | ErrorChunk | ToolCallChunk | PrefillProgressChunk
]()
with recv as token_chunks:
async for chunk in token_chunks:
yield chunk
if isinstance(chunk, PrefillProgressChunk):
continue
if chunk.finish_reason is not None:
break
except anyio.get_cancelled_exc_class():
command = TaskCancelled(cancelled_command_id=command_id)
with anyio.CancelScope(shield=True):
await self.command_sender.send(
ForwarderCommand(origin=self._system_id, command=command)
)
raise
finally:
await self._send(TaskFinished(finished_command_id=command_id))
if command_id in self._text_generation_queues:
del self._text_generation_queues[command_id]
async def _collect_text_generation_with_stats(
self, command_id: CommandId
) -> BenchChatCompletionResponse:
@@ -749,7 +790,7 @@ class API:
async with anyio.create_task_group() as tg:
tg.start_soon(sampler.run)
async for chunk in self._token_chunk_stream(command_id):
async for chunk in self.transport.stream_text(command_id):
if isinstance(chunk, PrefillProgressChunk):
continue
@@ -816,7 +857,7 @@ class API:
images = task_params.images
if not images:
command = TextGeneration(task_params=task_params)
await self._send(command)
await self.transport.send_command(command)
return command
hashes = [hashlib.sha256(img.encode("ascii")).hexdigest() for img in images]
@@ -833,7 +874,7 @@ class API:
new_images.append((idx, img))
if not new_images:
await self._send(command)
await self.transport.send_command(command)
return command
all_chunks: list[tuple[int, str]] = []
@@ -842,7 +883,7 @@ class API:
all_chunks.append((img_idx, img_data[i : i + EXO_MAX_CHUNK_SIZE]))
for global_idx, (img_idx, chunk_data) in enumerate(all_chunks):
await self._send(
await self.transport.send_command(
SendInputChunk(
chunk=InputImageChunk(
model=task_params.model,
@@ -855,7 +896,7 @@ class API:
)
)
await self._send(command)
await self.transport.send_command(command)
return command
async def chat_completions(
@@ -875,7 +916,7 @@ class API:
with_sse_keepalive(
generate_chat_stream(
command.command_id,
self._token_chunk_stream(command.command_id),
self.transport.stream_text(command.command_id),
),
),
media_type="text/event-stream",
@@ -889,7 +930,7 @@ class API:
return StreamingResponse(
collect_chat_response(
command.command_id,
self._token_chunk_stream(command.command_id),
self.transport.stream_text(command.command_id),
),
media_type="application/json",
)
@@ -918,7 +959,7 @@ class API:
with_sse_keepalive(
generate_chat_stream(
command.command_id,
self._token_chunk_stream(command.command_id),
self.transport.stream_text(command.command_id),
),
),
media_type="text/event-stream",
@@ -1024,7 +1065,7 @@ class API:
command = ImageGeneration(
task_params=payload,
)
await self._send(command)
await self.transport.send_command(command)
# Check if streaming is requested
if payload.stream and payload.partial_images and payload.partial_images > 0:
@@ -1060,105 +1101,85 @@ class API:
image_metadata: dict[tuple[int, bool], tuple[int | None, int | None]] = {}
images_complete = 0
try:
self._image_generation_queues[command_id], recv = channel[
ImageChunk | ErrorChunk
]()
with recv as chunks:
async for chunk in chunks:
if chunk.finish_reason == "error":
error_response = ErrorResponse(
error=ErrorInfo(
message=chunk.error_message or "Internal server error",
type="InternalServerError",
code=500,
)
)
yield f"data: {error_response.model_dump_json()}\n\n"
yield "data: [DONE]\n\n"
return
key = (chunk.image_index, chunk.is_partial)
if key not in image_chunks:
image_chunks[key] = {}
image_total_chunks[key] = chunk.total_chunks
image_metadata[key] = (
chunk.partial_index,
chunk.total_partials,
)
image_chunks[key][chunk.chunk_index] = chunk.data
# Check if this image is complete
if len(image_chunks[key]) == image_total_chunks[key]:
full_data = "".join(
image_chunks[key][i] for i in range(len(image_chunks[key]))
)
partial_idx, total_partials = image_metadata[key]
if chunk.is_partial:
# Yield partial image event (always use b64_json for partials)
event_data = {
"type": "partial",
"image_index": chunk.image_index,
"partial_index": partial_idx,
"total_partials": total_partials,
"format": str(chunk.format),
"data": {
"b64_json": full_data
if response_format == "b64_json"
else None,
},
}
yield f"data: {json.dumps(event_data)}\n\n"
else:
# Final image
if response_format == "url":
image_bytes = base64.b64decode(full_data)
content_type = _format_to_content_type(chunk.format)
stored = self._image_store.store(
image_bytes, content_type
)
url = self._build_image_url(request, stored.image_id)
event_data = {
"type": "final",
"image_index": chunk.image_index,
"format": str(chunk.format),
"data": {"url": url},
}
else:
event_data = {
"type": "final",
"image_index": chunk.image_index,
"format": str(chunk.format),
"data": {"b64_json": full_data},
}
yield f"data: {json.dumps(event_data)}\n\n"
images_complete += 1
if images_complete >= num_images:
yield "data: [DONE]\n\n"
break
# Clean up completed image chunks
del image_chunks[key]
del image_total_chunks[key]
del image_metadata[key]
except anyio.get_cancelled_exc_class():
command = TaskCancelled(cancelled_command_id=command_id)
with anyio.CancelScope(shield=True):
await self.command_sender.send(
ForwarderCommand(origin=self._system_id, command=command)
async for chunk in self.transport.stream_images(command_id):
if chunk.finish_reason == "error":
error_response = ErrorResponse(
error=ErrorInfo(
message=chunk.error_message or "Internal server error",
type="InternalServerError",
code=500,
)
)
raise
finally:
await self._send(TaskFinished(finished_command_id=command_id))
if command_id in self._image_generation_queues:
del self._image_generation_queues[command_id]
yield f"data: {error_response.model_dump_json()}\n\n"
yield "data: [DONE]\n\n"
return
key = (chunk.image_index, chunk.is_partial)
if key not in image_chunks:
image_chunks[key] = {}
image_total_chunks[key] = chunk.total_chunks
image_metadata[key] = (
chunk.partial_index,
chunk.total_partials,
)
image_chunks[key][chunk.chunk_index] = chunk.data
# Check if this image is complete
if len(image_chunks[key]) == image_total_chunks[key]:
full_data = "".join(
image_chunks[key][i] for i in range(len(image_chunks[key]))
)
partial_idx, total_partials = image_metadata[key]
if chunk.is_partial:
# Yield partial image event (always use b64_json for partials)
event_data = {
"type": "partial",
"image_index": chunk.image_index,
"partial_index": partial_idx,
"total_partials": total_partials,
"format": str(chunk.format),
"data": {
"b64_json": full_data
if response_format == "b64_json"
else None,
},
}
yield f"data: {json.dumps(event_data)}\n\n"
else:
# Final image
if response_format == "url":
image_bytes = base64.b64decode(full_data)
content_type = _format_to_content_type(chunk.format)
stored = self._image_store.store(image_bytes, content_type)
url = self._build_image_url(request, stored.image_id)
event_data = {
"type": "final",
"image_index": chunk.image_index,
"format": str(chunk.format),
"data": {"url": url},
}
else:
event_data = {
"type": "final",
"image_index": chunk.image_index,
"format": str(chunk.format),
"data": {"b64_json": full_data},
}
yield f"data: {json.dumps(event_data)}\n\n"
images_complete += 1
if images_complete >= num_images:
yield "data: [DONE]\n\n"
break
# Clean up completed image chunks
del image_chunks[key]
del image_total_chunks[key]
del image_metadata[key]
async def _collect_image_chunks(
self,
@@ -1177,74 +1198,55 @@ class API:
images_complete = 0
stats: ImageGenerationStats | None = None
try:
self._image_generation_queues[command_id], recv = channel[
ImageChunk | ErrorChunk
]()
while images_complete < num_images:
with recv as chunks:
async for chunk in chunks:
if chunk.finish_reason == "error":
raise HTTPException(
status_code=500,
detail=chunk.error_message or "Internal server error",
)
if chunk.is_partial:
continue
if chunk.image_index not in image_chunks:
image_chunks[chunk.image_index] = {}
image_total_chunks[chunk.image_index] = chunk.total_chunks
image_formats[chunk.image_index] = chunk.format
image_chunks[chunk.image_index][chunk.chunk_index] = chunk.data
if capture_stats and chunk.stats is not None:
stats = chunk.stats
if (
len(image_chunks[chunk.image_index])
== image_total_chunks[chunk.image_index]
):
images_complete += 1
if images_complete >= num_images:
break
images: list[ImageData] = []
for image_idx in range(num_images):
chunks_dict = image_chunks[image_idx]
full_data = "".join(chunks_dict[i] for i in range(len(chunks_dict)))
if response_format == "url" and request is not None:
image_bytes = base64.b64decode(full_data)
content_type = _format_to_content_type(image_formats.get(image_idx))
stored = self._image_store.store(image_bytes, content_type)
url = self._build_image_url(request, stored.image_id)
images.append(ImageData(b64_json=None, url=url))
else:
images.append(
ImageData(
b64_json=full_data
if response_format == "b64_json"
else None,
url=None,
)
while images_complete < num_images:
async for chunk in self.transport.stream_images(command_id):
if chunk.finish_reason == "error":
raise HTTPException(
status_code=500,
detail=chunk.error_message or "Internal server error",
)
return (images, stats if capture_stats else None)
except anyio.get_cancelled_exc_class():
command = TaskCancelled(cancelled_command_id=command_id)
with anyio.CancelScope(shield=True):
await self.command_sender.send(
ForwarderCommand(origin=self._system_id, command=command)
if chunk.is_partial:
continue
if chunk.image_index not in image_chunks:
image_chunks[chunk.image_index] = {}
image_total_chunks[chunk.image_index] = chunk.total_chunks
image_formats[chunk.image_index] = chunk.format
image_chunks[chunk.image_index][chunk.chunk_index] = chunk.data
if capture_stats and chunk.stats is not None:
stats = chunk.stats
if (
len(image_chunks[chunk.image_index])
== image_total_chunks[chunk.image_index]
):
images_complete += 1
if images_complete >= num_images:
break
images: list[ImageData] = []
for image_idx in range(num_images):
chunks_dict = image_chunks[image_idx]
full_data = "".join(chunks_dict[i] for i in range(len(chunks_dict)))
if response_format == "url" and request is not None:
image_bytes = base64.b64decode(full_data)
content_type = _format_to_content_type(image_formats.get(image_idx))
stored = self._image_store.store(image_bytes, content_type)
url = self._build_image_url(request, stored.image_id)
images.append(ImageData(b64_json=None, url=url))
else:
images.append(
ImageData(
b64_json=full_data if response_format == "b64_json" else None,
url=None,
)
)
raise
finally:
await self._send(TaskFinished(finished_command_id=command_id))
if command_id in self._image_generation_queues:
del self._image_generation_queues[command_id]
return (images, stats if capture_stats else None)
async def _collect_image_generation(
self,
@@ -1294,7 +1296,7 @@ class API:
command = ImageGeneration(
task_params=payload,
)
await self._send(command)
await self.transport.send_command(command)
return await self._collect_image_generation_with_stats(
request=request,
@@ -1357,7 +1359,7 @@ class API:
f"Sending input image: {len(image_data)} bytes in {total_chunks} chunks"
)
for chunk_index, chunk_data in enumerate(data_chunks):
await self._send(
await self.transport.send_command(
SendInputChunk(
chunk=InputImageChunk(
model=resolved_model,
@@ -1369,7 +1371,7 @@ class API:
)
)
await self._send(command)
await self.transport.send_command(command)
return command
async def image_edits(
@@ -1497,7 +1499,7 @@ class API:
generate_claude_stream(
command.command_id,
payload.model,
self._token_chunk_stream(command.command_id),
self.transport.stream_text(command.command_id),
),
),
media_type="text/event-stream",
@@ -1512,7 +1514,7 @@ class API:
collect_claude_response(
command.command_id,
payload.model,
self._token_chunk_stream(command.command_id),
self.transport.stream_text(command.command_id),
),
media_type="application/json",
)
@@ -1533,7 +1535,7 @@ class API:
generate_responses_stream(
command.command_id,
payload.model,
self._token_chunk_stream(command.command_id),
self.transport.stream_text(command.command_id),
),
),
media_type="text/event-stream",
@@ -1549,7 +1551,7 @@ class API:
collect_responses_response(
command.command_id,
payload.model,
self._token_chunk_stream(command.command_id),
self.transport.stream_text(command.command_id),
),
media_type="application/json",
)
@@ -1575,8 +1577,7 @@ class API:
if payload.stream:
return StreamingResponse(
generate_ollama_chat_stream(
command.command_id,
self._token_chunk_stream(command.command_id),
self.transport.stream_text(command.command_id),
),
media_type="application/x-ndjson",
headers={
@@ -1588,8 +1589,7 @@ class API:
else:
return StreamingResponse(
collect_ollama_chat_response(
command.command_id,
self._token_chunk_stream(command.command_id),
self.transport.stream_text(command.command_id),
),
media_type="application/json",
)
@@ -1611,8 +1611,7 @@ class API:
if payload.stream:
return StreamingResponse(
generate_ollama_generate_stream(
command.command_id,
self._token_chunk_stream(command.command_id),
self.transport.stream_text(command.command_id),
),
media_type="application/x-ndjson",
headers={
@@ -1624,8 +1623,7 @@ class API:
else:
return StreamingResponse(
collect_ollama_generate_response(
command.command_id,
self._token_chunk_stream(command.command_id),
self.transport.stream_text(command.command_id),
),
media_type="application/json",
)
@@ -1633,17 +1631,16 @@ class API:
async def ollama_tags(self) -> OllamaTagsResponse:
"""Returns list of models in Ollama tags format. We return the downloaded ones only."""
def none_if_empty(value: str) -> str | None:
return value or None
downloaded_model_ids: set[str] = set()
downloaded_model_ids: set[ModelId] = set()
for node_downloads in self.state.downloads.values():
for dl in node_downloads:
if isinstance(dl, DownloadCompleted):
downloaded_model_ids.add(dl.shard_metadata.model_card.model_id)
cards = [
c for c in await get_model_cards() if c.model_id in downloaded_model_ids
c
for c in await model_cards.card_cache.list_all()
if c.model_id in downloaded_model_ids
]
now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
@@ -1656,8 +1653,8 @@ class API:
size=card.storage_size.in_bytes,
digest="sha256:000000000000",
details=OllamaModelDetails(
family=none_if_empty(card.family),
quantization_level=none_if_empty(card.quantization),
family=card.family or None,
quantization_level=card.quantization or None,
),
)
for card in cards
@@ -1720,7 +1717,7 @@ class API:
async def get_models(self, status: str | None = Query(default=None)) -> ModelList:
"""Returns list of available models, optionally filtered by being downloaded."""
cards = await get_model_cards()
cards = await model_cards.card_cache.list_all()
if status == "downloaded":
downloaded_model_ids: set[str] = set()
@@ -1762,16 +1759,13 @@ class API:
status_code=400, detail=f"Failed to fetch model: {exc}"
) from exc
await self.command_sender.send(
ForwarderCommand(
origin=self._system_id,
command=AddCustomModelCard(model_card=card),
)
await self.transport.command_sender.send(
AddCustomModelCard(model_card=card).model_dump_json().encode("utf-8")
)
# Immediately update the local cache so the subsequent GET /models
# returns the new model without waiting for the event round-trip.
add_to_card_cache(card)
model_cards.card_cache.cc[card.model_id] = card
return ModelListModel(
id=card.model_id,
@@ -1787,15 +1781,12 @@ class API:
async def delete_custom_model(self, model_id: ModelId) -> JSONResponse:
"""Delete a user-added custom model card and sync deletion across the cluster."""
card = get_card(model_id)
card = model_cards.card_cache.get(model_id)
if card is None or not card.is_custom:
raise HTTPException(status_code=404, detail="Custom model card not found")
await self.command_sender.send(
ForwarderCommand(
origin=self._system_id,
command=DeleteCustomModelCard(model_id=model_id),
)
await self.transport.command_sender.send(
DeleteCustomModelCard(model_id=model_id).model_dump_json().encode("utf-8")
)
return JSONResponse(
@@ -1860,7 +1851,6 @@ class API:
shutdown_ev.set()
finally:
self._event_log.close()
self.command_sender.close()
self.event_receiver.close()
async def run_api(self, ev: anyio.Event):
@@ -1884,23 +1874,6 @@ class API:
self.state = apply(self.state, i_event)
event = i_event.event
if isinstance(event, ChunkGenerated):
if queue := self._image_generation_queues.get(
event.command_id, None
):
assert isinstance(event.chunk, ImageChunk)
try:
await queue.send(event.chunk)
except (BrokenResourceError, ClosedResourceError):
self._image_generation_queues.pop(event.command_id, None)
if queue := self._text_generation_queues.get(
event.command_id, None
):
assert not isinstance(event.chunk, ImageChunk)
try:
await queue.send(event.chunk)
except (BrokenResourceError, ClosedResourceError):
self._text_generation_queues.pop(event.command_id, None)
if isinstance(event, InstanceDeleted):
self._close_streams_for_instance(event.instance_id)
if isinstance(event, TracesMerged):
@@ -1915,10 +1888,7 @@ class API:
task, (TextGenerationTask, ImageGenerationTask, ImageEditsTask)
):
continue
if sender := self._text_generation_queues.pop(task.command_id, None):
sender.close()
if sender := self._image_generation_queues.pop(task.command_id, None):
sender.close()
self.transport.cancel(task.command_id)
def _save_merged_trace(self, event: TracesMerged) -> None:
traces = [
@@ -1939,7 +1909,7 @@ class API:
with self.election_receiver as ems:
async for message in ems:
if message.clock > self.last_completed_election:
self.paused = True
self.transport.paused = True
async def _cleanup_expired_images(self):
"""Periodically clean up expired images from the store."""
@@ -1950,13 +1920,6 @@ class API:
if removed > 0:
logger.debug(f"Cleaned up {removed} expired images")
async def _send(self, command: Command):
while self.paused:
await self.paused_ev.wait()
await self.command_sender.send(
ForwarderCommand(origin=self._system_id, command=command)
)
async def _send_download(self, command: DownloadCommand):
await self.download_command_sender.send(
ForwarderDownloadCommand(origin=self._system_id, command=command)
+10 -14
View File
@@ -1,11 +1,11 @@
# pyright: reportUnusedFunction=false, reportAny=false
from typing import Any
from unittest.mock import AsyncMock, MagicMock
from unittest.mock import AsyncMock
from fastapi import FastAPI
from fastapi.testclient import TestClient
from exo.api.main import API
from exo.api.main import API, Transport
from exo.shared.types.common import CommandId
@@ -15,9 +15,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.transport = object.__new__(Transport)
api.transport.cancel = AsyncMock()
api.transport.send_command = AsyncMock()
api._setup_exception_handlers() # pyright: ignore[reportPrivateUsage]
app.post("/v1/cancel/{command_id}")(api.cancel_command)
return api
@@ -43,16 +43,14 @@ def test_cancel_active_text_generation() -> None:
client = TestClient(api.app)
cid = CommandId("text-cmd-123")
sender = MagicMock()
api._text_generation_queues[cid] = sender
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()
api.transport.cancel.assert_called_once()
api.transport.send_command.assert_called_once()
task_cancelled = api._send.call_args[0][0]
assert task_cancelled.cancelled_command_id == cid
@@ -63,15 +61,13 @@ def test_cancel_active_image_generation() -> None:
client = TestClient(api.app)
cid = CommandId("img-cmd-456")
sender = MagicMock()
api._image_generation_queues[cid] = sender
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]
api.transport.cancel.assert_called_once()
api.transport.send_command.assert_called_once()
task_cancelled = api.transport.send_command.call_args[0][0]
assert task_cancelled.cancelled_command_id == cid
@@ -1,9 +1,10 @@
# pyright: reportUnusedFunction=false, reportAny=false
"""Tests that InstanceDeleted events close active generation streams."""
from typing import Any
from unittest.mock import MagicMock
from exo.api.main import API
from exo.api.main import API, Transport
from exo.api.types import ImageGenerationTaskParams
from exo.shared.types.common import CommandId, ModelId
from exo.shared.types.state import State
@@ -16,12 +17,11 @@ from exo.shared.types.text_generation import (
from exo.shared.types.worker.instances import InstanceId
def _make_api_with_state(state: State) -> API:
def _make_api_with_state(state: State) -> Any:
"""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]
api.transport = object.__new__(Transport)
return api
@@ -47,13 +47,10 @@ def test_close_streams_for_deleted_instance() -> None:
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)
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]
api.transport.cancel.assert_called_once()
assert api.transport.cancel.call_args[0][0] == command_id
def test_close_streams_ignores_unrelated_instances() -> None:
@@ -72,7 +69,6 @@ def test_close_streams_ignores_unrelated_instances() -> None:
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:
+3 -2
View File
@@ -16,7 +16,8 @@ from exo.download.download_utils import (
)
from exo.download.shard_downloader import ShardDownloader
from exo.shared.constants import EXO_DEFAULT_MODELS_DIR, EXO_MODELS_READ_ONLY_DIRS
from exo.shared.models.model_cards import ModelId, get_model_cards
from exo.shared.models import model_cards
from exo.shared.models.model_cards import ModelId
from exo.shared.types.commands import (
CancelDownload,
DeleteDownload,
@@ -422,7 +423,7 @@ class DownloadCoordinator:
)
# Scan read-only directories for pre-downloaded models
if EXO_MODELS_READ_ONLY_DIRS:
for card in await get_model_cards():
for card in await model_cards.card_cache.list_all():
mid = card.model_id
if mid in self.active_downloads:
continue
+89 -25
View File
@@ -1,11 +1,12 @@
import asyncio
import hashlib
import os
import random
import shutil
import ssl
import time
import traceback
from collections.abc import Awaitable
from collections.abc import Awaitable, Mapping
from datetime import timedelta
from pathlib import Path
from typing import Callable, Literal
@@ -55,6 +56,36 @@ class HuggingFaceAuthenticationError(Exception):
class HuggingFaceRateLimitError(Exception):
"""429 Huggingface code"""
def __init__(self, msg: str, retry_after: float | None = None) -> None:
super().__init__(msg)
self.retry_after = retry_after
def _parse_retry_after(headers: Mapping[str, str]) -> float | None:
"""Parse seconds-to-reset from HF's RateLimit header.
HF sends e.g. ``ratelimit: "api";r=0;t=52`` on 429s; ``t`` is the wait.
Returns ``None`` if the header is missing or has no ``t`` field.
"""
raw = headers.get("RateLimit") or headers.get("ratelimit")
if raw is None:
return None
for part in raw.split(";"):
key, _, val = part.strip().partition("=")
if key == "t":
try:
return float(val)
except ValueError:
return None
return None
# reset window is 5 min
_RATE_LIMIT_MAX_SLEEP_SECS = 300.0
# 24h. Manually clear the cache (or `delete_model`) to force a refresh.
_FILE_LIST_CACHE_TTL_SECS = 24 * 60 * 60
async def _build_auth_error_message(status_code: int, model_id: ModelId) -> str:
token = await get_hf_token()
@@ -348,9 +379,6 @@ async def _build_file_list_from_local_directory(
return None
_fetched_file_lists_this_session: set[str] = set()
async def fetch_file_list_with_cache(
model_id: ModelId,
revision: str = "main",
@@ -360,13 +388,16 @@ async def fetch_file_list_with_cache(
) -> list[FileListEntry]:
target_dir = await ensure_cache_dir(model_id)
cache_file = target_dir / f"{model_id.normalize()}--{revision}--file_list.json"
cache_key = f"{model_id.normalize()}--{revision}"
if cache_key in _fetched_file_lists_this_session and await aios.path.exists(
cache_file
):
async with aiofiles.open(cache_file, "r") as f:
return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
# cache survives process restarts so cold starts don't re-burst HF
if await aios.path.exists(cache_file):
try:
cache_age = time.time() - (await aios.stat(cache_file)).st_mtime
except OSError:
cache_age = float("inf")
if cache_age < _FILE_LIST_CACHE_TTL_SECS:
async with aiofiles.open(cache_file, "r") as f:
return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
if skip_internet:
if await aios.path.exists(cache_file):
@@ -395,7 +426,6 @@ async def fetch_file_list_with_cache(
await f.write(
TypeAdapter(list[FileListEntry]).dump_json(file_list).decode()
)
_fetched_file_lists_this_session.add(cache_key)
return file_list
except Exception as e:
logger.opt(exception=e).warning(
@@ -426,17 +456,29 @@ async def fetch_file_list_with_retry(
recursive: bool = False,
on_connection_lost: Callable[[], None] = lambda: None,
) -> list[FileListEntry]:
n_attempts = 3
n_attempts = 5
for attempt in range(n_attempts):
try:
return await _fetch_file_list(model_id, revision, path, recursive)
except HuggingFaceAuthenticationError:
raise
except HuggingFaceRateLimitError as e:
if attempt == n_attempts - 1:
raise
sleep_for = e.retry_after if e.retry_after is not None else 2.0**attempt
sleep_for = min(sleep_for, _RATE_LIMIT_MAX_SLEEP_SECS) + random.uniform(
0, 1
)
logger.warning(
f"Rate limited by HuggingFace fetching file list for {model_id}; "
f"sleeping {sleep_for:.1f}s before retry {attempt + 2}/{n_attempts}"
)
await asyncio.sleep(sleep_for)
except Exception as e:
on_connection_lost()
if attempt == n_attempts - 1:
raise e
await asyncio.sleep(2.0**attempt)
await asyncio.sleep(2.0**attempt + random.uniform(0, 1))
raise Exception(
f"Failed to fetch file list for {model_id=} {revision=} {path=} {recursive=}"
)
@@ -447,6 +489,9 @@ async def _fetch_file_list(
) -> list[FileListEntry]:
api_url = f"{get_hf_endpoint()}/api/models/{model_id}/tree/{revision}"
url = f"{api_url}/{path}" if path else api_url
# ?recursive=true returns the whole subtree in one request
if recursive:
url = f"{url}?recursive=true"
headers = await get_download_headers()
async with (
@@ -458,7 +503,8 @@ async def _fetch_file_list(
raise HuggingFaceAuthenticationError(msg)
elif response.status == 429:
raise HuggingFaceRateLimitError(
f"Couldn't download {model_id} because of HuggingFace rate limit."
f"HuggingFace rate limit hit fetching file list for {model_id}",
retry_after=_parse_retry_after(response.headers),
)
elif response.status == 200:
data_json = await response.text()
@@ -468,10 +514,14 @@ async def _fetch_file_list(
if item.type == "file":
files.append(FileListEntry.model_validate(item))
elif item.type == "directory" and recursive:
subfiles = await _fetch_file_list(
model_id, revision, item.path, recursive
)
files.extend(subfiles)
# already inlined by ?recursive=true
continue
if recursive and len(data) >= 1000:
# HF tree endpoint paginates at 1000; we don't follow cursors
logger.warning(
f"File list for {model_id} hit the 1000-entry page cap "
"and may be truncated; cursor pagination is not implemented"
)
return files
else:
raise Exception(f"Failed to fetch file list: {response.status}")
@@ -552,6 +602,11 @@ async def file_meta(
if r.status in [401, 403]:
msg = await _build_auth_error_message(r.status, model_id)
raise HuggingFaceAuthenticationError(msg)
if r.status == 429:
raise HuggingFaceRateLimitError(
f"HuggingFace rate limit hit fetching metadata for {model_id}/{path}",
retry_after=_parse_retry_after(r.headers),
)
content_length = int(
r.headers.get("x-linked-size") or r.headers.get("content-length") or 0
)
@@ -571,7 +626,7 @@ async def download_file_with_retry(
on_connection_lost: Callable[[], None] = lambda: None,
skip_internet: bool = False,
) -> Path:
n_attempts = 3
n_attempts = 5
for attempt in range(n_attempts):
try:
return await _download_file(
@@ -583,12 +638,16 @@ async def download_file_with_retry(
raise
except HuggingFaceRateLimitError as e:
if attempt == n_attempts - 1:
raise e
logger.error(
f"Download error on attempt {attempt}/{n_attempts} for {model_id=} {revision=} {path=} {target_dir=}"
raise
sleep_for = e.retry_after if e.retry_after is not None else 2.0**attempt
sleep_for = min(sleep_for, _RATE_LIMIT_MAX_SLEEP_SECS) + random.uniform(
0, 1
)
logger.error(traceback.format_exc())
await asyncio.sleep(2.0**attempt)
logger.warning(
f"Rate limited by HuggingFace downloading {model_id}/{path}; "
f"sleeping {sleep_for:.1f}s before retry {attempt + 2}/{n_attempts}"
)
await asyncio.sleep(sleep_for)
except Exception as e:
if attempt == n_attempts - 1:
on_connection_lost()
@@ -597,7 +656,7 @@ async def download_file_with_retry(
f"Download error on attempt {attempt + 1}/{n_attempts} for {model_id=} {revision=} {path=} {target_dir=}"
)
logger.error(traceback.format_exc())
await asyncio.sleep(2.0**attempt)
await asyncio.sleep(2.0**attempt + random.uniform(0, 1))
raise Exception(
f"Failed to download file {model_id=} {revision=} {path=} {target_dir=}"
)
@@ -665,6 +724,11 @@ async def _download_file(
if r.status in [401, 403]:
msg = await _build_auth_error_message(r.status, model_id)
raise HuggingFaceAuthenticationError(msg)
if r.status == 429:
raise HuggingFaceRateLimitError(
f"HuggingFace rate limit hit downloading {model_id}/{path}",
retry_after=_parse_retry_after(r.headers),
)
assert r.status in [200, 206], (
f"Failed to download {path} from {url}: {r.status}"
)
+2 -2
View File
@@ -11,11 +11,11 @@ from exo.download.download_utils import (
download_shard,
)
from exo.download.shard_downloader import ShardDownloader
from exo.shared.models import model_cards
from exo.shared.models.model_cards import (
ModelCard,
ModelId,
ModelTask,
get_model_cards,
)
from exo.shared.types.memory import Memory
from exo.shared.types.worker.shards import (
@@ -258,7 +258,7 @@ class ResumableShardDownloader(ShardDownloader):
tasks = [
create_task(download_with_semaphore(model_card))
for model_card in await get_model_cards()
for model_card in await model_cards.card_cache.list_all()
]
for task in asyncio.as_completed(tasks):
@@ -1,5 +1,7 @@
"""Tests for offline/air-gapped mode."""
import os
import time
from collections.abc import AsyncIterator
from pathlib import Path
from unittest.mock import AsyncMock, patch
@@ -231,3 +233,64 @@ class TestFetchFileListOffline:
raise FileNotFoundError."""
with pytest.raises(FileNotFoundError, match="No internet"):
await fetch_file_list_with_cache(model_id, "main", skip_internet=True)
class TestFileListCacheTTL:
async def test_uses_fresh_cache_without_fetching(
self, model_id: ModelId, temp_models_dir: Path
) -> None:
from pydantic import TypeAdapter
cache_dir = temp_models_dir / "caches" / model_id.normalize()
await aios.makedirs(cache_dir, exist_ok=True)
cached_list = [
FileListEntry(type="file", path="model.safetensors", size=1000),
]
cache_file = cache_dir / f"{model_id.normalize()}--main--file_list.json"
async with aiofiles.open(cache_file, "w") as f:
await f.write(
TypeAdapter(list[FileListEntry]).dump_json(cached_list).decode()
)
with patch(
"exo.download.download_utils.fetch_file_list_with_retry",
new_callable=AsyncMock,
) as mock_fetch:
result = await fetch_file_list_with_cache(model_id, "main")
assert result == cached_list
mock_fetch.assert_not_called()
async def test_refetches_when_cache_older_than_ttl(
self, model_id: ModelId, temp_models_dir: Path
) -> None:
from pydantic import TypeAdapter
from exo.download.download_utils import (
_FILE_LIST_CACHE_TTL_SECS, # pyright: ignore[reportPrivateUsage]
)
cache_dir = temp_models_dir / "caches" / model_id.normalize()
await aios.makedirs(cache_dir, exist_ok=True)
stale_list = [FileListEntry(type="file", path="stale.bin", size=1)]
cache_file = cache_dir / f"{model_id.normalize()}--main--file_list.json"
async with aiofiles.open(cache_file, "w") as f:
await f.write(
TypeAdapter(list[FileListEntry]).dump_json(stale_list).decode()
)
old_mtime = time.time() - _FILE_LIST_CACHE_TTL_SECS - 60
os.utime(cache_file, (old_mtime, old_mtime))
fresh_list = [FileListEntry(type="file", path="fresh.bin", size=2)]
with patch(
"exo.download.download_utils.fetch_file_list_with_retry",
new_callable=AsyncMock,
return_value=fresh_list,
) as mock_fetch:
result = await fetch_file_list_with_cache(model_id, "main")
assert result == fresh_list
mock_fetch.assert_called_once()
@@ -0,0 +1,355 @@
"""Tests for HuggingFace 429 rate-limit handling in download_utils."""
from collections.abc import AsyncIterator
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch
import aiofiles.os as aios
import pytest
from exo.download.download_utils import (
HuggingFaceRateLimitError,
_download_file, # pyright: ignore[reportPrivateUsage]
_fetch_file_list, # pyright: ignore[reportPrivateUsage]
_parse_retry_after, # pyright: ignore[reportPrivateUsage]
download_file_with_retry,
fetch_file_list_with_retry,
file_meta,
)
from exo.shared.types.common import ModelId
# captured from a real HF 429 on 2026-04-30 (header is lowercased by Cloudfront)
REAL_HF_429_HEADERS_2026_04_30 = {
"ratelimit": '"api";r=0;t=52',
"ratelimit-policy": '"fixed window";"api";q=500;w=300',
}
class TestParseRetryAfter:
def test_parses_documented_format(self) -> None:
assert _parse_retry_after({"RateLimit": '"api";r=0;t=243'}) == 243.0
def test_parses_real_hf_response(self) -> None:
assert _parse_retry_after(REAL_HF_429_HEADERS_2026_04_30) == 52.0
def test_parses_resolvers_bucket(self) -> None:
assert _parse_retry_after({"ratelimit": '"resolvers";r=0;t=120'}) == 120.0
def test_parses_pages_bucket(self) -> None:
assert _parse_retry_after({"ratelimit": '"pages";r=0;t=10'}) == 10.0
def test_returns_none_when_header_missing(self) -> None:
assert _parse_retry_after({}) is None
def test_returns_none_when_only_retry_after_present(self) -> None:
assert _parse_retry_after({"Retry-After": "60"}) is None
def test_returns_none_when_format_unrecognised(self) -> None:
assert _parse_retry_after({"ratelimit": "garbage"}) is None
def test_handles_extra_whitespace(self) -> None:
assert _parse_retry_after({"ratelimit": '"api"; r=0; t=42'}) == 42.0
class TestFetchFileListRetry:
async def test_uses_retry_after_from_error(self) -> None:
sleeps: list[float] = []
async def fake_sleep(seconds: float) -> None:
sleeps.append(seconds)
async def fake_fetch(*args: object, **kwargs: object) -> list[object]:
if not sleeps:
raise HuggingFaceRateLimitError("rate limited", retry_after=2.0)
return []
with (
patch(
"exo.download.download_utils._fetch_file_list", side_effect=fake_fetch
),
patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
):
result = await fetch_file_list_with_retry(ModelId("test/model"))
assert result == []
assert len(sleeps) == 1
assert 2.0 <= sleeps[0] < 3.0 # retry_after + jitter[0,1)
async def test_falls_back_to_exp_backoff_when_no_retry_after(self) -> None:
sleeps: list[float] = []
async def fake_sleep(seconds: float) -> None:
sleeps.append(seconds)
async def fake_fetch(*args: object, **kwargs: object) -> list[object]:
if not sleeps:
raise HuggingFaceRateLimitError("rate limited", retry_after=None)
return []
with (
patch(
"exo.download.download_utils._fetch_file_list", side_effect=fake_fetch
),
patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
):
await fetch_file_list_with_retry(ModelId("test/model"))
assert len(sleeps) == 1
assert 1.0 <= sleeps[0] < 2.0 # 2**0 + jitter[0,1)
async def test_caps_sleep_at_max_window(self) -> None:
sleeps: list[float] = []
async def fake_sleep(seconds: float) -> None:
sleeps.append(seconds)
async def fake_fetch(*args: object, **kwargs: object) -> list[object]:
if not sleeps:
raise HuggingFaceRateLimitError("rate limited", retry_after=10_000.0)
return []
with (
patch(
"exo.download.download_utils._fetch_file_list", side_effect=fake_fetch
),
patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
):
await fetch_file_list_with_retry(ModelId("test/model"))
assert len(sleeps) == 1
assert 300.0 <= sleeps[0] < 301.0 # cap + jitter[0,1)
async def test_retries_up_to_five_times(self) -> None:
sleeps: list[float] = []
async def fake_sleep(seconds: float) -> None:
sleeps.append(seconds)
async def fake_fetch(*args: object, **kwargs: object) -> list[object]:
raise HuggingFaceRateLimitError("rate limited", retry_after=1.0)
with (
patch(
"exo.download.download_utils._fetch_file_list", side_effect=fake_fetch
),
patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
pytest.raises(HuggingFaceRateLimitError),
):
await fetch_file_list_with_retry(ModelId("test/model"))
assert len(sleeps) == 4 # 5 attempts -> 4 sleeps before giving up
class TestDownloadFileRetry:
@pytest.fixture
async def target_dir(self, tmp_path: Path) -> AsyncIterator[Path]:
target = tmp_path / "downloads"
await aios.makedirs(target, exist_ok=True)
yield target
async def test_uses_retry_after_from_error(self, target_dir: Path) -> None:
sleeps: list[float] = []
results: list[Path] = [target_dir / "file.bin"]
async def fake_sleep(seconds: float) -> None:
sleeps.append(seconds)
async def fake_download(*args: object, **kwargs: object) -> Path:
if not sleeps:
raise HuggingFaceRateLimitError("rate limited", retry_after=5.0)
return results[0]
with (
patch(
"exo.download.download_utils._download_file",
side_effect=fake_download,
),
patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
):
result = await download_file_with_retry(
ModelId("test/model"), "main", "file.bin", target_dir
)
assert result == results[0]
assert len(sleeps) == 1
assert 5.0 <= sleeps[0] < 6.0
async def test_caps_sleep_at_max_window(self, target_dir: Path) -> None:
sleeps: list[float] = []
results: list[Path] = [target_dir / "file.bin"]
async def fake_sleep(seconds: float) -> None:
sleeps.append(seconds)
async def fake_download(*args: object, **kwargs: object) -> Path:
if not sleeps:
raise HuggingFaceRateLimitError("rate limited", retry_after=99_999.0)
return results[0]
with (
patch(
"exo.download.download_utils._download_file",
side_effect=fake_download,
),
patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
):
await download_file_with_retry(
ModelId("test/model"), "main", "file.bin", target_dir
)
assert len(sleeps) == 1
assert 300.0 <= sleeps[0] < 301.0
async def test_retries_up_to_five_times(self, target_dir: Path) -> None:
sleeps: list[float] = []
async def fake_sleep(seconds: float) -> None:
sleeps.append(seconds)
with (
patch(
"exo.download.download_utils._download_file",
new_callable=AsyncMock,
side_effect=HuggingFaceRateLimitError("rate limited", retry_after=1.0),
),
patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
pytest.raises(HuggingFaceRateLimitError),
):
await download_file_with_retry(
ModelId("test/model"), "main", "file.bin", target_dir
)
assert len(sleeps) == 4
def _make_mock_session_returning(
response_attrs: dict[str, object], method: str = "get"
) -> MagicMock:
"""Build a MagicMock that mimics ``create_http_session`` returning a
response whose ``status`` / ``headers`` are set from ``response_attrs``.
Mocks the chain ``create_http_session().__aenter__() -> session``, and
``session.<method>().__aenter__() -> response``.
"""
mock_response = MagicMock()
for k, v in response_attrs.items():
setattr(mock_response, k, v)
mock_session = MagicMock()
method_mock = getattr(mock_session, method) # pyright: ignore[reportAny]
method_mock.return_value.__aenter__ = AsyncMock( # pyright: ignore[reportAny]
return_value=mock_response
)
method_mock.return_value.__aexit__ = AsyncMock( # pyright: ignore[reportAny]
return_value=None
)
mock_factory = MagicMock()
mock_factory.return_value.__aenter__ = AsyncMock( # pyright: ignore[reportAny]
return_value=mock_session
)
mock_factory.return_value.__aexit__ = AsyncMock( # pyright: ignore[reportAny]
return_value=None
)
return mock_factory
REAL_HF_429_HEADER_DICT = {"ratelimit": '"api";r=0;t=52'}
class TestRateLimitAtHttpCallSites:
"""Verify each HF call site translates an HTTP 429 into a
``HuggingFaceRateLimitError`` carrying the parsed ``retry_after``.
These tests would catch regressions where (a) the 429 branch is
deleted, (b) ``_parse_retry_after`` stops being called, or
(c) the wrong header object is passed to it.
"""
async def test_fetch_file_list_maps_429_to_rate_limit_error(self) -> None:
mock_factory = _make_mock_session_returning(
{"status": 429, "headers": REAL_HF_429_HEADER_DICT}
)
with (
patch("exo.download.download_utils.create_http_session", mock_factory),
pytest.raises(HuggingFaceRateLimitError) as exc_info,
):
await _fetch_file_list(ModelId("test/model"), "main")
assert exc_info.value.retry_after == 52.0
async def test_file_meta_maps_429_to_rate_limit_error(self) -> None:
mock_factory = _make_mock_session_returning(
{"status": 429, "headers": REAL_HF_429_HEADER_DICT}, method="head"
)
with (
patch("exo.download.download_utils.create_http_session", mock_factory),
pytest.raises(HuggingFaceRateLimitError) as exc_info,
):
await file_meta(ModelId("test/model"), "main", "weights.safetensors")
assert exc_info.value.retry_after == 52.0
async def test_file_meta_maps_429_after_307_redirect(self) -> None:
"""When the initial HEAD 307s and the redirected HEAD then 429s,
the 429 must still surface as ``HuggingFaceRateLimitError``."""
# First HEAD -> 307 with a Location header pointing somewhere new.
first_response = MagicMock()
first_response.status = 307
first_response.headers = {"location": "/redirected/url"}
# Second HEAD (the recursive call) -> 429 with the real-HF header.
second_response = MagicMock()
second_response.status = 429
second_response.headers = REAL_HF_429_HEADER_DICT
responses = iter([first_response, second_response])
mock_session = MagicMock()
mock_session.head.return_value.__aenter__ = AsyncMock( # pyright: ignore[reportAny]
side_effect=lambda: next(responses)
)
mock_session.head.return_value.__aexit__ = AsyncMock( # pyright: ignore[reportAny]
return_value=None
)
mock_factory = MagicMock()
mock_factory.return_value.__aenter__ = AsyncMock( # pyright: ignore[reportAny]
return_value=mock_session
)
mock_factory.return_value.__aexit__ = AsyncMock( # pyright: ignore[reportAny]
return_value=None
)
with (
patch("exo.download.download_utils.create_http_session", mock_factory),
pytest.raises(HuggingFaceRateLimitError) as exc_info,
):
await file_meta(ModelId("test/model"), "main", "weights.safetensors")
assert exc_info.value.retry_after == 52.0
async def test_download_file_maps_429_to_rate_limit_error(
self, tmp_path: Path
) -> None:
target_dir = tmp_path / "downloads"
await aios.makedirs(target_dir, exist_ok=True)
# No local file -> _download_file goes straight to file_meta then GET.
# We need both calls to succeed enough to reach the GET branch:
# - file_meta returns a non-429 (size, etag) so we proceed.
# - the GET then 429s.
with (
patch(
"exo.download.download_utils.file_meta",
new_callable=AsyncMock,
return_value=(100, "abc123"),
),
patch(
"exo.download.download_utils.create_http_session",
_make_mock_session_returning(
{"status": 429, "headers": REAL_HF_429_HEADER_DICT}
),
),
pytest.raises(HuggingFaceRateLimitError) as exc_info,
):
await _download_file(
ModelId("test/model"), "main", "weights.safetensors", target_dir
)
assert exc_info.value.retry_after == 52.0
+16 -14
View File
@@ -5,8 +5,10 @@ import resource
import signal
from dataclasses import dataclass, field
from typing import Self
from uuid import uuid4
import anyio
from exo_net import PySession
from loguru import logger
from pydantic import PositiveInt
@@ -16,7 +18,7 @@ from exo.download.coordinator import DownloadCoordinator
from exo.download.impl_shard_downloader import exo_shard_downloader
from exo.master.main import Master
from exo.routing.event_router import EventRouter
from exo.routing.router import Router, get_node_id_keypair
from exo.routing.router import Router
from exo.shared.constants import EXO_LOG
from exo.shared.election import Election, ElectionResult
from exo.shared.logging import logger_cleanup, logger_setup
@@ -39,31 +41,31 @@ class Node:
api: API | None
node_id: NodeId
session: PySession
offline: bool
_api_port: int
_tg: TaskGroup = field(init=False, default_factory=TaskGroup)
@classmethod
async def create(cls, args: "Args") -> Self:
keypair = get_node_id_keypair()
node_id = NodeId(keypair.to_node_id())
node_id_bytes = uuid4()
node_id = NodeId(str(node_id_bytes))
session_id = SessionId(master_node_id=node_id, election_clock=0)
router = Router.create(
keypair,
router, session = Router.create(
node_id_bytes.bytes,
bootstrap_peers=args.bootstrap_peers,
listen_port=args.libp2p_port,
)
await router.register_topic(topics.GLOBAL_EVENTS)
await router.register_topic(topics.LOCAL_EVENTS)
await router.register_topic(topics.COMMANDS)
await router.register_topic(topics.ELECTION_MESSAGES)
await router.register_topic(topics.CONNECTION_MESSAGES)
await router.register_topic(topics.DOWNLOAD_COMMANDS)
event_router = EventRouter(
session_id,
command_sender=router.sender(topics.COMMANDS),
external_outbound=router.sender(topics.LOCAL_EVENTS),
external_inbound=router.receiver(topics.GLOBAL_EVENTS),
command_sender=session.net_sender("orchestrator"),
)
logger.info(f"Starting node {node_id}")
@@ -85,9 +87,9 @@ class Node:
node_id,
port=args.api_port,
event_receiver=event_router.receiver(),
command_sender=router.sender(topics.COMMANDS),
download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
election_receiver=router.receiver(topics.ELECTION_MESSAGES),
session=session,
)
else:
api = None
@@ -97,9 +99,9 @@ class Node:
node_id,
event_receiver=event_router.receiver(),
event_sender=event_router.sender(),
command_sender=router.sender(topics.COMMANDS),
download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
api_port=args.api_port,
session=session,
)
else:
worker = None
@@ -111,8 +113,8 @@ class Node:
event_sender=event_router.sender(),
global_event_sender=router.sender(topics.GLOBAL_EVENTS),
local_event_receiver=router.receiver(topics.LOCAL_EVENTS),
command_receiver=router.receiver(topics.COMMANDS),
download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
command_receiver=session.net_receiver("orchestrator")
)
er_send, er_recv = channel[ElectionResult]()
@@ -125,7 +127,6 @@ class Node:
election_message_sender=router.sender(topics.ELECTION_MESSAGES),
election_message_receiver=router.receiver(topics.ELECTION_MESSAGES),
connection_message_receiver=router.receiver(topics.CONNECTION_MESSAGES),
command_receiver=router.receiver(topics.COMMANDS),
election_result_sender=er_send,
)
@@ -139,6 +140,7 @@ class Node:
master,
api,
node_id,
session,
args.offline,
args.api_port,
)
@@ -188,7 +190,7 @@ class Node:
self.event_router.shutdown()
self.event_router = EventRouter(
result.session_id,
self.router.sender(topics.COMMANDS),
self.session.net_sender("orchestrator"),
self.router.receiver(topics.GLOBAL_EVENTS),
self.router.sender(topics.LOCAL_EVENTS),
)
@@ -209,10 +211,10 @@ class Node:
event_sender=self.event_router.sender(),
global_event_sender=self.router.sender(topics.GLOBAL_EVENTS),
local_event_receiver=self.router.receiver(topics.LOCAL_EVENTS),
command_receiver=self.router.receiver(topics.COMMANDS),
download_command_sender=self.router.sender(
topics.DOWNLOAD_COMMANDS
),
command_receiver=self.session.net_receiver("orchestrator"),
)
self._tg.start_soon(self.master.run)
elif (
@@ -248,11 +250,11 @@ class Node:
self.node_id,
event_receiver=self.event_router.receiver(),
event_sender=self.event_router.sender(),
command_sender=self.router.sender(topics.COMMANDS),
download_command_sender=self.router.sender(
topics.DOWNLOAD_COMMANDS
),
api_port=self._api_port,
session=self.session,
)
self._tg.start_soon(self.worker.run)
if self.api:
+283 -281
View File
@@ -1,7 +1,9 @@
from datetime import datetime, timedelta, timezone
import anyio
from exo_net import NetReceiver
from loguru import logger
from pydantic import TypeAdapter
from exo.master.placement import (
add_instance_to_placements,
@@ -15,11 +17,11 @@ 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,
Command,
CreateInstance,
DeleteCustomModelCard,
DeleteInstance,
DeleteInstanceLink,
ForwarderCommand,
ForwarderDownloadCommand,
ImageEdits,
ImageGeneration,
@@ -121,7 +123,7 @@ class Master:
node_id: NodeId,
session_id: SessionId,
*,
command_receiver: Receiver[ForwarderCommand],
command_receiver: NetReceiver, # todo: not this type
event_sender: Sender[Event],
local_event_receiver: Receiver[LocalForwarderEvent],
global_event_sender: Sender[GlobalForwarderEvent],
@@ -155,302 +157,302 @@ class Master:
self._event_log.close()
self.global_event_sender.close()
self.local_event_receiver.close()
self.command_receiver.close()
async def shutdown(self):
logger.info("Stopping Master")
self._tg.cancel_tasks()
async def _command_processor(self) -> None:
with self.command_receiver as commands:
async for forwarder_command in commands:
try:
logger.info(f"Executing command: {forwarder_command.command}")
while True:
data = await self.command_receiver.recv()
if not data:
break
try:
command = TypeAdapter[Command](Command).validate_json(data)
logger.info(f"Executing command: {command}")
generated_events: list[Event] = []
command = forwarder_command.command
instance_task_counts: dict[InstanceId, int] = {}
match command:
case TestCommand():
pass
case TextGeneration():
prefill_only: set[InstanceId] = set()
for link in self.state.instance_links.values():
prefill_only.update(link.prefill_instances)
for link in self.state.instance_links.values():
prefill_only.difference_update(link.decode_instances)
generated_events: list[Event] = []
instance_task_counts: dict[InstanceId, int] = {}
match command:
case TestCommand():
pass
case TextGeneration():
prefill_only: set[InstanceId] = set()
for link in self.state.instance_links.values():
prefill_only.update(link.prefill_instances)
for link in self.state.instance_links.values():
prefill_only.difference_update(link.decode_instances)
for instance in self.state.instances.values():
if (
instance.shard_assignments.model_id
== command.task_params.model
and instance.instance_id not in prefill_only
):
in_flight = {TaskStatus.Pending, TaskStatus.Running}
task_count = sum(
1
for task in self.state.tasks.values()
if task.instance_id == instance.instance_id
and task.task_status in in_flight
)
instance_task_counts[instance.instance_id] = (
task_count
)
if not instance_task_counts:
raise ValueError(
f"No instance found for model {command.task_params.model}"
)
available_instance_ids = sorted(
instance_task_counts.keys(),
key=lambda instance_id: instance_task_counts[
instance_id
],
)
decode_instance_id = available_instance_ids[0]
task_id = TaskId()
params = command.task_params.model_copy(
update={
"prefill_endpoint": _prefill_endpoint_for(
self.state, decode_instance_id
),
}
)
generated_events.append(
TaskCreated(
task_id=task_id,
task=TextGenerationTask(
task_id=task_id,
command_id=command.command_id,
instance_id=decode_instance_id,
task_status=TaskStatus.Pending,
task_params=params,
),
)
)
self.command_task_mapping[command.command_id] = task_id
case ImageGeneration():
for instance in self.state.instances.values():
if (
instance.shard_assignments.model_id
== command.task_params.model
):
in_flight = {TaskStatus.Pending, TaskStatus.Running}
task_count = sum(
1
for task in self.state.tasks.values()
if task.instance_id == instance.instance_id
and task.task_status in in_flight
)
instance_task_counts[instance.instance_id] = (
task_count
)
if not instance_task_counts:
raise ValueError(
f"No instance found for model {command.task_params.model}"
)
available_instance_ids = sorted(
instance_task_counts.keys(),
key=lambda instance_id: instance_task_counts[
instance_id
],
)
task_id = TaskId()
selected_instance_id = available_instance_ids[0]
generated_events.append(
TaskCreated(
task_id=task_id,
task=ImageGenerationTask(
task_id=task_id,
command_id=command.command_id,
instance_id=selected_instance_id,
task_status=TaskStatus.Pending,
task_params=command.task_params,
),
)
)
self.command_task_mapping[command.command_id] = task_id
if EXO_TRACING_ENABLED:
selected_instance = self.state.instances.get(
selected_instance_id
)
if selected_instance:
ranks = set(
shard.device_rank
for shard in selected_instance.shard_assignments.runner_to_shard.values()
)
self._expected_ranks[task_id] = ranks
case ImageEdits():
for instance in self.state.instances.values():
if (
instance.shard_assignments.model_id
== command.task_params.model
):
in_flight = {TaskStatus.Pending, TaskStatus.Running}
task_count = sum(
1
for task in self.state.tasks.values()
if task.instance_id == instance.instance_id
and task.task_status in in_flight
)
instance_task_counts[instance.instance_id] = (
task_count
)
if not instance_task_counts:
raise ValueError(
f"No instance found for model {command.task_params.model}"
)
available_instance_ids = sorted(
instance_task_counts.keys(),
key=lambda instance_id: instance_task_counts[
instance_id
],
)
task_id = TaskId()
selected_instance_id = available_instance_ids[0]
generated_events.append(
TaskCreated(
task_id=task_id,
task=ImageEditsTask(
task_id=task_id,
command_id=command.command_id,
instance_id=selected_instance_id,
task_status=TaskStatus.Pending,
task_params=command.task_params,
),
)
)
self.command_task_mapping[command.command_id] = task_id
if EXO_TRACING_ENABLED:
selected_instance = self.state.instances.get(
selected_instance_id
)
if selected_instance:
ranks = set(
shard.device_rank
for shard in selected_instance.shard_assignments.runner_to_shard.values()
)
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
for instance in self.state.instances.values():
if (
instance.shard_assignments.model_id
== command.task_params.model
and instance.instance_id not in prefill_only
):
await self.download_command_sender.send(
ForwarderDownloadCommand(
origin=self._system_id, command=cmd
)
in_flight = {TaskStatus.Pending, TaskStatus.Running}
task_count = sum(
1
for task in self.state.tasks.values()
if task.instance_id == instance.instance_id
and task.task_status in in_flight
)
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,
download_status=self.state.downloads,
)
transition_events = get_transition_events(
self.state.instances, placement, self.state.tasks
)
generated_events.extend(transition_events)
case CreateInstance():
placement = add_instance_to_placements(
command,
self.state.topology,
self.state.instances,
)
transition_events = get_transition_events(
self.state.instances, placement, self.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(
command.cancelled_command_id
)
) is not None:
generated_events.append(
TaskStatusUpdated(
task_status=TaskStatus.Cancelled,
task_id=task_id,
)
)
else:
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"
instance_task_counts[instance.instance_id] = (
task_count
)
case AddCustomModelCard():
generated_events.append(
CustomModelCardAdded(model_card=command.model_card)
if not instance_task_counts:
raise ValueError(
f"No instance found for model {command.task_params.model}"
)
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)
available_instance_ids = sorted(
instance_task_counts.keys(),
key=lambda instance_id: instance_task_counts[
instance_id
],
)
decode_instance_id = available_instance_ids[0]
task_id = TaskId()
params = command.task_params.model_copy(
update={
"prefill_endpoint": _prefill_endpoint_for(
self.state, decode_instance_id
),
decode_instances=list(
dict.fromkeys(command.decode_instances)
}
)
generated_events.append(
TaskCreated(
task_id=task_id,
task=TextGenerationTask(
task_id=task_id,
command_id=command.command_id,
instance_id=decode_instance_id,
task_status=TaskStatus.Pending,
task_params=params,
),
)
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
end = min(command.since_idx + 1000, len(self._event_log))
for i, event in enumerate(
self._event_log.read_range(command.since_idx, end),
start=command.since_idx,
)
self.command_task_mapping[command.command_id] = task_id
case ImageGeneration():
for instance in self.state.instances.values():
if (
instance.shard_assignments.model_id
== command.task_params.model
):
await self._send_event(IndexedEvent(idx=i, event=event))
for event in generated_events:
await self.event_sender.send(event)
except ValueError as e:
logger.opt(exception=e).warning("Error in command processor")
in_flight = {TaskStatus.Pending, TaskStatus.Running}
task_count = sum(
1
for task in self.state.tasks.values()
if task.instance_id == instance.instance_id
and task.task_status in in_flight
)
instance_task_counts[instance.instance_id] = (
task_count
)
if not instance_task_counts:
raise ValueError(
f"No instance found for model {command.task_params.model}"
)
available_instance_ids = sorted(
instance_task_counts.keys(),
key=lambda instance_id: instance_task_counts[
instance_id
],
)
task_id = TaskId()
selected_instance_id = available_instance_ids[0]
generated_events.append(
TaskCreated(
task_id=task_id,
task=ImageGenerationTask(
task_id=task_id,
command_id=command.command_id,
instance_id=selected_instance_id,
task_status=TaskStatus.Pending,
task_params=command.task_params,
),
)
)
self.command_task_mapping[command.command_id] = task_id
if EXO_TRACING_ENABLED:
selected_instance = self.state.instances.get(
selected_instance_id
)
if selected_instance:
ranks = set(
shard.device_rank
for shard in selected_instance.shard_assignments.runner_to_shard.values()
)
self._expected_ranks[task_id] = ranks
case ImageEdits():
for instance in self.state.instances.values():
if (
instance.shard_assignments.model_id
== command.task_params.model
):
in_flight = {TaskStatus.Pending, TaskStatus.Running}
task_count = sum(
1
for task in self.state.tasks.values()
if task.instance_id == instance.instance_id
and task.task_status in in_flight
)
instance_task_counts[instance.instance_id] = (
task_count
)
if not instance_task_counts:
raise ValueError(
f"No instance found for model {command.task_params.model}"
)
available_instance_ids = sorted(
instance_task_counts.keys(),
key=lambda instance_id: instance_task_counts[
instance_id
],
)
task_id = TaskId()
selected_instance_id = available_instance_ids[0]
generated_events.append(
TaskCreated(
task_id=task_id,
task=ImageEditsTask(
task_id=task_id,
command_id=command.command_id,
instance_id=selected_instance_id,
task_status=TaskStatus.Pending,
task_params=command.task_params,
),
)
)
self.command_task_mapping[command.command_id] = task_id
if EXO_TRACING_ENABLED:
selected_instance = self.state.instances.get(
selected_instance_id
)
if selected_instance:
ranks = set(
shard.device_rank
for shard in selected_instance.shard_assignments.runner_to_shard.values()
)
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,
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():
placement = add_instance_to_placements(
command,
self.state.topology,
self.state.instances,
)
transition_events = get_transition_events(
self.state.instances, placement, self.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(
command.cancelled_command_id
)
) is not None:
generated_events.append(
TaskStatusUpdated(
task_status=TaskStatus.Cancelled,
task_id=task_id,
)
)
else:
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():
end = len(self._event_log)
for i, event in enumerate(
self._event_log.read_range(command.since_idx, end),
start=command.since_idx,
):
await self._send_event(IndexedEvent(idx=i, event=event))
for event in generated_events:
await self.event_sender.send(event)
except ValueError as e:
logger.opt(exception=e).warning("Error in command processor")
# These plan loops are the cracks showing in our event sourcing architecture - more things could be commands
async def _plan(self) -> None:
+13 -2
View File
@@ -28,7 +28,7 @@ from exo.shared.types.events import (
TaskStatusUpdated,
)
from exo.shared.types.memory import Memory
from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo, NodeRdmaCtlStatus
from exo.shared.types.tasks import Task, TaskId, TaskStatus
from exo.shared.types.worker.downloads import (
DownloadCompleted,
@@ -105,6 +105,7 @@ def place_instance(
node_network: Mapping[NodeId, NodeNetworkInfo],
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]:
cycles = topology.get_cycles()
candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles))
@@ -166,8 +167,18 @@ def place_instance(
smallest_cycles = get_smallest_cycles(cycles_with_sufficient_memory)
rdma_ctl_status = node_rdma_ctl or {}
def _all_rdma_ctl_enabled(cycle: Cycle) -> bool:
return all(
((status := rdma_ctl_status.get(node_id)) is not None and status.enabled)
for node_id in cycle
)
smallest_rdma_cycles = [
cycle for cycle in smallest_cycles if topology.is_rdma_cycle(cycle)
cycle
for cycle in smallest_cycles
if topology.is_rdma_cycle(cycle) and _all_rdma_ctl_enabled(cycle)
]
if command.instance_meta == InstanceMeta.MlxJaccl:
+1 -3
View File
@@ -6,7 +6,6 @@ import pytest
from loguru import logger
from exo.master.main import Master
from exo.routing.router import get_node_id_keypair
from exo.shared.models.model_cards import ModelCard, ModelTask
from exo.shared.types.commands import (
CommandId,
@@ -47,8 +46,7 @@ from exo.utils.channels import channel
@pytest.mark.asyncio
async def test_master():
keypair = get_node_id_keypair()
node_id = NodeId(keypair.to_node_id())
node_id = NodeId("master test")
session_id = SessionId(master_node_id=node_id, election_clock=0)
ge_sender, global_event_receiver = channel[GlobalForwarderEvent]()
+144 -2
View File
@@ -21,7 +21,11 @@ from exo.shared.types.events import (
)
from exo.shared.types.memory import Memory
from exo.shared.types.multiaddr import Multiaddr
from exo.shared.types.profiling import NetworkInterfaceInfo, NodeNetworkInfo
from exo.shared.types.profiling import (
NetworkInterfaceInfo,
NodeNetworkInfo,
NodeRdmaCtlStatus,
)
from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
from exo.shared.types.text_generation import (
InputMessage,
@@ -439,8 +443,21 @@ def test_tensor_rdma_backend_connectivity_matrix(
min_nodes=1,
)
node_rdma_ctl = {
node_a: NodeRdmaCtlStatus(enabled=True),
node_b: NodeRdmaCtlStatus(enabled=True),
node_c: NodeRdmaCtlStatus(enabled=True),
}
# act
placements = place_instance(cic, topology, {}, node_memory, node_network)
placements = place_instance(
cic,
topology,
{},
node_memory,
node_network,
node_rdma_ctl=node_rdma_ctl,
)
# assert
assert len(placements) == 1
@@ -482,6 +499,131 @@ def test_tensor_rdma_backend_connectivity_matrix(
assert len(ip_part.split(".")) == 4
def _build_three_node_rdma_topology() -> tuple[
Topology, NodeId, NodeId, NodeId, dict[NodeId, NodeNetworkInfo]
]:
topology = Topology()
node_a = NodeId()
node_b = NodeId()
node_c = NodeId()
ethernet_interface = NetworkInterfaceInfo(name="en0", ip_address="10.0.0.1")
ethernet_conn = SocketConnection(
sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000")
)
node_network = {
node_a: NodeNetworkInfo(interfaces=[ethernet_interface]),
node_b: NodeNetworkInfo(interfaces=[ethernet_interface]),
node_c: NodeNetworkInfo(interfaces=[ethernet_interface]),
}
for n in (node_a, node_b, node_c):
topology.add_node(n)
rdma_pairs = [
(node_a, node_b, 3),
(node_b, node_a, 3),
(node_b, node_c, 4),
(node_c, node_b, 4),
(node_a, node_c, 5),
(node_c, node_a, 5),
]
for src, sink, iface in rdma_pairs:
topology.add_connection(
Connection(source=src, sink=sink, edge=create_rdma_connection(iface))
)
socket_pairs = [
(node_a, node_b),
(node_b, node_c),
(node_c, node_a),
(node_a, node_c),
(node_b, node_a),
(node_c, node_b),
]
for src, sink in socket_pairs:
topology.add_connection(Connection(source=src, sink=sink, edge=ethernet_conn))
return topology, node_a, node_b, node_c, node_network
def test_place_mlx_jaccl_rejects_when_a_node_has_rdma_ctl_disabled(
model_card: ModelCard,
):
# arrange
model_card = model_card.model_copy(
update={"n_layers": 12, "storage_size": Memory.from_bytes(1500)}
)
topology, node_a, node_b, node_c, node_network = _build_three_node_rdma_topology()
node_memory = {
node_a: create_node_memory(500),
node_b: create_node_memory(500),
node_c: create_node_memory(500),
}
node_rdma_ctl = {
node_a: NodeRdmaCtlStatus(enabled=True),
node_b: NodeRdmaCtlStatus(enabled=True),
node_c: NodeRdmaCtlStatus(enabled=False),
}
cic = PlaceInstance(
sharding=Sharding.Tensor,
instance_meta=InstanceMeta.MlxJaccl,
command_id=CommandId(),
model_card=model_card,
min_nodes=3,
)
# act / assert
with pytest.raises(
ValueError, match="Requested RDMA \\(MlxJaccl\\) but no RDMA-connected cycles"
):
place_instance(
cic,
topology,
{},
node_memory,
node_network,
node_rdma_ctl=node_rdma_ctl,
)
def test_place_mlx_jaccl_rejects_when_node_rdma_ctl_missing(model_card: ModelCard):
"""A node with no observed rdma_ctl status must not participate in RDMA placement."""
# arrange
model_card = model_card.model_copy(
update={"n_layers": 12, "storage_size": Memory.from_bytes(1500)}
)
topology, node_a, node_b, node_c, node_network = _build_three_node_rdma_topology()
node_memory = {
node_a: create_node_memory(500),
node_b: create_node_memory(500),
node_c: create_node_memory(500),
}
# node_c has no rdma_ctl entry at all
node_rdma_ctl = {
node_a: NodeRdmaCtlStatus(enabled=True),
node_b: NodeRdmaCtlStatus(enabled=True),
}
cic = PlaceInstance(
sharding=Sharding.Tensor,
instance_meta=InstanceMeta.MlxJaccl,
command_id=CommandId(),
model_card=model_card,
min_nodes=3,
)
# act / assert
with pytest.raises(ValueError):
place_instance(
cic,
topology,
{},
node_memory,
node_network,
node_rdma_ctl=node_rdma_ctl,
)
def _make_task(
instance_id: InstanceId,
status: TaskStatus = TaskStatus.Running,
+2 -4
View File
@@ -1,15 +1,13 @@
from exo_pyo3_bindings import PyFromSwarm
from exo_net import PyFromSwarm
from exo.shared.types.common import NodeId
from exo.utils.pydantic_ext import FrozenModel
"""Serialisable types for Connection Updates/Messages"""
class ConnectionMessage(FrozenModel):
node_id: NodeId
connected: bool
@classmethod
def from_update(cls, update: PyFromSwarm.Connection) -> "ConnectionMessage":
return cls(node_id=NodeId(update.peer_id), connected=update.connected)
return cls(connected=update.connected)
+6 -6
View File
@@ -4,9 +4,10 @@ from random import random
import anyio
from anyio import BrokenResourceError, ClosedResourceError
from anyio.abc import CancelScope
from exo_net import NetSender
from loguru import logger
from exo.shared.types.commands import ForwarderCommand, RequestEventLog
from exo.shared.types.commands import RequestEventLog
from exo.shared.types.common import SessionId, SystemId
from exo.shared.types.events import (
Event,
@@ -23,7 +24,7 @@ from exo.utils.task_group import TaskGroup
@dataclass
class EventRouter:
session_id: SessionId
command_sender: Sender[ForwarderCommand]
command_sender: NetSender
external_inbound: Receiver[GlobalForwarderEvent]
external_outbound: Sender[LocalForwarderEvent]
_system_id: SystemId = field(init=False, default_factory=SystemId)
@@ -152,10 +153,9 @@ class EventRouter:
f"Nack attempt {self._nack_attempts}: Requesting Event Log from {since_idx}"
)
await self.command_sender.send(
ForwarderCommand(
origin=self._system_id,
command=RequestEventLog(since_idx=since_idx),
)
RequestEventLog(since_idx=since_idx)
.model_dump_json()
.encode("utf-8")
)
finally:
if self._nack_cancel_scope is scope:
+12 -66
View File
@@ -2,8 +2,6 @@ from collections.abc import Sequence
from copy import copy
from itertools import count
from math import inf
from os import PathLike
from pathlib import Path
from typing import cast
from anyio import (
@@ -12,18 +10,9 @@ from anyio import (
move_on_after,
sleep_forever,
)
from exo_pyo3_bindings import (
AllQueuesFullError,
Keypair,
MessageTooLargeError,
NetworkingHandle,
NoPeersSubscribedToTopicError,
PyFromSwarm,
)
from filelock import FileLock
from exo_net import NetworkingHandle, PyFromSwarm, PySession
from loguru import logger
from exo.shared.constants import EXO_NODE_ID_KEYPAIR
from exo.utils.channels import Receiver, Sender, channel
from exo.utils.pydantic_ext import FrozenModel
from exo.utils.task_group import TaskGroup
@@ -105,13 +94,14 @@ class Router:
@classmethod
def create(
cls,
identity: Keypair,
identity: bytes,
bootstrap_peers: Sequence[str] = (),
listen_port: int = 0,
) -> "Router":
return cls(
handle=NetworkingHandle(identity, list(bootstrap_peers), listen_port)
) -> "tuple[Router, PySession]":
handle, session = NetworkingHandle.new(
identity, list(bootstrap_peers), listen_port
)
return cls(handle=handle), session
def __init__(self, handle: NetworkingHandle):
self.topic_routers: dict[str, TopicRouter[FrozenModel]] = {}
@@ -191,10 +181,8 @@ class Router:
from_swarm = await self._net.recv()
logger.debug(from_swarm)
match from_swarm:
case PyFromSwarm.Message(origin, topic, data):
logger.trace(
f"Received message on {topic} from {origin} with payload {data}"
)
case PyFromSwarm.Message(topic, data):
logger.trace(f"Received message on {topic} with payload {data}")
if topic not in self.topic_routers:
logger.warning(
f"Received message on unknown or inactive topic {topic}"
@@ -225,51 +213,9 @@ class Router:
async def _networking_publish(self):
with self.networking_receiver as networked_items:
async for topic, data in networked_items:
try:
logger.trace(f"Sending message on {topic} with payload {data}")
if len(data) > 1024 * 1024:
logger.warning(
"Sending overlarge payload, network performance may be temporarily degraded"
)
await self._net.gossipsub_publish(topic, data)
except NoPeersSubscribedToTopicError:
pass
except AllQueuesFullError:
logger.warning(f"All peer queues full, dropping message on {topic}")
except MessageTooLargeError:
logger.trace(f"Sending message on {topic} with payload {data}")
if len(data) > 1024 * 1024:
logger.warning(
f"Message too large for gossipsub on {topic} ({len(data)} bytes), dropping"
"Sending overlarge payload, network performance may be temporarily degraded"
)
def get_node_id_keypair(
path: str | bytes | PathLike[str] | PathLike[bytes] = EXO_NODE_ID_KEYPAIR,
) -> Keypair:
"""
Obtains the :class:`Keypair` associated with this node-ID.
Obtain the :class:`PeerId` by from it.
"""
# TODO(evan): bring back node id persistence once we figure out how to deal with duplicates
return Keypair.generate()
def lock_path(path: str | bytes | PathLike[str] | PathLike[bytes]) -> Path:
return Path(str(path) + ".lock")
# operate with cross-process lock to avoid race conditions
with FileLock(lock_path(path)):
with open(path, "a+b") as f: # opens in append-mode => starts at EOF
# if non-zero EOF, then file exists => use to get node-ID
if f.tell() != 0:
f.seek(0) # go to start & read protobuf-encoded bytes
protobuf_encoded = f.read()
try: # if decoded successfully, save & return
return Keypair.from_bytes(protobuf_encoded)
except ValueError as e: # on runtime error, assume corrupt file
logger.warning(f"Encountered error when trying to get keypair: {e}")
# if no valid credentials, create new ones and persist
with open(path, "w+b") as f:
keypair = Keypair.generate()
f.write(keypair.to_bytes())
return keypair
await self._net.gossipsub_publish(topic, data)
+50 -3
View File
@@ -4,7 +4,8 @@ from datetime import datetime
from loguru import logger
from exo.shared.types.common import NodeId
from exo.shared.models.model_cards import ModelCard
from exo.shared.types.common import ModelId, NodeId
from exo.shared.types.events import (
ChunkGenerated,
CustomModelCardAdded,
@@ -65,6 +66,18 @@ 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:
@@ -75,10 +88,12 @@ def event_apply(event: Event, state: State) -> State:
| InputChunkReceived()
| TracesCollected()
| TracesMerged()
| CustomModelCardAdded()
| CustomModelCardDeleted()
): # 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():
@@ -397,6 +412,9 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
for nid in state.node_thunderbolt
for tb_ident in state.node_thunderbolt[nid].interfaces
}
source_is_rdma_enabled = _is_rdma_ctl_enabled(
event.node_id, state.node_rdma_ctl
)
as_rdma_conns = [
Connection(
source=event.node_id,
@@ -409,6 +427,10 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
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():
@@ -432,6 +454,12 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
**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)
return state.model_copy(update=update)
@@ -447,3 +475,22 @@ def apply_topology_edge_deleted(event: TopologyEdgeDeleted, state: State) -> Sta
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,
}
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
}
return state.model_copy(update={"custom_model_cards": new_cards})
-14
View File
@@ -9,7 +9,6 @@ from anyio import (
from loguru import logger
from exo.routing.connection_message import ConnectionMessage
from exo.shared.types.commands import ForwarderCommand
from exo.shared.types.common import NodeId, SessionId
from exo.utils.channels import Receiver, Sender
from exo.utils.pydantic_ext import FrozenModel
@@ -22,7 +21,6 @@ class ElectionMessage(FrozenModel):
clock: int
seniority: int
proposed_session: SessionId
commands_seen: int
# Could eventually include a list of neighbour nodes for centrality
def __lt__(self, other: Self) -> bool:
@@ -30,8 +28,6 @@ class ElectionMessage(FrozenModel):
return self.clock < other.clock
if self.seniority != other.seniority:
return self.seniority < other.seniority
elif self.commands_seen != other.commands_seen:
return self.commands_seen < other.commands_seen
else:
return (
self.proposed_session.master_node_id
@@ -54,7 +50,6 @@ class Election:
election_message_sender: Sender[ElectionMessage],
election_result_sender: Sender[ElectionResult],
connection_message_receiver: Receiver[ConnectionMessage],
command_receiver: Receiver[ForwarderCommand],
is_candidate: bool = True,
seniority: int = 0,
):
@@ -64,7 +59,6 @@ class Election:
self.seniority = seniority if is_candidate else -1
self.clock = 0
self.node_id = node_id
self.commands_seen = 0
# Every node spawns as master
self.current_session: SessionId = SessionId(
master_node_id=node_id, election_clock=0
@@ -75,7 +69,6 @@ class Election:
self._em_receiver = election_message_receiver
self._er_sender = election_result_sender
self._cm_receiver = connection_message_receiver
self._co_receiver = command_receiver
# Campaign state
self._candidates: list[ElectionMessage] = []
@@ -89,7 +82,6 @@ class Election:
async with self._tg as tg:
tg.start_soon(self._election_receiver)
tg.start_soon(self._connection_receiver)
tg.start_soon(self._command_counter)
# And start an election immediately, that instantly resolves
candidates: list[ElectionMessage] = []
@@ -179,11 +171,6 @@ class Election:
logger.debug("Campaign started")
logger.debug("Connection message added")
async def _command_counter(self) -> None:
with self._co_receiver as commands:
async for _command in commands:
self.commands_seen += 1
async def _campaign(
self, candidates: list[ElectionMessage], campaign_timeout: float
) -> None:
@@ -261,5 +248,4 @@ class Election:
),
clock=c,
seniority=self.seniority,
commands_seen=self.commands_seen,
)
+2 -1
View File
@@ -46,7 +46,8 @@ class _InterceptHandler(logging.Handler):
def logger_setup(log_file: Path | None, verbosity: int = 0):
"""Set up logging for this process - formatting, file handles, verbosity and output"""
logging.getLogger("exo_pyo3_bindings").setLevel(logging.WARNING)
logging.getLogger("exo_net").setLevel(logging.INFO)
logging.getLogger("networking").setLevel(logging.INFO)
logging.getLogger("httpx").setLevel(logging.WARNING)
logging.getLogger("httpcore").setLevel(logging.WARNING)
+54 -52
View File
@@ -39,7 +39,57 @@ _BUILTIN_CARD_DIRS = [
Path(RESOURCES_DIR) / "image_model_cards",
]
_card_cache: dict[ModelId, "ModelCard"] = {}
class _CardCache:
def __init__(self):
self.cc: dict[ModelId, "ModelCard"] = {}
def get(self, model_id: ModelId) -> "ModelCard | None":
return self.cc.get(model_id)
async def save(self, card: "ModelCard"):
self.cc[card.model_id] = card
try:
await card.save_to_custom_dir()
except OSError as e:
logger.warning(f"failed to save custom model card ({e.strerror})")
async def pop(self, model_id: ModelId) -> "ModelCard | None":
"""Delete a user-added custom model card. Returns True if deleted."""
card_path = _custom_cards_dir / (ModelId(model_id).normalize() + ".toml")
try:
if await card_path.exists():
await card_path.unlink()
return self.cc.pop(model_id, None)
except OSError as e:
logger.warning(f"failed to delete custom model card ({e.strerror})")
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)]
async def _load_cards_from_dir(self, directory: Path, *, is_custom: bool) -> None:
"""Load all TOML model cards from a directory into the cache."""
async for toml_file in directory.rglob("*.toml"):
try:
card = await ModelCard.load_from_path(toml_file)
if is_custom:
card = card.model_copy(update={"is_custom": True})
if self.get(card.model_id) is None:
self.cc[card.model_id] = card
except (ValidationError, TOMLKitError):
pass
async def refresh(self) -> None:
for path in _BUILTIN_CARD_DIRS:
await self._load_cards_from_dir(path, is_custom=False)
await self._load_cards_from_dir(_custom_cards_dir, is_custom=True)
card_cache = _CardCache()
def detect_vision_from_config(model_id: ModelId) -> "VisionCardConfig | None":
@@ -59,42 +109,10 @@ def detect_vision_from_config(model_id: ModelId) -> "VisionCardConfig | None":
return None
async def _load_cards_from_dir(directory: Path, *, is_custom: bool) -> None:
"""Load all TOML model cards from a directory into the cache."""
async for toml_file in directory.rglob("*.toml"):
try:
card = await ModelCard.load_from_path(toml_file)
if is_custom:
card = card.model_copy(update={"is_custom": True})
if card.model_id not in _card_cache:
_card_cache[card.model_id] = card
except (ValidationError, TOMLKitError):
pass
async def _refresh_card_cache() -> None:
for path in _BUILTIN_CARD_DIRS:
await _load_cards_from_dir(path, is_custom=False)
await _load_cards_from_dir(_custom_cards_dir, is_custom=True)
def _is_image_card(card: "ModelCard") -> bool:
return any(t in (ModelTask.TextToImage, ModelTask.ImageToImage) for t in card.tasks)
def get_card(model_id: ModelId) -> "ModelCard | None":
"""Look up a single model card from the cache by ID."""
return _card_cache.get(model_id)
async def get_model_cards() -> list["ModelCard"]:
if len(_card_cache) == 0:
await _refresh_card_cache()
if EXO_ENABLE_IMAGE_MODELS:
return list(_card_cache.values())
return [c for c in _card_cache.values() if not _is_image_card(c)]
class ModelTask(str, Enum):
TextGeneration = "TextGeneration"
TextToImage = "TextToImage"
@@ -196,14 +214,13 @@ class ModelCard(FrozenModel):
# Is it okay that model card.load defaults to network access if the card doesn't exist? do we want to be more explicit here?
@staticmethod
async def load(model_id: ModelId) -> "ModelCard":
if model_id not in _card_cache:
await _refresh_card_cache()
if (mc := _card_cache.get(model_id)) is not None:
if card_cache.get(model_id) is None:
await card_cache.refresh()
if (mc := card_cache.get(model_id)) is not None:
return mc
mc = await ModelCard.fetch_from_hf(model_id)
await mc.save_to_custom_dir()
_card_cache[model_id] = mc
return mc
@staticmethod
@@ -233,21 +250,6 @@ class ModelCard(FrozenModel):
)
def add_to_card_cache(card: "ModelCard") -> None:
"""Add or update a model card in the in-memory cache."""
_card_cache[card.model_id] = card
async def delete_custom_card(model_id: ModelId) -> bool:
"""Delete a user-added custom model card. Returns True if deleted."""
card_path = _custom_cards_dir / (ModelId(model_id).normalize() + ".toml")
if await card_path.exists():
await card_path.unlink()
_card_cache.pop(model_id, None)
return True
return False
class ConfigData(BaseModel):
model_config = {"extra": "ignore"} # Allow unknown fields
@@ -0,0 +1,44 @@
from exo.shared.apply import apply
from exo.shared.models.model_cards import ModelCard, ModelTask
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],
)
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 == {}
@@ -0,0 +1,231 @@
from datetime import datetime, timezone
from exo.shared.apply import apply_node_gathered_info
from exo.shared.topology import Topology
from exo.shared.types.common import NodeId
from exo.shared.types.events import NodeGatheredInfo
from exo.shared.types.profiling import (
NodeRdmaCtlStatus,
NodeThunderboltInfo,
)
from exo.shared.types.state import State
from exo.shared.types.thunderbolt import ThunderboltConnection, ThunderboltIdentifier
from exo.shared.types.topology import RDMAConnection
from exo.utils.info_gatherer.info_gatherer import (
MacThunderboltConnections,
RdmaCtlStatus,
)
def _now() -> str:
return datetime.now(timezone.utc).isoformat()
def _make_state_with_thunderbolt_idents(
*node_ids_and_uuids: tuple[NodeId, str, str],
rdma_ctl: dict[NodeId, NodeRdmaCtlStatus] | None = None,
) -> State:
"""Build a State with Thunderbolt identifiers per node so the apply MacThunderboltConnections
case can resolve uuid -> (node, iface)."""
node_thunderbolt = {
nid: NodeThunderboltInfo(
interfaces=[ThunderboltIdentifier(rdma_interface=iface, domain_uuid=uuid)]
)
for nid, uuid, iface in node_ids_and_uuids
}
return State(
node_thunderbolt=node_thunderbolt,
node_rdma_ctl=rdma_ctl or {},
)
def _has_rdma_edge(topology: Topology, source: NodeId, sink: NodeId) -> bool:
return any(
isinstance(edge, RDMAConnection)
for edge in topology.get_all_connections_between(source, sink)
)
def test_mac_thunderbolt_connections_emits_rdma_when_both_endpoints_enabled():
node_a = NodeId()
node_b = NodeId()
state = _make_state_with_thunderbolt_idents(
(node_a, "uuid-a", "rdma_en1"),
(node_b, "uuid-b", "rdma_en1"),
rdma_ctl={
node_a: NodeRdmaCtlStatus(enabled=True),
node_b: NodeRdmaCtlStatus(enabled=True),
},
)
event = NodeGatheredInfo(
node_id=node_a,
when=_now(),
info=MacThunderboltConnections(
conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
),
)
new_state = apply_node_gathered_info(event, state)
assert _has_rdma_edge(new_state.topology, node_a, node_b)
def test_mac_thunderbolt_connections_skips_rdma_when_source_rdma_ctl_disabled():
node_a = NodeId()
node_b = NodeId()
state = _make_state_with_thunderbolt_idents(
(node_a, "uuid-a", "rdma_en1"),
(node_b, "uuid-b", "rdma_en1"),
rdma_ctl={
node_a: NodeRdmaCtlStatus(enabled=False),
node_b: NodeRdmaCtlStatus(enabled=True),
},
)
event = NodeGatheredInfo(
node_id=node_a,
when=_now(),
info=MacThunderboltConnections(
conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
),
)
new_state = apply_node_gathered_info(event, state)
assert not _has_rdma_edge(new_state.topology, node_a, node_b)
def test_mac_thunderbolt_connections_skips_rdma_when_sink_rdma_ctl_disabled():
node_a = NodeId()
node_b = NodeId()
state = _make_state_with_thunderbolt_idents(
(node_a, "uuid-a", "rdma_en1"),
(node_b, "uuid-b", "rdma_en1"),
rdma_ctl={
node_a: NodeRdmaCtlStatus(enabled=True),
node_b: NodeRdmaCtlStatus(enabled=False),
},
)
event = NodeGatheredInfo(
node_id=node_a,
when=_now(),
info=MacThunderboltConnections(
conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
),
)
new_state = apply_node_gathered_info(event, state)
assert not _has_rdma_edge(new_state.topology, node_a, node_b)
def test_mac_thunderbolt_connections_skips_rdma_when_rdma_ctl_status_missing():
"""Missing rdma_ctl status defaults to not-enabled — node is RDMA-incapable."""
node_a = NodeId()
node_b = NodeId()
state = _make_state_with_thunderbolt_idents(
(node_a, "uuid-a", "rdma_en1"),
(node_b, "uuid-b", "rdma_en1"),
rdma_ctl={
node_a: NodeRdmaCtlStatus(enabled=True),
# node_b intentionally absent
},
)
event = NodeGatheredInfo(
node_id=node_a,
when=_now(),
info=MacThunderboltConnections(
conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
),
)
new_state = apply_node_gathered_info(event, state)
assert not _has_rdma_edge(new_state.topology, node_a, node_b)
def test_rdma_ctl_status_disabled_purges_existing_rdma_edges():
"""When a node reports rdma_ctl disabled, all RDMA edges touching it must be removed."""
node_a = NodeId()
node_b = NodeId()
# Start with both nodes RDMA-enabled and existing RDMA edges in the topology.
state = _make_state_with_thunderbolt_idents(
(node_a, "uuid-a", "rdma_en1"),
(node_b, "uuid-b", "rdma_en1"),
rdma_ctl={
node_a: NodeRdmaCtlStatus(enabled=True),
node_b: NodeRdmaCtlStatus(enabled=True),
},
)
state = apply_node_gathered_info(
NodeGatheredInfo(
node_id=node_a,
when=_now(),
info=MacThunderboltConnections(
conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
),
),
state,
)
state = apply_node_gathered_info(
NodeGatheredInfo(
node_id=node_b,
when=_now(),
info=MacThunderboltConnections(
conns=[ThunderboltConnection(source_uuid="uuid-b", sink_uuid="uuid-a")]
),
),
state,
)
assert _has_rdma_edge(state.topology, node_a, node_b)
assert _has_rdma_edge(state.topology, node_b, node_a)
# Now node_a flips to rdma_ctl disabled — both directions of RDMA edge must drop.
state = apply_node_gathered_info(
NodeGatheredInfo(
node_id=node_a, when=_now(), info=RdmaCtlStatus(enabled=False)
),
state,
)
assert not _has_rdma_edge(state.topology, node_a, node_b)
assert not _has_rdma_edge(state.topology, node_b, node_a)
assert state.node_rdma_ctl[node_a].enabled is False
def test_topology_remove_all_rdma_connections_touching_keeps_socket_edges():
"""Purging RDMA edges for a disabled node must not affect non-RDMA edges."""
from exo.shared.types.multiaddr import Multiaddr
from exo.shared.types.topology import Connection, SocketConnection
topology = Topology()
node_a = NodeId()
node_b = NodeId()
topology.add_node(node_a)
topology.add_node(node_b)
topology.add_connection(
Connection(
source=node_a,
sink=node_b,
edge=RDMAConnection(
source_rdma_iface="rdma_en1", sink_rdma_iface="rdma_en1"
),
)
)
socket_edge = SocketConnection(
sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000")
)
topology.add_connection(Connection(source=node_a, sink=node_b, edge=socket_edge))
topology.remove_all_rdma_connections_touching(node_a)
assert not _has_rdma_edge(topology, node_a, node_b)
# Socket edge survives.
assert any(
isinstance(edge, SocketConnection)
for edge in topology.get_all_connections_between(node_a, node_b)
)
+1 -1
View File
@@ -327,7 +327,7 @@ async def test_connection_message_triggers_new_round_broadcast() -> None:
tg.start_soon(election.run)
# Send any connection message object; we close quickly to cancel before result creation
await cm_tx.send(ConnectionMessage(node_id=NodeId(), connected=True))
await cm_tx.send(ConnectionMessage(connected=True))
# Expect a broadcast for the new round at clock=1
while True:
+16
View File
@@ -169,6 +169,22 @@ class Topology:
for conn in new_connections:
self.add_connection(conn)
def remove_all_rdma_connections_touching(self, node_id: NodeId) -> None:
"""Remove every RDMA edge incident to ``node_id`` (incoming or outgoing)."""
if node_id not in self._vertex_indices:
return
rx_idx = self._vertex_indices[node_id]
rdma_edge_idxs = [
edge_idx
for edge_idx in (
*self._graph.out_edge_indices(rx_idx),
*self._graph.in_edge_indices(rx_idx),
)
if isinstance(self._graph.get_edge_data_by_index(edge_idx), RDMAConnection)
]
for edge_idx in rdma_edge_idxs:
self._graph.remove_edge_from_index(edge_idx)
def remove_connection(self, conn: Connection) -> None:
if (
conn.source not in self._vertex_indices
+14 -2
View File
@@ -32,15 +32,21 @@ class TokenChunk(BaseChunk):
class ErrorChunk(BaseChunk):
error_message: str
finish_reason: Literal["error"] = "error"
@property
def finish_reason(self) -> Literal["error"]:
return "error"
class ToolCallChunk(BaseChunk):
tool_calls: list[ToolCallItem]
usage: Usage | None
finish_reason: Literal["tool_calls"] = "tool_calls"
stats: GenerationStats | None = None
@property
def finish_reason(self) -> Literal["tool_calls"]:
return "tool_calls"
class ImageChunk(BaseChunk):
data: str
@@ -84,7 +90,13 @@ class PrefillProgressChunk(BaseChunk):
processed_tokens: int
total_tokens: int
@property
def finish_reason(self) -> FinishReason | None:
return None
StatusChunk = PrefillProgressChunk
GenerationChunk = TokenChunk | ImageChunk | ToolCallChunk | ErrorChunk
TextGenerationChunk = TokenChunk | ToolCallChunk | ErrorChunk
ImageGenerationChunk = ImageChunk | ErrorChunk
Chunk = StatusChunk | GenerationChunk
+2 -2
View File
@@ -159,10 +159,10 @@ Event = (
| NodeTimedOut
| NodeGatheredInfo
| NodeDownloadProgress
| ChunkGenerated
| InputChunkReceived
| TopologyEdgeCreated
| TopologyEdgeDeleted
| ChunkGenerated
| InputChunkReceived
| TracesCollected
| TracesMerged
| CustomModelCardAdded
+5 -1
View File
@@ -5,8 +5,9 @@ from typing import Any, cast
from pydantic import ConfigDict, Field, field_serializer, field_validator
from pydantic.alias_generators import to_camel
from exo.shared.models.model_cards import ModelCard
from exo.shared.topology import Topology, TopologySnapshot
from exo.shared.types.common import NodeId
from exo.shared.types.common import ModelId, NodeId
from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
from exo.shared.types.profiling import (
DiskUsage,
@@ -65,6 +66,9 @@ class State(FrozenModel):
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] = {}
@field_serializer("topology", mode="plain")
def _encode_topology(self, value: Topology) -> TopologySnapshot:
return value.to_snapshot()
+2 -2
View File
@@ -135,9 +135,9 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
prefill_endpoint: str | None = None
def with_card_sampling_defaults(self) -> "TextGenerationTaskParams":
from exo.shared.models.model_cards import get_card
from exo.shared.models import model_cards
card = get_card(self.model)
card = model_cards.card_cache.get(self.model)
if card is None:
return self
+1 -1
View File
@@ -38,7 +38,7 @@ def print_startup_banner(port: int) -> None:
🌐 Dashboard & API Ready
Dashboard & API Ready
{dashboard_url}{" " * (69 - len(dashboard_url))}
+6 -1
View File
@@ -143,6 +143,7 @@ class ImageEngine(Engine):
Generator[tuple[TaskId, Chunk | FinishedResponse | CancelledResponse]] | None
) = field(init=False, default=None)
queue: deque[ImageTask] = field(init=False, default_factory=deque)
_cancelled_tasks: set[TaskId] = field(init=False, default_factory=set)
def warmup(self) -> None:
image = warmup_image_generator(model=self.image_model)
@@ -168,7 +169,11 @@ class ImageEngine(Engine):
task = self.queue.popleft()
self.current_gen = self._run_image_task(task.task_id, task.task_params)
resp = next(self.current_gen, None)
return (resp,) if resp is not None else ()
return (
(resp,)
if resp is not None and _is_primary_output_node(self.shard_metadata)
else ()
)
def close(self) -> None:
with contextlib.suppress(NameError, AttributeError):
+22 -17
View File
@@ -4,24 +4,22 @@ from datetime import datetime, timezone
import anyio
from anyio import fail_after, to_thread
from exo_net import PySession
from loguru import logger
from exo.api.types import ImageEditsTaskParams
from exo.download.download_utils import is_read_only_model_dir, resolve_existing_model
from exo.shared.apply import apply
from exo.shared.constants import EXO_MAX_INSTANCE_RETRIES
from exo.shared.models.model_cards import ModelId, add_to_card_cache, delete_custom_card
from exo.shared.models.model_cards import ModelId, card_cache
from exo.shared.types.chunks import InputImageChunk
from exo.shared.types.commands import (
DeleteInstance,
ForwarderCommand,
ForwarderDownloadCommand,
StartDownload,
)
from exo.shared.types.common import CommandId, NodeId, SystemId
from exo.shared.types.events import (
CustomModelCardAdded,
CustomModelCardDeleted,
Event,
IndexedEvent,
InputChunkReceived,
@@ -67,16 +65,17 @@ class Worker:
*,
event_receiver: Receiver[IndexedEvent],
event_sender: Sender[Event],
session: PySession,
# This is for requesting updates. It doesn't need to be a general command sender right now,
# but I think it's the correct way to be thinking about commands
command_sender: Sender[ForwarderCommand],
download_command_sender: Sender[ForwarderDownloadCommand],
api_port: int,
):
self.node_id: NodeId = node_id
self.event_receiver = event_receiver
self.event_sender = event_sender
self.command_sender = command_sender
self.session = session
self.command_sender = session.net_sender("orchestrator")
self.download_command_sender = download_command_sender
self.api_port = api_port
@@ -110,11 +109,12 @@ class Worker:
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)
finally:
# Actual shutdown code - waits for all tasks to complete before executing.
logger.info("Stopping Worker")
self.event_sender.close()
self.command_sender.close()
self.download_command_sender.close()
for runner in self.runners.values():
runner.shutdown()
@@ -151,7 +151,6 @@ class Worker:
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]
@@ -172,12 +171,18 @@ class Worker:
)
] = img
if isinstance(event, CustomModelCardAdded):
await event.model_card.save_to_custom_dir()
add_to_card_cache(event.model_card)
async def _reconcile_custom_cards(self) -> None:
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:
continue
await card_cache.save(card)
if isinstance(event, CustomModelCardDeleted):
await delete_custom_card(event.model_id)
for card in await card_cache.list_all():
if card.model_id not in target:
await card_cache.pop(card.model_id)
async def plan_step(self):
while True:
@@ -204,10 +209,9 @@ class Worker:
f"Instance {iid} exceeded {EXO_MAX_INSTANCE_RETRIES} retries, requesting deletion"
)
await self.command_sender.send(
ForwarderCommand(
origin=self._system_id,
command=DeleteInstance(instance_id=iid),
)
DeleteInstance(instance_id=iid)
.model_dump_json()
.encode("utf-8")
)
continue
@@ -370,6 +374,7 @@ class Worker:
runner = RunnerSupervisor.create(
bound_instance=task.bound_instance,
event_sender=self.event_sender.clone(),
session=self.session,
)
self.runners[task.bound_instance.bound_runner_id] = runner
self._tg.start_soon(runner.run)
@@ -138,8 +138,10 @@ class SequentialGenerator(Engine):
def agree_on_tasks(self) -> None:
"""Agree between all ranks about the task ordering (some may have received in different order or not at all)."""
agreed, different = mx_all_gather_tasks(self._maybe_queue, self.group)
self._queue.extend(task for task in self._maybe_queue if task in agreed)
self._maybe_queue = [task for task in self._maybe_queue if task in different]
# Extend from `agreed` (sorted by task_id on all ranks) to guarantee every
# rank enqueues tasks in the same order, preventing TP collective deadlocks.
self._queue.extend(agreed)
self._maybe_queue = list(different)
def agree_on_cancellations(self) -> None:
"""Agree between all ranks about which tasks to cancel."""
@@ -197,9 +199,14 @@ class SequentialGenerator(Engine):
self._active = None
raise
return itertools.chain(
output,
map(lambda task: (task, CancelledResponse()), self._cancelled_tasks),
return filter(
lambda chunk: (
not isinstance(chunk[1], GenerationChunk) or self.device_rank == 0
),
itertools.chain(
output,
map(lambda task: (task, CancelledResponse()), self._cancelled_tasks),
),
)
def _start_next(self) -> None:
@@ -368,8 +375,10 @@ class BatchGenerator(Engine):
def agree_on_tasks(self) -> None:
"""Agree between all ranks about the task ordering (some may have received in different order or not at all)."""
agreed, different = mx_all_gather_tasks(self._maybe_queue, self.group)
self._queue.extend(task for task in self._maybe_queue if task in agreed)
self._maybe_queue = [task for task in self._maybe_queue if task in different]
# Extend from `agreed` (sorted by task_id on all ranks) to guarantee every
# rank enqueues tasks in the same order, preventing TP collective deadlocks.
self._queue.extend(agreed)
self._maybe_queue = list(different)
def agree_on_cancellations(self) -> None:
"""Agree between all ranks about which tasks to cancel."""
@@ -449,7 +458,12 @@ class BatchGenerator(Engine):
output.append((task.task_id, FinishedResponse()))
del self._active_tasks[uid]
return itertools.chain(output, self._apply_cancellations())
return filter(
lambda chunk: (
not isinstance(chunk[1], GenerationChunk) or self.device_rank == 0
),
itertools.chain(output, self._apply_cancellations()),
)
def _apply_cancellations(
self,
+3 -3
View File
@@ -86,7 +86,7 @@ class Runner:
self,
bound_instance: BoundInstance,
builder: Builder,
event_sender: MpSender[Event],
event_sender: MpSender[Event | ChunkGenerated],
task_receiver: MpReceiver[Task],
):
self.event_sender = event_sender
@@ -390,5 +390,5 @@ class Runner:
chunk: Chunk,
command_id: CommandId,
):
if self.device_rank == 0:
self.event_sender.send(ChunkGenerated(command_id=command_id, chunk=chunk))
assert isinstance(self.generator, Engine)
self.event_sender.send(ChunkGenerated(command_id=command_id, chunk=chunk))
+33 -15
View File
@@ -10,9 +10,11 @@ from anyio import (
ClosedResourceError,
to_thread,
)
from exo_net import NetSender, PySession
from loguru import logger
from exo.shared.types.chunks import ErrorChunk
from exo.shared.types.commands import CommandId
from exo.shared.types.events import (
ChunkGenerated,
Event,
@@ -45,20 +47,18 @@ from exo.utils.channels import MpReceiver, MpSender, Sender, mp_channel
from exo.utils.task_group import TaskGroup
from exo.worker.runner.bootstrap import entrypoint
PREFILL_TIMEOUT_SECONDS = 60
DECODE_TIMEOUT_SECONDS = 5
@dataclass(eq=False)
class RunnerSupervisor:
shard_metadata: ShardMetadata
bound_instance: BoundInstance
runner_process: mp.Process
initialize_timeout: float
_ev_recv: MpReceiver[Event]
_task_sender: MpSender[Task]
_event_sender: Sender[Event]
_cancel_sender: MpSender[TaskId]
session: PySession
_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)
@@ -75,7 +75,7 @@ class RunnerSupervisor:
*,
bound_instance: BoundInstance,
event_sender: Sender[Event],
initialize_timeout: float = 400,
session: PySession,
) -> Self:
ev_send, ev_recv = mp_channel[Event]()
task_sender, task_recv = mp_channel[Task]()
@@ -99,11 +99,11 @@ class RunnerSupervisor:
bound_instance=bound_instance,
shard_metadata=shard_metadata,
runner_process=runner_process,
initialize_timeout=initialize_timeout,
_ev_recv=ev_recv,
_task_sender=task_sender,
_cancel_sender=cancel_sender,
_event_sender=event_sender,
session=session,
)
return self
@@ -210,9 +210,26 @@ class RunnerSupervisor:
await self._check_runner(TimeoutError("cancel pipe blocked"))
async def _forward_events(self):
pubs: dict[CommandId, NetSender] = {}
try:
with self._ev_recv as events:
async for event in events:
if isinstance(event, ChunkGenerated):
if (pub := pubs.get(event.command_id, None)) is None:
pub = pubs[event.command_id] = self.session.net_sender(
f"runners/{self.bound_instance.bound_runner_id}/active_tasks/{event.command_id}/chunks"
)
sent = await pub.send(
event.chunk.model_dump_json().encode("utf-8")
)
if not sent:
logger.warning(
"api node closed communication, dropping chunk"
)
if event.chunk.finish_reason is not None:
pubs.pop(event.command_id, None)
continue
if isinstance(event, RunnerStatusUpdated):
self.status = event.runner_status
if isinstance(event, TaskAcknowledged):
@@ -275,17 +292,18 @@ class RunnerSupervisor:
for task in self.in_progress.values():
if isinstance(task, (TextGeneration, ImageGeneration, ImageEdits)):
with anyio.CancelScope(shield=True):
await self._event_sender.send(
ChunkGenerated(
command_id=task.command_id,
chunk=ErrorChunk(
model=self.shard_metadata.model_card.model_id,
error_message=(
"Runner shutdown before completing command "
f"({cause})"
),
send = self.session.net_sender(
f"runners/{self.bound_instance.bound_runner_id}/active_tasks/{task.command_id}/chunks"
)
await send.send(
ErrorChunk(
model=self.shard_metadata.model_card.model_id,
error_message=(
f"Runner shutdown before completing command ({cause})"
),
)
.model_dump_json()
.encode("utf-8")
)
try:
@@ -16,7 +16,7 @@ from exo.download.download_utils import (
fetch_file_list_with_cache,
resolve_model_dir,
)
from exo.shared.models.model_cards import ModelCard, ModelId, get_model_cards
from exo.shared.models.model_cards import ModelCard, ModelId, card_cache
from exo.worker.engines.mlx.utils_mlx import (
get_eos_token_ids_for_model,
load_tokenizer_for_model_id,
@@ -76,7 +76,7 @@ def get_test_models() -> list[ModelCard]:
"""Get a representative sample of models to test."""
# Pick one model from each family to test
families: dict[str, ModelCard] = {}
for card in asyncio.run(get_model_cards()):
for card in asyncio.run(card_cache.list_all()):
# Extract family name (e.g., "llama-3.1" from "llama-3.1-8b")
parts = card.model_id.short().split("-")
family = "-".join(parts[:2]) if len(parts) >= 2 else parts[0]
@@ -298,7 +298,7 @@ async def test_tokenizer_special_tokens(model_card: ModelCard) -> None:
async def test_kimi_tokenizer_specifically():
"""Test Kimi tokenizer with its specific patches and quirks."""
kimi_models = [
card for card in await get_model_cards() if "kimi" in card.model_id.lower()
card for card in await card_cache.list_all() if "kimi" in card.model_id.lower()
]
if not kimi_models:
@@ -350,7 +350,7 @@ async def test_glm_tokenizer_specifically():
glm_model_cards = [
card
for card in await get_model_cards()
for card in await card_cache.list_all()
if contains(card, "glm")
and not contains(card, "-5")
and not contains(card, "4.7")
+15
View File
@@ -0,0 +1,15 @@
import anyio
from exo_net import StateProxy
async def main():
sp = await StateProxy.init()
while True:
data = await sp.snapshot()
if data != "{}":
print(data)
await anyio.sleep(1)
if __name__ == "__main__":
anyio.run(main)
Generated
+16 -16
View File
@@ -22,7 +22,7 @@ prerelease-mode = "allow"
members = [
"exo",
"exo-bench",
"exo-pyo3-bindings",
"exo-net",
]
constraints = [{ name = "transformers", specifier = ">=5.6.2" }]
overrides = [
@@ -385,7 +385,7 @@ dependencies = [
{ name = "aiofiles", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "aiohttp", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "anyio", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "exo-pyo3-bindings", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "exo-net", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "fastapi", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "filelock", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "httpx", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
@@ -394,7 +394,7 @@ dependencies = [
{ name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "mflux", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx", version = "0.31.1", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx", version = "0.32.0.dev20260427+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-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx", version = "0.32.0.dev20260429+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-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx-lm", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx-vlm", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "msgspec", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
@@ -416,21 +416,21 @@ build = [
]
cpu = [
{ name = "mlx", version = "0.31.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx", version = "0.32.0.dev20260427+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-cpu') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx", version = "0.32.0.dev20260429+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-cpu') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx-cpu", marker = "sys_platform == 'linux'" },
{ name = "mlx-lm", marker = "sys_platform == 'linux'" },
{ name = "mlx-vlm", marker = "sys_platform == 'linux'" },
]
cuda12 = [
{ name = "mlx", version = "0.31.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx", version = "0.32.0.dev20260427+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-cuda12') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx", version = "0.32.0.dev20260429+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-cuda12') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx-cuda-12", marker = "sys_platform == 'linux'" },
{ name = "mlx-lm", marker = "sys_platform == 'linux'" },
{ name = "mlx-vlm", marker = "sys_platform == 'linux'" },
]
cuda13 = [
{ name = "mlx", version = "0.31.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx", version = "0.32.0.dev20260427+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-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx", version = "0.32.0.dev20260429+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-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx-cuda-13", marker = "sys_platform == 'linux'" },
{ name = "mlx-lm", marker = "sys_platform == 'linux'" },
{ name = "mlx-vlm", marker = "sys_platform == 'linux'" },
@@ -451,7 +451,7 @@ requires-dist = [
{ name = "aiofiles", specifier = ">=24.1.0" },
{ name = "aiohttp", specifier = ">=3.12.14" },
{ name = "anyio", specifier = "==4.11.0" },
{ name = "exo-pyo3-bindings", editable = "rust/exo_pyo3_bindings" },
{ name = "exo-net", editable = "rust/exo_net" },
{ name = "fastapi", specifier = ">=0.116.1" },
{ name = "filelock", specifier = ">=3.18.0" },
{ name = "httpx", specifier = ">=0.28.1" },
@@ -541,13 +541,13 @@ requires-dist = [
]
[[package]]
name = "exo-pyo3-bindings"
version = "0.2.1"
source = { editable = "rust/exo_pyo3_bindings" }
name = "exo-net"
version = "0.3.0"
source = { editable = "rust/exo_net" }
[package.dev-dependencies]
dev = [
{ name = "exo-pyo3-bindings", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "exo-net", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "pytest", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "pytest-asyncio", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
]
@@ -556,7 +556,7 @@ dev = [
[package.metadata.requires-dev]
dev = [
{ name = "exo-pyo3-bindings", editable = "rust/exo_pyo3_bindings" },
{ name = "exo-net", editable = "rust/exo_net" },
{ name = "pytest", specifier = ">=8.4.0" },
{ name = "pytest-asyncio", specifier = ">=1.0.0" },
]
@@ -1213,7 +1213,7 @@ dependencies = [
{ name = "hf-transfer", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "huggingface-hub", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "matplotlib", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx", version = "0.32.0.dev20260427+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-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx", version = "0.32.0.dev20260429+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-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "numpy", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "opencv-python", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "piexif", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
@@ -1263,7 +1263,7 @@ wheels = [
[[package]]
name = "mlx"
version = "0.32.0.dev20260427+cc3f3e60"
version = "0.32.0.dev20260429+cc3f3e60"
source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }
resolution-markers = [
"sys_platform == 'darwin'",
@@ -1315,7 +1315,7 @@ source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Fdeepseek-v4
dependencies = [
{ name = "jinja2", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx", version = "0.31.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx", version = "0.32.0.dev20260427+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-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx", version = "0.32.0.dev20260429+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-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "numpy", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "protobuf", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "pyyaml", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
@@ -1332,7 +1332,7 @@ dependencies = [
{ name = "fastapi", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "miniaudio", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx", version = "0.31.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx", version = "0.32.0.dev20260427+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-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx", version = "0.32.0.dev20260429+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-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "mlx-lm", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "numpy", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
{ name = "opencv-python", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },